"""
api_server.py — FastAPI JSON API + Streamlit reverse proxy.

This is the NEW entrypoint for the HF Space, replacing `exec streamlit run ui.py`
in start.sh. It:
  - Serves /api/health and /api/generate directly (Chrome extension targets these)
  - Launches Streamlit as a subprocess on port 8501
  - Reverse-proxies all other HTTP traffic to Streamlit via httpx
  - Proxies the Streamlit WebSocket (/_stcore/stream) via the websockets library

No ATS or LLM logic is implemented here — only wiring calls into existing src/ modules.
"""

import asyncio
import base64
import hashlib
import hmac
import io
import os
import shutil
import subprocess
import sys
import tempfile

import httpx
import uvicorn
from fastapi import FastAPI, Header, HTTPException, UploadFile, Form, Request, WebSocket, WebSocketDisconnect
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse
from starlette.responses import Response as StarletteResponse

# ── App init ──────────────────────────────────────────────────────────────────
app = FastAPI(title="ATS Resume API", version="1.0")

# ── CORS middleware ───────────────────────────────────────────────────────────
# Dev: allow all origins (*). Override CORS_ORIGINS with the published
# chrome-extension://<ID> value after publishing to the Chrome Web Store.
ALLOWED_ORIGINS = os.getenv("CORS_ORIGINS", "*").split(",")

app.add_middleware(
    CORSMiddleware,
    allow_origins=ALLOWED_ORIGINS,
    allow_methods=["POST", "GET", "OPTIONS"],
    allow_headers=["X-Api-Token", "Content-Type"],
    expose_headers=["Content-Disposition"],
    allow_credentials=False,  # credentials=True requires specific origin, not *
)

# ── Auth token guard (constant-time comparison) ───────────────────────────────
_SECRET = os.getenv("API_SECRET_TOKEN", "")


def _check_token(token: str | None) -> None:
    """Raise HTTP 401 if API_SECRET_TOKEN is set and the header doesn't match."""
    if not _SECRET:
        return  # No token configured → open (development mode only)
    if not token or not hmac.compare_digest(token.strip(), _SECRET.strip()):
        raise HTTPException(status_code=401, detail="auth_failed")


# ── In-memory resume parse cache (single-user, per-process) ──────────────────
# sha256_hex -> parsed Resume model (avoids redundant pdfplumber parsing)
_RESUME_CACHE: dict[str, object] = {}


# ── Core generation helper ────────────────────────────────────────────────────
def generate_resume_for_api(
    pdf_bytes: bytes,
    jd_text: str,
    job_title: str,
    company: str,
) -> tuple[str, dict]:
    """
    Run the full ATS pipeline for an uploaded resume PDF + JD text.

    Returns (docx_filepath, v2_report_dict). Raises on hard failure.
    The caller is responsible for cleaning up the returned filepath's parent
    temp directory after base64-encoding the file.
    """
    from src.resume_parser_v2 import parse_resume_pdf
    from src.resume_customizer import ResumeCustomizer
    from src.llm_client import LLMClient
    from src.providers import build_provider_chain
    import config

    # 1. Parse the uploaded PDF — use in-memory cache keyed by SHA-256 to
    #    avoid redundant pdfplumber parsing across repeated requests.
    pdf_hash = hashlib.sha256(pdf_bytes).hexdigest()
    if pdf_hash in _RESUME_CACHE:
        base_resume = _RESUME_CACHE[pdf_hash]
    else:
        with tempfile.NamedTemporaryFile(suffix=".pdf", delete=False) as tmp_pdf:
            tmp_pdf.write(pdf_bytes)
            tmp_pdf_path = tmp_pdf.name
        try:
            base_resume = parse_resume_pdf(tmp_pdf_path)
        finally:
            try:
                os.unlink(tmp_pdf_path)
            except OSError:
                pass
        _RESUME_CACHE[pdf_hash] = base_resume

    # 2. Build ResumeCustomizer pointed at a fresh temp output dir.
    tmp_out = tempfile.mkdtemp()
    llm = LLMClient()
    provider_chain = build_provider_chain(llm)

    # Use the first tailor-capable model config for the customizer constructor
    # (same logic as the batch pipeline).
    primary_cfg = next(
        (m for m in config.ASSESSMENT_MODELS if m.get("tailor") and m.get("api_key")),
        None,
    )
    customizer = ResumeCustomizer(llm, base_resume.to_flat_text(), tmp_out,
                                  fast_model_cfg=primary_cfg)

    # 3. Build a minimal job dict (same shape the batch pipeline uses).
    job: dict = {
        "title": job_title,
        "company": company,
        "description": jd_text[:16000],  # truncate to ~4000 words to avoid LLM token overflow
        "ats_keywords": "",
        "_raw_assessment": {},
    }

    # 4. Call _run_provider_chain (full provider fallback + READY gate + best-attempt
    #    selection). Always pass base_resume_override to skip the hardcoded disk path.
    filepath = tempfile.mktemp(suffix=".docx", dir=tmp_out)
    result_path = customizer._run_provider_chain(
        job, filepath, provider_chain, base_resume_override=base_resume
    )

    v2_report = job.get("_v2_report", {}) or {}
    return result_path, v2_report


