# api/server.py
import os
import time
from typing import Dict

from fastapi import FastAPI, UploadFile, File, Form, Request
from fastapi.responses import FileResponse, JSONResponse
from fastapi.staticfiles import StaticFiles
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel

from api.config import DEFAULT_COURSE_TOPICS, DEFAULT_MODEL
from api.syllabus_utils import extract_course_topics_from_file
from api.rag_engine import build_rag_chunks_from_file, retrieve_relevant_chunks
from api.clare_core import (
    detect_language,
    chat_with_clare,
    update_weaknesses_from_message,
    update_cognitive_state_from_message,
    render_session_status,
    export_conversation,
    summarize_conversation,
)

# ----------------------------
# Paths / Constants
# ----------------------------
API_DIR = os.path.dirname(__file__)

MODULE10_PATH = os.path.join(API_DIR, "module10_responsible_ai.pdf")
MODULE10_DOC_TYPE = "Literature Review / Paper"

# Vite build output in your repo is "web/build"
WEB_DIST = os.path.abspath(os.path.join(API_DIR, "..", "web", "build"))
WEB_INDEX = os.path.join(WEB_DIST, "index.html")
WEB_ASSETS = os.path.join(WEB_DIST, "assets")

# ----------------------------
# App
# ----------------------------
app = FastAPI(title="Clare API")

# Same-origin for Docker Space doesn't need CORS, but leaving it open helps if you later split FE/BE.
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# ----------------------------
# Static hosting (Vite build)
# ----------------------------
# Mount /assets so <script src="/assets/..."> works.
if os.path.isdir(WEB_ASSETS):
    app.mount("/assets", StaticFiles(directory=WEB_ASSETS), name="assets")

# Optional: serve other static files in build root (e.g., favicon) under /static
if os.path.isdir(WEB_DIST):
    app.mount("/static", StaticFiles(directory=WEB_DIST), name="static")


@app.get("/")
def index():
    if os.path.exists(WEB_INDEX):
        return FileResponse(WEB_INDEX)
    return JSONResponse(
        {"detail": "web/build not found. Build frontend first (web/build/index.html)."},
        status_code=500,
    )


# ----------------------------
# In-memory session store (MVP)
# ----------------------------
SESSIONS: Dict[str, Dict] = {}


def _preload_module10_chunks():
    if os.path.exists(MODULE10_PATH):
        try:
            return build_rag_chunks_from_file(MODULE10_PATH, MODULE10_DOC_TYPE) or []
        except Exception as e:
            print(f"[preload] module10 parse failed: {repr(e)}")
            return []
    return []


MODULE10_CHUNKS_CACHE = _preload_module10_chunks()


def _get_session(user_id: str) -> Dict:
    if user_id not in SESSIONS:
        SESSIONS[user_id] = {
            "user_id": user_id,
            "name": "",
            "history": [],
            "weaknesses": [],
            "cognitive_state": {"confusion": 0, "mastery": 0},
            "course_outline": DEFAULT_COURSE_TOPICS,
            # preload base reading
            "rag_chunks": list(MODULE10_CHUNKS_CACHE),
            "model_name": DEFAULT_MODEL,
        }
    return SESSIONS[user_id]


# ----------------------------
# Schemas
# ----------------------------
class LoginReq(BaseModel):
    name: str
    user_id: str


class ChatReq(BaseModel):
    user_id: str
    message: str
    learning_mode: str
    language_preference: str = "Auto"
    doc_type: str = "Syllabus"


class ExportReq(BaseModel):
    user_id: str
    learning_mode: str


class SummaryReq(BaseModel):
    user_id: str
    learning_mode: str
    language_preference: str = "Auto"


# ----------------------------
# API Routes
# ----------------------------
@app.post("/api/login")
def login(req: LoginReq):
    user_id = (req.user_id or "").strip()
    name = (req.name or "").strip()
    if not user_id or not name:
        return JSONResponse({"ok": False, "error": "Missing name/user_id"}, status_code=400)

    sess = _get_session(user_id)
    sess["name"] = name
    return {"ok": True, "user": {"name": name, "user_id": user_id}}


