import os
import json
import pickle
import tempfile
import numpy as np
import pandas as pd
import pyotp
from datetime import datetime
from typing import Optional, List, Dict, Any

from fastapi import FastAPI, UploadFile, File, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import Response
from pydantic import BaseModel
from fpdf import FPDF
from deep_translator import GoogleTranslator

from inference_engine import ClinicalAI, get_logic_driven_questions
from audio_processing import AudioAnalyzer
from database_manager import init_db, save_session, get_all_data_as_csv

# ── Startup ────────────────────────────────────────────────────────────────────
app = FastAPI(title="CogniDetect API", version="1.0.0")

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

MODEL_DIR = "models"
AI_ENGINE: Optional[ClinicalAI] = None
AUDIO_TOOL: Optional[AudioAnalyzer] = None
RF_MODELS: Dict[str, Any] = {}


@app.on_event("startup")
def startup():
    global AI_ENGINE, AUDIO_TOOL, RF_MODELS
    init_db()
    AI_ENGINE = ClinicalAI()
    AUDIO_TOOL = AudioAnalyzer()
    rf_names = ["RF_risk.pkl", "rf_ADHD_sev.pkl", "rf_ASD_sev.pkl",
                "rf_SPCD_sev.pkl", "rf_DEP_sev.pkl", "rf_ANX_sev.pkl"]
    keys = ["risk", "ADHD", "ASD", "SPCD", "DEP", "ANX"]
    for k, name in zip(keys, rf_names):
        path = os.path.join(MODEL_DIR, name)
        if os.path.exists(path):
            RF_MODELS[k] = pickle.load(open(path, "rb"))


def load_config() -> dict:
    with open("config.json", "r") as f:
        return json.load(f)


# ── Request / Response models ──────────────────────────────────────────────────

class TranslateRequest(BaseModel):
    texts: List[str]
    target_lang: str


class QuestionnaireRequest(BaseModel):
    responses: List[int]


class SymptomFlagsRequest(BaseModel):
    s1: bool = False
    s2: bool = False
    s3: bool = False
    s4: bool = False
    s5: bool = False
    custom_symptom: Optional[str] = None


class NLPAnalyzeRequest(BaseModel):
    text: str
    category: Optional[str] = None


class MetaFusionRequest(BaseModel):
    rf_probs: Dict[str, float]
    nlp_probs_list: List[List[float]]


class SessionSaveRequest(BaseModel):
    user_data: Dict[str, Any]
    rf_answers: List[int]
    rf_risk: Dict[str, float]
    nlp_data: Optional[List[Dict]] = None
    final_diagnosis: str


class ReportRequest(BaseModel):
    user_data: Dict[str, Any]
    rf_data: Dict[str, Any]
    nlp_data: List[Dict]
    final_result: Dict[str, Any]


class AdminPasswordRequest(BaseModel):
    password: str


class AdminTOTPRequest(BaseModel):
    code: str


class ConfigUpdateRequest(BaseModel):
    config: Dict[str, Any]


# ── Utility ───────────────────────────────────────────────────────────────────

def _clean(text: Any) -> str:
    return str(text).encode("ascii", "ignore").decode("ascii") if text else ""


# ── Routes ────────────────────────────────────────────────────────────────────

@app.get("/health")
def health():
    return {"status": "ok", "models_loaded": AI_ENGINE is not None}


@app.get("/api/config")
def get_config():
    cfg = load_config()
    return {
        "questions": cfg["questions"],
        "symptoms": cfg["symptoms"],
        "nlp_questions": cfg["nlp_questions"],
    }


@app.post("/api/translate")
def translate(req: TranslateRequest):
    if req.target_lang == "en":
        return {"translated": req.texts}
    results = []
    for text in req.texts:
        try:
            t = GoogleTranslator(source="en", target=req.target_lang).translate(text)
            results.append(t)
        except Exception:
            results.append(text)
    return {"translated": results}