# ── /api/health ───────────────────────────────────────────────────────────────
@app.get("/api/health")
async def health():
    """Health check. Returns 200 with {status: ok} when the server is up."""
    return {
        "status": "ok",
        "version": "1.0",
        "streamlit_pid": _streamlit_proc.pid if _streamlit_proc else None,
    }


# ── /api/generate ─────────────────────────────────────────────────────────────
@app.post("/api/generate")
async def generate(
    jd_text: str = Form(...),
    job_title: str = Form(""),
    company: str = Form(""),
    resume: UploadFile = None,
    x_api_token: str = Header(None),
):
    """
    Generate a tailored resume DOCX (and optional PDF) for the supplied JD.

    Request: multipart/form-data
      - jd_text   (required) — full job description text
      - job_title (optional) — for recruiter pitch header
      - company   (optional) — for recruiter pitch header
      - resume    (required) — PDF file bytes of the candidate's resume
      Header: X-Api-Token — must match API_SECRET_TOKEN env var (when set)

    Response: JSON with docx_b64, pdf_b64, status, scores, quality_flag, etc.
    """
    _check_token(x_api_token)

    if resume is None:
        raise HTTPException(status_code=422, detail="resume_required")

    pdf_bytes = await resume.read()
    if not pdf_bytes:
        raise HTTPException(status_code=422, detail="resume_required")

    tmp_out: str | None = None
    result_path: str | None = None

    try:
        # Run the blocking pipeline in a thread pool to keep the event loop free.
        loop = asyncio.get_event_loop()
        result_path, v2_report = await loop.run_in_executor(
            None, generate_resume_for_api, pdf_bytes, jd_text, job_title, company
        )

        if not result_path:
            return JSONResponse(
                {"error": "generation_failed", "detail": "pipeline returned no output"},
                status_code=500,
            )

        # Encode DOCX as base64.
        with open(result_path, "rb") as f:
            docx_bytes = f.read()
        docx_b64 = base64.b64encode(docx_bytes).decode("ascii")

        # Look for an optional PDF sidecar (same stem, .pdf) from pdf_writer.
        pdf_b64: str | None = None
        pdf_sidecar = os.path.splitext(result_path)[0] + ".pdf"
        if os.path.exists(pdf_sidecar):
            with open(pdf_sidecar, "rb") as f:
                pdf_b64 = base64.b64encode(f.read()).decode("ascii")

        # Map report keys to the response contract (see 06-RESEARCH.md Section 2).
        est_scores = v2_report.get("estimated_scores", {}) or {}
        response_payload = {
            "status": v2_report.get("status", "NEEDS_USER_INPUT"),
            "download_allowed": bool(v2_report.get("download_allowed", False)),
            "scores": {
                "jd_match": est_scores.get("jd_match"),
                "ats_readability": est_scores.get("ats_readability"),
                "independent_jd_match": v2_report.get("independent_jd_match"),
            },
            "quality_flag": v2_report.get("quality_flag", ""),
            "review_terms": v2_report.get("review_terms", []),
            "high_risk_terms": v2_report.get("high_risk_terms", []),
            "docx_b64": docx_b64,
            "pdf_b64": pdf_b64,
        }
        return JSONResponse(response_payload)

    except HTTPException:
        raise
    except Exception as exc:
        return JSONResponse(
            {"error": "generation_failed", "detail": str(exc)},
            status_code=500,
        )
    finally:
        # Clean up the temp directory that generate_resume_for_api created.
        if result_path:
            parent = os.path.dirname(result_path)
            shutil.rmtree(parent, ignore_errors=True)