@app.post("/api/chat")
def chat(req: ChatReq):
    user_id = (req.user_id or "").strip()
    msg = (req.message or "").strip()

    if not user_id:
        return JSONResponse({"error": "Missing user_id"}, status_code=400)

    sess = _get_session(user_id)

    if not msg:
        return {
            "reply": "",
            "session_status_md": render_session_status(
                req.learning_mode, sess["weaknesses"], sess["cognitive_state"]
            ),
            "refs": [],
            "latency_ms": 0.0,
        }

    resolved_lang = detect_language(msg, req.language_preference)

    sess["weaknesses"] = update_weaknesses_from_message(msg, sess["weaknesses"])
    sess["cognitive_state"] = update_cognitive_state_from_message(msg, sess["cognitive_state"])

    # RAG
    rag_context_text, rag_used_chunks = retrieve_relevant_chunks(msg, sess["rag_chunks"])

    start_ts = time.time()
    try:
        answer, new_history = chat_with_clare(
            message=msg,
            history=sess["history"],
            model_name=sess["model_name"],
            language_preference=resolved_lang,
            learning_mode=req.learning_mode,
            doc_type=req.doc_type,
            course_outline=sess["course_outline"],
            weaknesses=sess["weaknesses"],
            cognitive_state=sess["cognitive_state"],
            rag_context=rag_context_text,
        )
    except Exception as e:
        print(f"[chat] error: {repr(e)}")
        return JSONResponse({"error": f"chat failed: {repr(e)}"}, status_code=500)

    latency_ms = (time.time() - start_ts) * 1000.0
    sess["history"] = new_history

    refs = [
        {"source_file": c.get("source_file"), "section": c.get("section")}
        for c in (rag_used_chunks or [])
    ]

    return {
        "reply": answer,
        "session_status_md": render_session_status(
            req.learning_mode, sess["weaknesses"], sess["cognitive_state"]
        ),
        "refs": refs,
        "latency_ms": latency_ms,
    }


@app.post("/api/upload")
async def upload(
    user_id: str = Form(...),
    doc_type: str = Form(...),
    file: UploadFile = File(...),
):
    user_id = (user_id or "").strip()
    doc_type = (doc_type or "").strip()

    if not user_id:
        return JSONResponse({"ok": False, "error": "Missing user_id"}, status_code=400)
    if not file or not file.filename:
        return JSONResponse({"ok": False, "error": "Missing file"}, status_code=400)

    sess = _get_session(user_id)

    # Save to /tmp (sanitize filename)
    safe_name = os.path.basename(file.filename).replace("..", "_")
    tmp_path = os.path.join("/tmp", safe_name)

    content = await file.read()
    with open(tmp_path, "wb") as f:
        f.write(content)

    # Update topics only for syllabus
    if doc_type == "Syllabus":
        class _F:
            pass

        fo = _F()
        fo.name = tmp_path
        try:
            sess["course_outline"] = extract_course_topics_from_file(fo, doc_type)
        except Exception as e:
            print(f"[upload] syllabus parse error: {repr(e)}")

    # Update rag chunks for any doc
    try:
        new_chunks = build_rag_chunks_from_file(tmp_path, doc_type) or []
        sess["rag_chunks"] = (sess["rag_chunks"] or []) + new_chunks
    except Exception as e:
        print(f"[upload] rag build error: {repr(e)}")
        new_chunks = []

    status_md = f"✅ Loaded base reading + uploaded {doc_type} file."
    return {"ok": True, "added_chunks": len(new_chunks), "status_md": status_md}


@app.post("/api/export")
def api_export(req: ExportReq):
    user_id = (req.user_id or "").strip()
    if not user_id:
        return JSONResponse({"error": "Missing user_id"}, status_code=400)

    sess = _get_session(user_id)
    md = export_conversation(
        sess["history"],
        sess["course_outline"],
        req.learning_mode,
        sess["weaknesses"],
        sess["cognitive_state"],
    )
    return {"markdown": md}


@app.post("/api/summary")
def api_summary(req: SummaryReq):
    user_id = (req.user_id or "").strip()
    if not user_id:
        return JSONResponse({"error": "Missing user_id"}, status_code=400)

    sess = _get_session(user_id)
    md = summarize_conversation(
        sess["history"],
        sess["course_outline"],
        sess["weaknesses"],
        sess["cognitive_state"],
        sess["model_name"],
        req.language_preference,
    )
    return {"markdown": md}


@app.get("/api/memoryline")
def memoryline(user_id: str):
    _ = _get_session((user_id or "").strip())
    # v1: 写死也没问题；前端只渲染
    return {"next_review_label": "T+7", "progress_pct": 0.4}


# ----------------------------
# SPA Fallback (important!)
# ----------------------------
# If user refreshes /some/route, FE router needs index.html.
@app.get("/{full_path:path}")
def spa_fallback(full_path: str, request: Request):
    # Do not hijack API/static paths
    if (
        full_path.startswith("api/")
        or full_path.startswith("assets/")
        or full_path.startswith("static/")
    ):
        return JSONResponse({"detail": "Not Found"}, status_code=404)

    if os.path.exists(WEB_INDEX):
        return FileResponse(WEB_INDEX)

    return JSONResponse(
        {"detail": "web/build not found. Build frontend first (web/build/index.html)."},
        status_code=500,
    )