@app.post("/api/questionnaire/analyze")
def analyze_questionnaire(req: QuestionnaireRequest):
    if len(req.responses) != 27:
        raise HTTPException(400, f"Expected 27 responses, got {len(req.responses)}")

    responses = req.responses
    feature_names = (
        [f"ADHD_Q{i}" for i in range(1, 8)]
        + [f"ASD_Q{i}" for i in range(1, 7)]
        + [f"SPCD_Q{i}" for i in range(1, 5)]
        + [f"DEP_Q{i}" for i in range(1, 6)]
        + [f"ANX_Q{i}" for i in range(1, 6)]
    )

    try:
        X_df = pd.DataFrame([responses], columns=feature_names)
        risk_model = RF_MODELS.get("risk")
        rf_pred = risk_model.predict(X_df)[0] if risk_model else [0, 0, 0, 0, 0]

        rf_probs = {"ADHD": 0.01, "Depression": 0.01, "Anxiety": 0.01, "Autism": 0.01}

        if rf_pred[0] == 1:
            rf_probs["ADHD"] = 0.85
        if rf_pred[1] == 1 or rf_pred[2] == 1:
            rf_probs["Autism"] = 0.85
        if rf_pred[3] == 1:
            rf_probs["Depression"] = 0.85
        if rf_pred[4] == 1:
            rf_probs["Anxiety"] = 0.85

        if sum(responses[0:7]) >= 15:
            rf_probs["ADHD"] = max(rf_probs["ADHD"], 0.65)
        if sum(responses[7:13]) >= 12:
            rf_probs["Autism"] = max(rf_probs["Autism"], 0.65)
        if sum(responses[17:22]) >= 10:
            rf_probs["Depression"] = max(rf_probs["Depression"], 0.65)
        if sum(responses[22:27]) >= 10:
            rf_probs["Anxiety"] = max(rf_probs["Anxiety"], 0.65)

        detected = [k for k, v in rf_probs.items() if v > 0.5]
        return {"rf_probs": rf_probs, "detected": detected}
    except Exception as e:
        raise HTTPException(500, str(e))


@app.post("/api/nlp/questions")
def get_nlp_questions(req: SymptomFlagsRequest):
    flags = {
        "S1": req.s1, "S2": req.s2, "S3": req.s3,
        "S4": req.s4, "S5": req.s5,
    }

    if req.custom_symptom and len(req.custom_symptom) > 3:
        pred = AI_ENGINE.predict_symptom_category(req.custom_symptom)
        if pred == "ADHD":
            flags["S1"] = True
        elif pred == "Depression":
            flags["S2"] = True
        elif pred == "Anxiety":
            flags["S3"] = True
        elif pred == "Autism":
            flags["S5"] = True

    questions = get_logic_driven_questions(flags)

    if req.custom_symptom and len(req.custom_symptom) > 3:
        questions.append({
            "text": f"You mentioned '{req.custom_symptom}'. Please describe how this affects your daily life and why you thought about mentioning it.",
            "cat": "General",
        })

    return {"questions": questions}


@app.post("/api/nlp/analyze")
def analyze_nlp(req: NLPAnalyzeRequest):
    result = AI_ENGINE.analyze_single_nlp_response(req.text, req.category)
    return {
        "diagnosis": result["diagnosis"],
        "confidence": float(result["confidence"]),
        "probs": result["probs"].tolist(),
        "analysis": result["analysis"],
    }


@app.post("/api/audio/transcribe")
async def transcribe_audio(file: UploadFile = File(...)):
    try:
        ct = (file.content_type or "audio/webm").lower()
        if "ogg" in ct:
            suffix = ".ogg"
        elif "mp4" in ct or "m4a" in ct:
            suffix = ".mp4"
        elif "wav" in ct:
            suffix = ".wav"
        else:
            suffix = ".webm"
        with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tmp:
            content = await file.read()
            tmp.write(content)
            tmp_path = tmp.name

        result = AUDIO_TOOL.process_audio(tmp_path)
        os.unlink(tmp_path)

        return {
            "text": result.get("text", ""),
            "speech_rate": result.get("speech_rate", None),
        }
    except Exception as e:
        raise HTTPException(500, str(e))


@app.post("/api/inference/fuse")
def run_meta_fusion(req: MetaFusionRequest):
    nlp_arrays = [np.array(p) for p in req.nlp_probs_list]
    result = AI_ENGINE.run_meta_fusion(req.rf_probs, nlp_arrays)

    severity = "Moderate"
    if result["diagnosis"] == "No Significant Risk":
        severity = "None"
    elif result["confidence"] > 0.8:
        severity = "High"
    elif result["confidence"] < 0.5:
        severity = "Mild"

    result["severity"] = severity
    result["suggestion"] = AI_ENGINE.get_suggestions(result["diagnosis"], severity)
    result["all_probs"] = {k: float(v) for k, v in result["all_probs"].items()}
    result["confidence"] = float(result["confidence"])
    result["method"] = "Meta-Fusion Algorithmic AI Process"

    return result


