"""
FinSim — Full FastAPI Application
Routes: Auth, Scenarios, Generator, Chat, MCQ Quiz, Akinator 2.0
"""

import json
import uuid
import bcrypt
from datetime import datetime
from fastapi import FastAPI, HTTPException, Request, Depends, UploadFile, File, Form
from fastapi.middleware.cors import CORSMiddleware
from fastapi.staticfiles import StaticFiles
from fastapi.responses import FileResponse, JSONResponse
from pydantic import BaseModel, Field
import psycopg2.extras
from typing import Optional
from dotenv import load_dotenv

import services.database as db
from services.zscore_engine import identify_events
from services.scenario_gen import generate_scenarios
from services.chat_agent import get_chat_response
from services.akinator import get_akinator_response
from models import GenerateRequest, ChatRequest

load_dotenv()

# Initialize DB tables
try:
    db.init_db()
except Exception as e:
    print(f"Warning: Could not initialize DB on startup. Error: {e}")

app = FastAPI(title="FinSim", version="2.0.0")

# ── CORS — allow Flutter web (and any dev origin) ────────────────
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# ── In-memory session store ──────────────────────────────────────
_tokens: dict[str, dict] = {}  # token -> {user_id, username, role, ...}


# ── Request Models ───────────────────────────────────────────────
class AkinatorRequest(BaseModel):
    session_id: str = Field(default="default")
    message: str
    panel_mode: bool = Field(default=False)
    knowledge_base_id: Optional[str] = None

class LoginRequest(BaseModel):
    email: str
    password: str

class RegisterRequest(BaseModel):
    username: str
    email: str
    password: str

class AnswerRequest(BaseModel):
    scenario_id: str
    selected_answer: str

class StartQuizRequest(BaseModel):
    num_questions: int = Field(default=5, ge=1, le=20)
    difficulty: Optional[str] = None
    category: Optional[str] = None


# ── Auth Helper ──────────────────────────────────────────────────
def get_current_user(request: Request) -> dict:
    auth = request.headers.get("Authorization", "")
    if not auth.startswith("Bearer "):
        raise HTTPException(status_code=401, detail="Not authenticated")
    token = auth[7:]
    user = _tokens.get(token)
    if not user:
        raise HTTPException(status_code=401, detail="Invalid or expired token")
    return user


# ── Serve UI ─────────────────────────────────────────────────────
app.mount("/static", StaticFiles(directory="static"), name="static")

@app.get("/")
def get_ui():
    return FileResponse("static/index.html")


# ══════════════════════════════════════════════════════════════════
# AUTH ENDPOINTS
# ══════════════════════════════════════════════════════════════════

@app.post("/api/login")
def login(req: LoginRequest):
    from services.database import get_connection
    with get_connection() as conn:
        cursor = conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor)
        cursor.execute("SELECT * FROM users WHERE email = %s", (req.email,))
        user = cursor.fetchone()

    if not user:
        raise HTTPException(status_code=401, detail="Invalid email or password")

    stored_pw = user["password"]
    # Support both bcrypt and plain-text legacy passwords
    valid = False
    if stored_pw.startswith("$2"):
        try:
            valid = bcrypt.checkpw(req.password.encode(), stored_pw.encode())
        except Exception:
            valid = False
    else:
        valid = (req.password == stored_pw)

    if not valid:
        raise HTTPException(status_code=401, detail="Invalid email or password")

    token = uuid.uuid4().hex
    _tokens[token] = {
        "user_id": user["id"],
        "username": user["username"],
        "email": user["email"],
        "role": user["role"],
        "avatar": user["avatar"],
        "total_score": user["total_score"],
        "quizzes_taken": user["quizzes_taken"],
    }
    return {"token": token, "user": _tokens[token]}