# ── Streamlit subprocess ──────────────────────────────────────────────────────
_streamlit_proc: subprocess.Popen | None = None
_STREAMLIT_PORT = 8501


def start_streamlit() -> None:
    """Launch Streamlit on an internal port so the proxy can reach it."""
    global _streamlit_proc
    _streamlit_proc = subprocess.Popen(
        [
            sys.executable,
            "-m",
            "streamlit",
            "run",
            "ui.py",
            "--server.port",
            str(_STREAMLIT_PORT),
            "--server.address",
            "127.0.0.1",
            "--server.headless",
            "true",
        ]
    )


# ── WebSocket proxy for /_stcore/stream (Streamlit live updates) ──────────────
@app.websocket("/_stcore/stream")
async def ws_proxy(websocket: WebSocket):
    """
    Bidirectional WebSocket proxy to Streamlit's internal port.
    httpx does NOT handle WebSocket upgrades — this route handles that separately.
    """
    import websockets as _ws
    import websockets.exceptions as _ws_exc

    await websocket.accept()
    upstream_url = (
        f"ws://127.0.0.1:{_STREAMLIT_PORT}/_stcore/stream"
        f"?{websocket.scope.get('query_string', b'').decode()}"
    )

    try:
        async with _ws.connect(upstream_url) as upstream:

            async def client_to_upstream():
                try:
                    while True:
                        try:
                            data = await websocket.receive_bytes()
                            await upstream.send(data)
                        except Exception:
                            try:
                                text = await websocket.receive_text()
                                await upstream.send(text)
                            except Exception:
                                break
                except WebSocketDisconnect:
                    pass

            async def upstream_to_client():
                try:
                    async for message in upstream:
                        if isinstance(message, bytes):
                            await websocket.send_bytes(message)
                        else:
                            await websocket.send_text(message)
                except _ws_exc.ConnectionClosed:
                    pass

            await asyncio.gather(client_to_upstream(), upstream_to_client())

    except Exception:
        pass
    finally:
        try:
            await websocket.close()
        except Exception:
            pass


# ── HTTP reverse proxy catch-all → Streamlit ─────────────────────────────────
# FastAPI routes (/api/*) take priority; everything else is proxied to Streamlit.
@app.api_route(
    "/{path:path}",
    methods=["GET", "POST", "PUT", "DELETE", "OPTIONS", "HEAD"],
)
async def proxy(request: Request, path: str):
    """Proxy all non-/api/* requests to the internal Streamlit server."""
    url = f"http://127.0.0.1:{_STREAMLIT_PORT}/{path}"
    params = dict(request.query_params)
    headers = {
        k: v for k, v in request.headers.items() if k.lower() != "host"
    }
    body = await request.body()

    async with httpx.AsyncClient(timeout=60.0) as client:
        try:
            resp = await client.request(
                method=request.method,
                url=url,
                headers=headers,
                content=body,
                params=params,
            )
            return StarletteResponse(
                content=resp.content,
                status_code=resp.status_code,
                headers=dict(resp.headers),
            )
        except httpx.ConnectError:
            # Streamlit may still be starting up — return a friendly retry message.
            return StarletteResponse(
                content=b"Streamlit is starting, please wait...",
                status_code=503,
                headers={"Retry-After": "5"},
            )


# ── Entrypoint ────────────────────────────────────────────────────────────────
if __name__ == "__main__":
    start_streamlit()
    uvicorn.run(app, host="0.0.0.0", port=7860, log_level="info")