@app.post("/api/session/save")
def save_session_endpoint(req: SessionSaveRequest):
    save_session(
        user_data=req.user_data,
        rf_answers=req.rf_answers,
        rf_risk=req.rf_risk,
        nlp_data=req.nlp_data,
        final_diag=req.final_diagnosis,
    )
    return {"status": "saved"}


@app.post("/api/report/generate")
def generate_report(req: ReportRequest):
    class PDF(FPDF):
        def header(self):
            self.set_font("Arial", "B", 16)
            self.cell(0, 10, "CogniDetectAI - Clinical Screening Report", 0, 1, "C")
            self.ln(5)

        def footer(self):
            self.set_y(-15)
            self.set_font("Arial", "I", 8)
            self.cell(0, 10, f"Downloaded from CogniDetectAI System | Page {self.page_no()}", 0, 0, "C")

    pdf = PDF()
    pdf.add_page()

    u = req.user_data
    fr = req.final_result
    rf_data = req.rf_data
    nlp_data = req.nlp_data

    pdf.set_font("Arial", "", 11)
    pdf.cell(0, 8, f"Date: {datetime.now().strftime('%Y-%m-%d')}", 0, 1)
    pdf.cell(0, 8, f"Patient: {u.get('age', 'N/A')} yrs | {u.get('gender', 'N/A')}", 0, 1)
    pdf.cell(0, 8, f"Location: {_clean(u.get('country', 'N/A'))}", 0, 1)

    if u.get("med_hist"):
        pdf.set_font("Arial", "B", 11)
        pdf.cell(0, 8, f"Medical History: {_clean(', '.join(u['med_hist']))}", 0, 1)

    if u.get("symptom_data"):
        pdf.ln(3)
        pdf.set_font("Arial", "B", 11)
        pdf.cell(0, 8, "Reported Symptoms & Durations:", 0, 1)
        pdf.set_font("Arial", "", 10)
        for cat, data in u["symptom_data"].items():
            if data.get("symptoms"):
                symp_str = _clean(", ".join(data["symptoms"]))
                dur_str = _clean(data.get("duration", ""))
                pdf.multi_cell(0, 6, f"- {cat} ({dur_str}): {symp_str}")

    pdf.ln(5)
    pdf.set_font("Arial", "B", 14)
    pdf.cell(0, 10, "Final Assessment", 0, 1)
    pdf.set_font("Arial", "", 12)
    pdf.cell(0, 8, f"Primary Indication: {_clean(fr['diagnosis'])}", 0, 1)
    pdf.cell(0, 8, f"Confidence: {fr['confidence'] * 100:.1f}%", 0, 1)
    pdf.cell(0, 8, f"Severity: {_clean(fr.get('severity', 'N/A'))}", 0, 1)
    pdf.cell(0, 8, f"Method: {_clean(fr.get('method', 'N/A'))}", 0, 1)

    # ── Disorder Probability Bar Chart ────────────────────────────────────────
    pdf.ln(6)
    pdf.set_font("Arial", "B", 12)
    pdf.cell(0, 8, "Disorder Probability Overview:", 0, 1)
    pdf.ln(3)

    all_probs = fr.get("all_probs", {})
    chart_data = [
        ("ADHD",       all_probs.get("ADHD", 0.0),       (102, 126, 234)),
        ("Depression", all_probs.get("Depression", 0.0),  (239,  68,  68)),
        ("Anxiety",    all_probs.get("Anxiety", 0.0),     (245, 158,  11)),
        ("Autism",     all_probs.get("Autism", 0.0),      ( 16, 185, 129)),
    ]

    label_w   = 36   # mm — label column width
    bar_x     = 12 + label_w   # left edge of bars
    max_bar_w = 115  # mm — maximum bar width (100 % probability)
    bar_h     = 7    # mm — bar height
    bar_gap   = 5    # mm — gap between bars

    for label, prob, (r, g, b) in chart_data:
        y = pdf.get_y()

        # Disorder label
        pdf.set_font("Arial", "B", 9)
        pdf.set_text_color(60, 60, 80)
        pdf.set_xy(12, y + 1)
        pdf.cell(label_w, bar_h, label, 0, 0)

        # Background track
        pdf.set_fill_color(220, 220, 238)
        pdf.rect(bar_x, y + 1, max_bar_w, bar_h - 2, "F")

        # Filled probability bar
        if prob > 0.005:
            pdf.set_fill_color(r, g, b)
            filled_w = max(max_bar_w * float(prob), 2)
            pdf.rect(bar_x, y + 1, filled_w, bar_h - 2, "F")

        # Percentage label after the bar
        pdf.set_font("Arial", "B", 9)
        pdf.set_text_color(40, 40, 60)
        pdf.set_xy(bar_x + max_bar_w + 3, y + 1)
        pdf.cell(22, bar_h - 2, f"{float(prob) * 100:.1f}%", 0, 0)

        pdf.set_y(y + bar_h + bar_gap)

    # Reset text colour to black
    pdf.set_text_color(0, 0, 0)
    pdf.ln(4)

    # ── Suggestions ───────────────────────────────────────────────────────────
    pdf.set_font("Arial", "B", 12)
    pdf.cell(0, 10, "Suggestions:", 0, 1)
    pdf.set_font("Arial", "", 11)
    pdf.multi_cell(0, 8, _clean(fr.get("suggestion", "")))

    # Annexure 1 — Questionnaire
    pdf.add_page()
    pdf.set_font("Arial", "B", 14)
    pdf.cell(0, 10, "Annexure 1: Questionnaire Responses", 0, 1)
    pdf.set_font("Arial", "", 10)

    cfg = load_config()
    questions_list = cfg["questions"]
    answers = rf_data.get("responses", [])
    score_map = {0: "Never", 1: "Rarely", 2: "Sometimes", 3: "Often", 4: "Very Often"}

    for i, ans_score in enumerate(answers):
        ans_text = score_map.get(ans_score, str(ans_score))
        q_text = questions_list[i] if i < len(questions_list) else f"Question {i + 1}"
        pdf.set_font("Arial", "B", 9)
        pdf.multi_cell(0, 5, f"Q{i + 1}: {_clean(q_text)}")
        pdf.set_font("Arial", "", 9)
        pdf.cell(0, 5, f"Answer: {ans_text}", 0, 1)
        pdf.ln(2)

    # Annexure 2 — NLP Interview
    if nlp_data:
        pdf.add_page()
        pdf.set_font("Arial", "B", 14)
        pdf.cell(0, 10, "Annexure 2: NLP Interview Transcript", 0, 1)
        pdf.set_font("Arial", "", 10)
        for item in nlp_data:
            pdf.set_font("Arial", "B", 10)
            pdf.multi_cell(0, 6, f"Q: {_clean(item.get('q', ''))}")
            pdf.set_font("Arial", "", 10)
            pdf.multi_cell(
                0, 6,
                f"A: {_clean(item.get('a', ''))}\n(Analysis: {_clean(item.get('analysis', ''))})"
            )
            pdf.ln(4)

    pdf_bytes = pdf.output(dest="S").encode("latin-1", "replace")
    return Response(
        content=pdf_bytes,
        media_type="application/pdf",
        headers={"Content-Disposition": "attachment; filename=CogniDetectAI_Report.pdf"},
    )


# ── Admin Routes ──────────────────────────────────────────────────────────────

@app.post("/api/admin/verify-password")
def verify_password(req: AdminPasswordRequest):
    cfg = load_config()
    correct = os.environ.get("ADMIN_PASSWORD", cfg.get("admin_password", ""))
    if req.password == correct:
        return {"verified": True}
    raise HTTPException(401, "Invalid password")


@app.post("/api/admin/verify-totp")
def verify_totp(req: AdminTOTPRequest):
    cfg = load_config()
    secret = os.environ.get("TOTP_SECRET", cfg.get("totp_secret", ""))
    if pyotp.TOTP(secret).verify(req.code):
        return {"verified": True}
    raise HTTPException(401, "Invalid OTP")


@app.get("/api/admin/config")
def get_admin_config():
    return load_config()


@app.put("/api/admin/config")
def update_admin_config(req: ConfigUpdateRequest):
    with open("config.json", "w") as f:
        json.dump(req.config, f, indent=4)
    return {"status": "updated"}


@app.get("/api/admin/sessions")
def export_sessions():
    csv_data = get_all_data_as_csv()
    return Response(
        content=csv_data,
        media_type="text/csv",
        headers={"Content-Disposition": "attachment; filename=cognidetect_sessions.csv"},
    )