@app.post("/api/register")
def register(req: RegisterRequest):
    from services.database import get_connection
    hashed = bcrypt.hashpw(req.password.encode(), bcrypt.gensalt()).decode()
    try:
        with get_connection() as conn:
            cursor = conn.cursor()
            cursor.execute(
                "INSERT INTO users (username, email, password, role) VALUES (%s, %s, %s, 'user') RETURNING id",
                (req.username, req.email, hashed)
            )
            user_id = cursor.fetchone()[0]
    except Exception as e:
        if "Duplicate" in str(e):
            raise HTTPException(status_code=409, detail="Username or email already exists")
        raise HTTPException(status_code=500, detail=str(e))

    token = uuid.uuid4().hex
    _tokens[token] = {
        "user_id": user_id,
        "username": req.username,
        "email": req.email,
        "role": "user",
        "avatar": None,
        "total_score": 0,
        "quizzes_taken": 0,
    }
    return {"token": token, "user": _tokens[token]}


@app.get("/api/me")
def get_me(user: dict = Depends(get_current_user)):
    # Refresh from DB
    from services.database import get_connection
    with get_connection() as conn:
        cursor = conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor)
        cursor.execute("SELECT total_score, quizzes_taken FROM users WHERE id = %s", (user["user_id"],))
        fresh = cursor.fetchone()
    if fresh:
        user["total_score"] = fresh["total_score"]
        user["quizzes_taken"] = fresh["quizzes_taken"]
    return {"user": user}


@app.post("/api/logout")
def logout(request: Request):
    auth = request.headers.get("Authorization", "")
    if auth.startswith("Bearer "):
        token = auth[7:]
        _tokens.pop(token, None)
    return {"ok": True}


# ══════════════════════════════════════════════════════════════════
# DATA INGESTION & KNOWLEDGE BASES (RAG)
# ══════════════════════════════════════════════════════════════════

@app.get("/api/rag/kbs")
def get_kbs():
    from services.rag_engine import list_knowledge_bases
    try:
        kbs = list_knowledge_bases()
        return {"status": "success", "knowledge_bases": [kb["name"] for kb in kbs]}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.post("/api/rag/upload")
def upload_to_kb(
    kb_name: str = Form(...),
    urls: Optional[str] = Form(None),
    files: Optional[list[UploadFile]] = File(None)
):
    import os
    import tempfile
    from services.rag_engine import ingest_document
    
    try:
        total_chunks = 0
        if urls:
            url_list = [u.strip() for u in urls.split(",") if u.strip()]
            for u in url_list:
                total_chunks += ingest_document(kb_name, u, "url")
                
        if files:
            for file in files:
                ext = file.filename.split(".")[-1].lower() if "." in file.filename else "txt"
                if ext in ["doc", "docx"]:
                    ext = "docx"
                elif ext not in ["pdf", "txt", "json", "docx"]:
                    raise HTTPException(status_code=400, detail=f"Unsupported file extension: {ext}. Allowed: PDF, TXT, JSON, DOCX.")
                
                suffix = f".{ext}"
                with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tmp:
                    content = file.file.read()
                    tmp.write(content)
                    source = tmp.name
                    
                total_chunks += ingest_document(kb_name, source, ext)
                
                # Cleanup temp file if necessary
                if os.path.exists(source):
                    os.remove(source)
                    
        if not urls and not files:
            raise HTTPException(status_code=400, detail="Must provide at least one URL or one file.")
            
        return {"status": "success", "message": f"Ingested {total_chunks} chunks into Knowledge Base '{kb_name}'"}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


# ══════════════════════════════════════════════════════════════════
# SCENARIO GENERATOR
# ══════════════════════════════════════════════════════════════════

