import gradio as gr
from transformers import pipeline
from huggingface_hub import login
from openai import OpenAI
import os
import json

# --- 1. SETUP ---
hf_token = os.getenv("HF_TOKEN")
if hf_token:
    login(token=hf_token)

# --- 2. LOAD TOOL ---
print("Loading Mental-Longformer...")
model_name = "avtak/erisk-longformer-depression-v1"
classifier = pipeline("text-classification", model=model_name, truncation=True, max_length=4096, top_k=None)

# --- 3. MCP TOOLS ---

def get_crisis_resources(location: str = "Global") -> str:
    """Returns mental health resources based on location."""
    resources = {
        "US": "🇺🇸 **US:** Crisis Text Line: 741741 | Suicide Lifeline: 988",
        "Malaysia": "🇲🇾 **Malaysia:** Befrienders KL: 03-76272929 | Talian Kasih: 15999",
        "Global": "🌍 **International:** [befrienders.org](https://www.befrienders.org)"
    }
    for key in resources:
        if location and key.lower() in location.lower():
            return resources[key]
    return resources["Global"]

def detect_depression_risk(text: str) -> dict:
    """Analyzes text using Mental-Longformer (eRisk 2025)."""
    
    # --- THESIS LOGIC: AGGREGATION ---
    # 1. We replace single newlines with double newlines for the MODEL
    processed_text = text.replace("\n", "\n\n") 
    
    results = classifier(processed_text)[0]
    prob = next((r['score'] for r in results if r['label'] == 'LABEL_1'), 0)
    
    # Thesis Thresholds (Figure 4.15)
    if prob < 0.40: 
        level = "Low Risk"
        biomarker = "Healthy External Focus"
        desc = "Matches 'Isolated Control' group. High lexical diversity, focus on hobbies/events."
        color = "#10b981" # Green
    elif 0.40 <= prob < 0.60: 
        level = "Moderate Risk"
        biomarker = "Echo Chamber Interaction"
        desc = "Matches 'Interactive Non-Depressed' group. Engaging in support forums but likely not clinically depressed."
        color = "#f59e0b" # Yellow
    else: 
        level = "High Risk"
        biomarker = "Nocturnal & High-Effort"
        desc = "Matches 'Depressed' cohort. Indicators: Nocturnal posting spikes (00-05 UTC), high-effort/low-frequency posting."
        color = "#ef4444" # Red
        
    return {
        "probability": prob, 
        "risk_level": level, 
        "biomarker": biomarker,
        "description": desc,
        "color": color,
        "word_count": len(processed_text.split())
    }

# --- 4. AGENT REASONING (SambaNova + Nebius Kimi) ---
def agent_reasoning(text, risk_data, provider="SambaNova"):
    """
    Uses Sponsor APIs to generate the analysis report.
    """
    client = None
    model_id = None
    
    # IMPROVED SYSTEM PROMPT
    system_prompt = f"""
    You are 'MindSight', a Clinical Research Agent specialized in linguistic biomarker analysis.
    
    CONTEXTUAL DATA:
    - Detected Risk: {risk_data['risk_level']} ({risk_data['probability']:.1%})
    - Key Biomarker: {risk_data['biomarker']}
    - Dataset Context: Based on eRisk longitudinal analysis (2017-2022).
    
    USER TEXT SNIPPET: "{text[:1000]}..."
    
    YOUR MISSION:
    Synthesize the clinical data with the user's personal narrative to provide a supportive, research-backed insight.
    
    GUIDELINES:
    1. **Validate:** Start by acknowledging the specific struggles or sentiments expressed in the text (e.g., "I hear your exhaustion...").
    2. **Connect:** Explain *why* the risk level was triggered using the Biomarker. 
       - If High Risk: Link sleep/energy complaints to the "Nocturnal Posting" pattern (activity spikes 00:00-05:00 UTC).
       - If Moderate Risk: Link their focus on others to the "Supportive Responder" effect found in online echo chambers.
       - If Low Risk: Affirm their "Healthy External Focus" (hobbies, work, events).
    3. **Bridge:** Offer a gentle, non-medical bridge to the resources below.
    
    CONSTRAINTS:
    - Be warm but objective.
    - NO medical diagnosis (use "suggests alignment with...").
    - Max 80 words.
    """

    try:
        # --- SPONSOR 1: NEBIUS (Kimi K2) ---
        if provider == "Nebius (Kimi K2)":
            api_key = os.getenv("NEBIUS_API_KEY")
            if not api_key: return "⚠️ Nebius API Key missing."
            
            client = OpenAI(
                base_url="https://api.tokenfactory.nebius.com/v1/", 
                api_key=api_key
            )
            model_id = "moonshotai/Kimi-K2-Instruct" 

        # --- SPONSOR 2: SAMBANOVA ---
        else: # Default to SambaNova
            api_key = os.getenv("SAMBANOVA_API_KEY")
            if not api_key: return "⚠️ SambaNova API Key missing."
            
            client = OpenAI(
                base_url="https://api.sambanova.ai/v1", 
                api_key=api_key
            )
            model_id = "Meta-Llama-3.3-70B-Instruct" 

        # EXECUTE
        response = client.chat.completions.create(
            model=model_id,
            messages=[
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": "Analyze this."}
            ],
            temperature=0.6,
            max_tokens=300
        )
        return response.choices[0].message.content

    except Exception as e:
        return f"Reasoning Error ({provider}): {str(e)}"

