import os
import requests
import json
import pandas as pd
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel

# Load DSM-5 Dataset
file_path = "dsm5_final_cleaned.csv"
df = pd.read_csv(file_path)

# OpenRouter API Configuration (DeepSeek Model)
OPENROUTER_API_KEY = os.getenv("OPENROUTER_API_KEY")  # Use environment variable for security
if not OPENROUTER_API_KEY:
    raise ValueError("OPENROUTER_API_KEY is missing. Set it as an environment variable.")
OPENROUTER_BASE_URL = "https://openrouter.ai/api/v1/chat/completions"

# Initialize FastAPI
app = FastAPI()

# Pydantic Models
class ChatRequest(BaseModel):
    message: str

class SummaryRequest(BaseModel):
    chat_history: list

def deepseek_request(prompt, max_tokens=300):
    """Helper function to send a request to DeepSeek API and handle response."""
    headers = {"Authorization": f"Bearer {OPENROUTER_API_KEY}", "Content-Type": "application/json"}
    payload = {
        "model": "deepseek/deepseek-r1-distill-llama-8b",
        "messages": [{"role": "user", "content": prompt}],
        "max_tokens": max_tokens,
        "temperature": 0.7
    }
    response = requests.post(OPENROUTER_BASE_URL, headers=headers, data=json.dumps(payload))
    if response.status_code == 200:
        response_json = response.json()
        if "choices" in response_json and response_json["choices"]:
            return response_json["choices"][0].get("message", {}).get("content", "").strip()
    return "Error: Unable to process the request."

def match_disorders(chat_history):
    """Match user symptoms with DSM-5 disorders based on keyword occurrence."""
    disorder_scores = {}
    for _, row in df.iterrows():
        disorder = row["Disorder"]
        keywords = row["Criteria"].split(", ")
        match_count = sum(1 for word in keywords if word in chat_history.lower())
        if match_count > 0:
            disorder_scores[disorder] = match_count
    sorted_disorders = sorted(disorder_scores, key=disorder_scores.get, reverse=True)
    return sorted_disorders[:3] if sorted_disorders else []

@app.post("/detect_disorders")
def detect_disorders(request: SummaryRequest):
    """Detect psychiatric disorders using DSM-5 keyword matching + DeepSeek validation."""
    full_chat = " ".join(request.chat_history)
    matched_disorders = match_disorders(full_chat)
    
    prompt = f"""
    The following is a psychiatric conversation:
    {full_chat}
    
    Based on DSM-5 diagnostic criteria, analyze the symptoms and determine the most probable psychiatric disorders.
    Here are possible disorder matches from DSM-5 keyword analysis: {', '.join(matched_disorders) if matched_disorders else 'None found'}.
    If no clear matches exist, diagnose based purely on symptom patterns and clinical reasoning.
    Return a **list** of disorders, separated by commas, without extra text.
    """
    
    response = deepseek_request(prompt, max_tokens=150)
    disorders = [disorder.strip() for disorder in response.split(",")] if response and response.lower() != "unspecified disorder" else matched_disorders
    return {"disorders": disorders if disorders else ["Unspecified Disorder"]}

@app.post("/get_treatment")
def get_treatment(request: SummaryRequest):
    """Retrieve structured treatment recommendations based on detected disorders."""
    detected_disorders = detect_disorders(request)["disorders"]
    disorders_text = ", ".join(detected_disorders)
    prompt = f"""
    The user has been diagnosed with: {disorders_text}.
    Provide a structured, evidence-based treatment plan including:
    - Therapy recommendations (e.g., CBT, DBT, psychotherapy).
    - Possible medications if applicable (e.g., SSRIs, anxiolytics, sleep aids).
    - Lifestyle and self-care strategies (e.g., sleep hygiene, mindfulness, exercise).
    If the user has suicidal thoughts, emphasize **immediate crisis intervention and emergency medical support.**
    """
    treatment_response = deepseek_request(prompt, max_tokens=200)
    return {"treatments": treatment_response}

@app.post("/summarize_chat")
def summarize_chat(request: SummaryRequest):
    """Generate a structured summary of the psychiatric consultation."""
    full_chat = " ".join(request.chat_history)
    detected_disorders = detect_disorders(request)["disorders"]
    treatment_response = get_treatment(request)["treatments"]
    prompt = f"""
    Summarize the following psychiatric conversation:
    {full_chat}
    
    - **Detected Disorders:** {', '.join(detected_disorders)}
    - **Suggested Treatments:** {treatment_response}
    
    The summary should include:
    - Main concerns reported by the user.
    - Key symptoms observed.
    - Possible underlying psychological conditions.
    - Recommended next steps, including professional consultation and self-care strategies.
    If suicidal thoughts were mentioned, highlight the **need for immediate crisis intervention and professional support.**
    """
    summary = deepseek_request(prompt, max_tokens=300)
    return {"summary": summary}

@app.post("/chat")
def chat(request: ChatRequest):
    """Generate AI psychiatric response for user input."""
    prompt = f"""
    You are an AI psychiatrist conducting a mental health consultation.
    The user is discussing their concerns and symptoms. Engage in a supportive conversation,
    ask relevant follow-up questions, and maintain an empathetic tone.
    
    User input:
    {request.message}
    
    Provide a meaningful response and a follow-up question if necessary.
    If the user mentions suicidal thoughts, respond with an urgent and compassionate tone,
    suggesting that they seek immediate professional help while providing emotional support.
    """
    response = deepseek_request(prompt, max_tokens=200)
    return {"response": response}