@app.post("/api/generate")
def api_generate_scenarios(req: GenerateRequest, user: dict = Depends(get_current_user)):
    try:
        zscore_result = identify_events(
            symbol=req.stock_symbol,
            window=req.zscore_window,
            trigger_min=req.zscore_trigger_min,
            trigger_max=req.zscore_trigger_max
        )
        if zscore_result.total_events == 0:
            return {"status": "success", "message": f"No significant events found for {req.stock_symbol}."}

        generation_result = generate_scenarios(req.stock_symbol, zscore_result, req.knowledge_base_id)
        inserted_count = db.insert_all_scenarios(generation_result)

        return {
            "status": "success",
            "stock_symbol": req.stock_symbol or "Mixed Basket",
            "events_analyzed": zscore_result.total_events,
            "scenarios_generated": len(generation_result.scenarios),
            "scenarios_saved_to_db": inserted_count
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


# Keep old endpoint for backwards compat
@app.post("/generate")
def api_generate_scenarios_legacy(req: GenerateRequest):
    try:
        zscore_result = identify_events(
            symbol=req.stock_symbol, window=req.zscore_window,
            trigger_min=req.zscore_trigger_min, trigger_max=req.zscore_trigger_max
        )
        if zscore_result.total_events == 0:
            return {"status": "success", "message": f"No significant events found for {req.stock_symbol}."}
        generation_result = generate_scenarios(req.stock_symbol, zscore_result)
        inserted_count = db.insert_all_scenarios(generation_result)
        return {
            "status": "success", "stock_symbol": req.stock_symbol or "Mixed Basket",
            "events_analyzed": zscore_result.total_events,
            "scenarios_generated": len(generation_result.scenarios),
            "scenarios_saved_to_db": inserted_count
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


# ══════════════════════════════════════════════════════════════════
# CHAT
# ══════════════════════════════════════════════════════════════════

@app.post("/api/chat")
def api_chat(req: ChatRequest, user: dict = Depends(get_current_user)):
    try:
        response = get_chat_response(req.session_id, req.message, req.knowledge_base_id)
        if isinstance(response, dict):
            if response.get("rate_limited"):
                return response
            response["session_id"] = req.session_id
            return response
        return {"session_id": req.session_id, "reply": response}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.post("/chat")
def api_chat_legacy(req: ChatRequest):
    try:
        response = get_chat_response(req.session_id, req.message)
        if isinstance(response, dict):
            if response.get("rate_limited"):
                return response
            response["session_id"] = req.session_id
            return response
        return {"session_id": req.session_id, "reply": response}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


# ══════════════════════════════════════════════════════════════════
# AKINATOR 2.0 — MULTI-AGENT PREDICTION ENGINE
# ══════════════════════════════════════════════════════════════════

@app.post("/api/akinator")
def api_akinator(req: AkinatorRequest, user: dict = Depends(get_current_user)):
    try:
        response = get_akinator_response(req.session_id, req.message, req.panel_mode, req.knowledge_base_id)
        if isinstance(response, dict) and response.get("rate_limited"):
            return response
        # Rich response from LangGraph pipeline
        if isinstance(response, dict):
            return {
                "session_id": req.session_id,
                "reply": response.get("reply", ""),
                "confidence_score": response.get("confidence_score", 0),
                "news_sentiment": response.get("news_sentiment", {}),
                "jargon_definitions": response.get("jargon_definitions", {}),
                "memo": response.get("memo", ""),
                "whatif_result": response.get("whatif_result", ""),
                "critique_note": response.get("critique_note", ""),
            }
        return {"session_id": req.session_id, "reply": response}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))



# ══════════════════════════════════════════════════════════════════
# SCENARIOS BROWSING
# ══════════════════════════════════════════════════════════════════