# --- 5. PIPELINE ---
def full_analysis_pipeline(user_text, location, provider):
    if not user_text.strip(): return "Please enter text."
    
    # 1. Run Tool
    risk_data = detect_depression_risk(user_text)
    
    # 2. Get Resources
    resources = get_crisis_resources(location)
    
    # 3. Agent Reasoning
    explanation = agent_reasoning(user_text, risk_data, provider)
    
    # 4. Color Logic (Fixed)
    color = risk_data['color']
    
    # 5. Build HTML/Markdown Report
    return f"""
    <div style="border-left: 6px solid {color}; padding-left: 15px; background-color: {color}10; border-radius: 5px; padding: 20px;">
        <h2 style="color:{color}; margin:0;">{risk_data['risk_level']}</h2>
        <h3 style="margin-top:5px;">Probability: {risk_data['probability']*100:.1f}%</h3>
    </div>
    
    ### 👁️ MindSight Analysis ({provider})
    <div style="white-space: pre-wrap; font-family: sans-serif; font-size: 1.1em;">
    {explanation}
    </div>
    
    ---
    ### 🆘 Recommended Resources
    {resources}
    """

# --- 6. EXAMPLES ---
example_low = """The new update for the Linux kernel (6.8) finally addressed the driver issues I was seeing with Realtek cards. I read the changelog on kernel.org and noticed they patched the module specific to the rtl8821ce chipset.

On a different note, the defensive stats for the basketball team have improved significantly. Allowing 15% fewer points per possession is a game changer.

I also spent some time woodworking. I'm building a walnut coffee table and the grain is tricky to plane. I might switch to a polyurethane finish for durability."""

example_mod = """Things have been pretty busy at work. We're in the middle of a big project and deadlines are tight. Staying a bit later than usual but that's just how it goes sometimes. Team is handling it well overall.

Sleep could be better. Been staying up too late scrolling my phone. Need to work on that. Usually feel okay once I get moving in the morning though.

Went hiking with some friends last weekend which was nice. Good to get outside and move around. We're talking about doing another trip next month. Weather should be better by then."""

example_high = """I don't know why I even bother getting out of bed anymore. I slept for 12 hours yesterday and I'm still exhausted. It feels like my limbs weigh a thousand pounds.

I've been avoiding my friends for weeks. They keep texting me, but I can't bring myself to reply. The thought of socializing is terrifying.

Everything feels like a shade of grey. I can't concentrate on my work. I feel like I'm drowning while everyone else is breathing fine. I was up until 4 AM again last night just staring at the ceiling."""

# --- 7. UI ---
with gr.Blocks(title="Depression Risk Agent") as demo:
    gr.Markdown("# 🧠 Early Depression Detection Agent (MCP)")
    gr.Markdown("Agentic system using **Mental-Longformer** (Tool) + **Multi-Provider Reasoning** (Nebius/SambaNova).")
    
    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("### 1. User History")
            input_text = gr.Textbox(
                label="User Timeline", 
                lines=10, 
                placeholder="[Post 1] ...\n\n[Post 2] ...\n\n(Context >100 words recommended)"
            )
            
            gr.Markdown("### 2. Thesis Patterns (Click to Load)")
            with gr.Row():
                btn_low = gr.Button("🟢 Low Risk", size="sm", variant="secondary")
                btn_mod = gr.Button("🟡 Moderate", size="sm", variant="secondary")
                btn_high = gr.Button("🔴 High Risk", size="sm", variant="secondary")
            
            gr.Markdown("### 3. Agent Settings")
            with gr.Row():
                # RENAMED to be clearer
                loc_dropdown = gr.Dropdown(["Global", "US", "Malaysia"], value="Malaysia", label="Crisis Resource Region")
                provider_dropdown = gr.Dropdown(["SambaNova(Llama 3.3)", "Nebius (Kimi K2)"], value="SambaNova", label="Reasoning Brain")
            
            submit = gr.Button("🚀 Run Analysis Agent", variant="primary", size="lg")
            
        with gr.Column(scale=1):
            gr.Markdown("### 4. Agent Response")
            output = gr.Markdown(label="Response")
            
            with gr.Accordion("🔧 MCP Tools Exposed", open=True):
                gr.Markdown("""
                The following functions are exposed to the MCP Client:
                - `detect_depression_risk(text)`: Returns probability & thesis classification.
                - `get_crisis_resources(location)`: Returns localized help.
                """)

    submit.click(full_analysis_pipeline, inputs=[input_text, loc_dropdown, provider_dropdown], outputs=output)
    
    btn_low.click(lambda: example_low, None, input_text)
    btn_mod.click(lambda: example_mod, None, input_text)
    btn_high.click(lambda: example_high, None, input_text)

if __name__ == "__main__":
    demo.launch(mcp_server=True, theme=gr.themes.Soft())