@app.get("/api/scenarios")
def get_scenarios(page: int = 1, limit: int = 12, difficulty: str = "", category: str = "", risk: str = ""):
    from services.database import get_connection
    offset = (page - 1) * limit
    where_clauses = ["is_active = 1"]
    params = []

    if difficulty:
        where_clauses.append("difficulty = %s")
        params.append(difficulty)
    if category:
        where_clauses.append("category = %s")
        params.append(category)
    if risk:
        where_clauses.append("risk_level = %s")
        params.append(risk)

    where_sql = " AND ".join(where_clauses)

    with get_connection() as conn:
        cursor = conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor)
        cursor.execute(f"SELECT COUNT(*) as total FROM scenarios WHERE {where_sql}", params)
        total = cursor.fetchone()["total"]

        cursor.execute(
            f"""SELECT id, title, short_description, event_type, difficulty, category, risk_level, created_at
                FROM scenarios WHERE {where_sql}
                ORDER BY created_at DESC LIMIT %s OFFSET %s""",
            params + [limit, offset]
        )
        scenarios = cursor.fetchall()

    # Serialize datetimes
    for s in scenarios:
        if s.get("created_at"):
            s["created_at"] = s["created_at"].isoformat()

    return {"scenarios": scenarios, "total": total, "page": page, "pages": (total + limit - 1) // limit}


@app.get("/api/scenarios/{scenario_id}")
def get_scenario_detail(scenario_id: str):
    from services.database import get_connection
    with get_connection() as conn:
        cursor = conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor)
        cursor.execute("SELECT * FROM scenarios WHERE id = %s", (scenario_id,))
        scenario = cursor.fetchone()
    if not scenario:
        raise HTTPException(status_code=404, detail="Scenario not found")
    if scenario.get("created_at"):
        scenario["created_at"] = scenario["created_at"].isoformat()
    if scenario.get("updated_at"):
        scenario["updated_at"] = scenario["updated_at"].isoformat()
    if scenario.get("givens_table") and isinstance(scenario["givens_table"], str):
        try:
            scenario["givens_table"] = json.loads(scenario["givens_table"])
        except:
            pass
    return scenario


# ══════════════════════════════════════════════════════════════════
# MCQ QUIZ SYSTEM
# ══════════════════════════════════════════════════════════════════

@app.post("/api/quiz/start")
def start_quiz(req: StartQuizRequest, user: dict = Depends(get_current_user)):
    from services.database import get_connection
    where_clauses = ["is_active = 1"]
    params = []

    if req.difficulty:
        where_clauses.append("difficulty = %s")
        params.append(req.difficulty)
    if req.category:
        where_clauses.append("category = %s")
        params.append(req.category)

    where_sql = " AND ".join(where_clauses)

    with get_connection() as conn:
        cursor = conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor)

        # Get random scenarios
        cursor.execute(
            f"SELECT id, title, short_description, scenario_paragraph, givens_table, "
            f"best_answer, best_answer_rationale, "
            f"other_option1, other_option1_exp, other_option2, other_option2_exp, "
            f"other_option3, other_option3_exp, event_type, difficulty, category, risk_level "
            f"FROM scenarios WHERE {where_sql} ORDER BY RAND() LIMIT %s",
            params + [req.num_questions]
        )
        scenarios = cursor.fetchall()

        if not scenarios:
            raise HTTPException(status_code=404, detail="No scenarios found matching your criteria")

        # Create quiz
        cursor.execute(
            "INSERT INTO quizzes (title, description, time_limit, passing_score, created_by) VALUES (%s, %s, %s, %s, %s) RETURNING id",
            (f"Quick Quiz - {datetime.now().strftime('%b %d, %H:%M')}", f"{len(scenarios)} questions", 0, 60, user["user_id"])
        )
        quiz_id = cursor.fetchone()['id']

        # Link scenarios
        for i, s in enumerate(scenarios):
            cursor.execute(
                "INSERT INTO quiz_scenarios (quiz_id, scenario_id, question_order) VALUES (%s, %s, %s)",
                (quiz_id, s["id"], i + 1)
            )

        # Create attempt
        cursor.execute(
            "INSERT INTO quiz_attempts (user_id, quiz_id, total_questions, status) VALUES (%s, %s, %s, 'in_progress') RETURNING id",
            (user["user_id"], quiz_id, len(scenarios))
        )
        attempt_id = cursor.fetchone()['id']

    # Process scenarios for frontend
    questions = []
    for s in scenarios:
        givens = s.get("givens_table")
        if givens and isinstance(givens, str):
            try:
                givens = json.loads(givens)
            except:
                givens = {}

        # Build shuffled options
        import random
        options = [
            {"text": s["best_answer"], "key": "correct"},
            {"text": s["other_option1"], "key": "wrong1"},
            {"text": s["other_option2"], "key": "wrong2"},
            {"text": s["other_option3"], "key": "wrong3"},
        ]
        random.shuffle(options)
        labels = ["A", "B", "C", "D"]
        labeled_options = []
        correct_label = ""
        for j, opt in enumerate(options):
            labeled_options.append({"label": labels[j], "text": opt["text"]})
            if opt["key"] == "correct":
                correct_label = labels[j]

        questions.append({
            "scenario_id": s["id"],
            "title": s["title"],
            "short_description": s["short_description"],
            "scenario_paragraph": s["scenario_paragraph"],
            "givens_table": givens,
            "risk_level": s["risk_level"],
            "difficulty": s["difficulty"],
            "event_type": s["event_type"],
            "options": labeled_options,
            "correct_label": correct_label,
            "best_answer_rationale": s["best_answer_rationale"],
            "other_explanations": {
                s["other_option1"]: s["other_option1_exp"],
                s["other_option2"]: s["other_option2_exp"],
                s["other_option3"]: s["other_option3_exp"],
            }
        })

    return {
        "quiz_id": quiz_id,
        "attempt_id": attempt_id,
        "total_questions": len(questions),
        "questions": questions,
    }


@app.post("/api/quiz/{attempt_id}/submit")
def submit_quiz(attempt_id: int, answers: list[AnswerRequest], user: dict = Depends(get_current_user)):
    from services.database import get_connection

    with get_connection() as conn:
        cursor = conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor)

        # Verify attempt belongs to user
        cursor.execute("SELECT * FROM quiz_attempts WHERE id = %s AND user_id = %s", (attempt_id, user["user_id"]))
        attempt = cursor.fetchone()
        if not attempt:
            raise HTTPException(status_code=404, detail="Quiz attempt not found")

        correct_count = 0
        results = []

        for ans in answers:
            # Get scenario
            cursor.execute("SELECT best_answer, best_answer_rationale, other_option1, other_option1_exp, other_option2, other_option2_exp, other_option3, other_option3_exp FROM scenarios WHERE id = %s", (ans.scenario_id,))
            scenario = cursor.fetchone()
            if not scenario:
                continue

            is_correct = ans.selected_answer.strip() == scenario["best_answer"].strip()
            if is_correct:
                correct_count += 1

            # Record history
            cursor.execute(
                "INSERT INTO user_scenario_history (user_id, scenario_id, selected_answer, is_correct, attempt_id) VALUES (%s, %s, %s, %s, %s)",
                (user["user_id"], ans.scenario_id, ans.selected_answer, is_correct, attempt_id)
            )

            # Build explanation
            explanation = ""
            if is_correct:
                explanation = scenario["best_answer_rationale"]
            else:
                # Find what they picked
                for opt_key in ["other_option1", "other_option2", "other_option3"]:
                    if ans.selected_answer.strip() == scenario[opt_key].strip():
                        explanation = scenario[opt_key + "_exp"]
                        break

            results.append({
                "scenario_id": ans.scenario_id,
                "selected_answer": ans.selected_answer,
                "correct_answer": scenario["best_answer"],
                "is_correct": is_correct,
                "explanation": explanation,
                "correct_rationale": scenario["best_answer_rationale"],
            })

        # Update attempt
        score = int((correct_count / len(answers)) * 100) if answers else 0
        cursor.execute(
            "UPDATE quiz_attempts SET score = %s, correct_answers = %s, status = 'completed', completed_at = NOW(), answers = %s WHERE id = %s",
            (score, correct_count, json.dumps([r for r in results]), attempt_id)
        )

        # Update user stats
        cursor.execute(
            "UPDATE users SET total_score = total_score + %s, quizzes_taken = quizzes_taken + 1 WHERE id = %s",
            (score, user["user_id"])
        )

    return {
        "score": score,
        "correct": correct_count,
        "total": len(answers),
        "percentage": score,
        "results": results,
    }


# ══════════════════════════════════════════════════════════════════
# USER STATS / HISTORY
# ══════════════════════════════════════════════════════════════════

@app.get("/api/stats")
def get_user_stats(user: dict = Depends(get_current_user)):
    from services.database import get_connection
    with get_connection() as conn:
        cursor = conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor)

        # Get overall stats
        cursor.execute("SELECT total_score, quizzes_taken FROM users WHERE id = %s", (user["user_id"],))
        stats = cursor.fetchone()

        # Get recent attempts
        cursor.execute(
            """SELECT qa.id, qa.score, qa.correct_answers, qa.total_questions, qa.status, qa.started_at, qa.completed_at
               FROM quiz_attempts qa
               WHERE qa.user_id = %s
               ORDER BY qa.started_at DESC LIMIT 10""",
            (user["user_id"],)
        )
        attempts = cursor.fetchall()
        for a in attempts:
            if a.get("started_at"):
                a["started_at"] = a["started_at"].isoformat()
            if a.get("completed_at"):
                a["completed_at"] = a["completed_at"].isoformat()

        # Get accuracy
        cursor.execute(
            "SELECT COUNT(*) as total, SUM(is_correct::int) as correct FROM user_scenario_history WHERE user_id = %s",
            (user["user_id"],)
        )
        accuracy = cursor.fetchone()

        # Get category breakdown
        cursor.execute(
            """SELECT s.category, COUNT(*) as attempts, SUM(ush.is_correct::int) as correct
               FROM user_scenario_history ush
               JOIN scenarios s ON ush.scenario_id = s.id
               WHERE ush.user_id = %s
               GROUP BY s.category""",
            (user["user_id"],)
        )
        categories = cursor.fetchall()

    return {
        "total_score": stats["total_score"] if stats else 0,
        "quizzes_taken": stats["quizzes_taken"] if stats else 0,
        "avg_score": int(stats["total_score"] / stats["quizzes_taken"]) if stats and stats["quizzes_taken"] > 0 else 0,
        "total_questions_answered": accuracy["total"] if accuracy else 0,
        "total_correct": int(accuracy["correct"]) if accuracy and accuracy["correct"] else 0,
        "accuracy": round((int(accuracy["correct"]) / accuracy["total"]) * 100, 1) if accuracy and accuracy["total"] and accuracy["correct"] else 0,
        "recent_attempts": attempts,
        "categories": categories,
    }


@app.get("/api/leaderboard")
def get_leaderboard():
    from services.database import get_connection
    with get_connection() as conn:
        cursor = conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor)
        cursor.execute(
            """SELECT username, total_score, quizzes_taken, avatar,
                      CASE WHEN quizzes_taken > 0 THEN ROUND(total_score / quizzes_taken) ELSE 0 END as avg_score
               FROM users
               WHERE quizzes_taken > 0
               ORDER BY total_score DESC LIMIT 20"""
        )
        leaders = cursor.fetchall()
    return {"leaderboard": leaders}


# ══════════════════════════════════════════════════════════════════
# FILTERS
# ══════════════════════════════════════════════════════════════════

@app.get("/api/filters")
def get_filters():
    from services.database import get_connection
    with get_connection() as conn:
        cursor = conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor)
        cursor.execute("SELECT DISTINCT category FROM scenarios WHERE is_active = 1 AND category IS NOT NULL")
        categories = [r["category"] for r in cursor.fetchall()]
        cursor.execute("SELECT DISTINCT difficulty FROM scenarios WHERE is_active = 1 AND difficulty IS NOT NULL")
        difficulties = [r["difficulty"] for r in cursor.fetchall()]
        cursor.execute("SELECT DISTINCT risk_level FROM scenarios WHERE is_active = 1 AND risk_level IS NOT NULL")
        risks = [r["risk_level"] for r in cursor.fetchall()]
    return {"categories": categories, "difficulties": difficulties, "risk_levels": risks}
