"""
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 re
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")


def _truthy(v: str | None) -> bool:
    """Parse a multipart form flag into a bool ('1'/'true'/'yes'/'on')."""
    return str(v or "").strip().lower() in ("1", "true", "yes", "on")


def _term_list(v: str | None) -> list:
    """Split a comma/newline/semicolon term list into a clean list."""
    return [t.strip() for t in re.split(r"[\n,;]+", v or "") if t.strip()]


# ── 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,
    maximum_ats_mode: bool = False,
    confirmed_terms: list | None = None,
) -> 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": {},
        # Maximum ATS Mode flows into _generate_resume_v4 via the job dict so
        # normal PM/AI craft terms are treated as user-confirmed and woven in.
        "_maximum_ats_mode": bool(maximum_ats_mode),
        "_confirmed_terms": confirmed_terms or [],
    }

    # 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(""),
    maximum_ats_mode: str = Form(""),       # "1"/"true" enables Maximum ATS Mode
    user_confirmed_expansion: str = Form(""),  # alias for maximum_ats_mode
    confirmed_terms: str = Form(""),        # optional comma/newline list
    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")

    max_ats = _truthy(maximum_ats_mode) or _truthy(user_confirmed_expansion)
    conf_terms = _term_list(confirmed_terms)

    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,
            max_ats, conf_terms,
        )

        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",
                                            v2_report.get("high_risk_terms_for_confirmation", [])),
            "maximum_ats_mode": max_ats,
            "external_coverage": v2_report.get("external_coverage", {}),
            "external_coverage_pct": v2_report.get("external_coverage_pct"),
            "coverage_report": v2_report.get("coverage_report", {}),
            "keyword_coverage": v2_report.get("keyword_coverage", {}),
            "docx_b64": docx_b64,
            "pdf_b64": pdf_b64,
        }
        print(f"[api/generate] maximum_ats_mode={max_ats} status={response_payload['status']} "
              f"internal={response_payload['scores'].get('jd_match')} "
              f"external_cov={response_payload.get('external_coverage_pct')}")
        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)


# ── External ATS Feedback Repair helper ──────────────────────────────────────
def repair_resume_for_api(
    pdf_bytes: bytes, jd_text: str, job_title: str, company: str,
    feedback: str, missing_keywords: list, external_score,
    maximum_ats_mode: bool = False, confirmed_terms: list | None = None,
    target_external_score=None,
) -> dict:
    """Re-tailor using pasted external-checker feedback. Honest: every keyword is
    risk-classified; BLOCKED + unconfirmed HIGH terms are never added. Returns the
    enriched repair result dict (incl. resume_path + coverage_report)."""
    from src.resume_parser_v2 import parse_resume_pdf
    from src.llm_client import LLMClient
    from src.jobalytics_repair import repair_with_external_feedback

    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

    tmp_out = tempfile.mkdtemp()
    job = {
        "title": job_title, "company": company,
        "description": (jd_text or "")[:16000], "ats_keywords": "", "_raw_assessment": {},
    }
    result = repair_with_external_feedback(
        job, feedback_text=feedback or None,
        missing_keywords=missing_keywords or None,
        external_score=external_score, llm=LLMClient(),
        base_resume=base_resume, output_dir=tmp_out,
        maximum_ats_mode=maximum_ats_mode, confirmed_terms=confirmed_terms,
        target_external_score=target_external_score,
    )
    # Persist user-confirmed expansion terms to the vault so future resumes treat
    # them as safe (only the terms actually added under confirmation).
    try:
        if maximum_ats_mode and not result.get("error"):
            from src.candidate_vault import confirm_expansion_terms
            confirm_expansion_terms(result.get("added_terms", []))
            if confirmed_terms:
                confirm_expansion_terms(confirmed_terms)
            result["vault_added"] = result.get("added_terms", [])
    except Exception:
        pass
    return result


# ── /api/repair-with-feedback ─────────────────────────────────────────────────
@app.post("/api/repair-with-feedback")
async def repair_with_feedback(
    jd_text: str = Form(...),
    job_title: str = Form(""),
    company: str = Form(""),
    feedback: str = Form(""),
    missing_keywords: str = Form(""),   # optional comma/newline list
    external_score: str = Form(""),
    maximum_ats_mode: str = Form(""),
    user_confirmed_expansion: str = Form(""),
    confirmed_terms: str = Form(""),
    target_external_score: str = Form(""),
    resume: UploadFile = None,
    x_api_token: str = Header(None),
):
    """External ATS Feedback Repair Mode. Paste Jobalytics/Simplify feedback (or
    an explicit missing-keyword list); re-tailor honestly and re-score from the
    re-parsed export. BLOCKED + unconfirmed HIGH terms are never fabricated."""
    _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")
    if not feedback.strip() and not missing_keywords.strip():
        raise HTTPException(status_code=422, detail="feedback_or_keywords_required")

    mk = _term_list(missing_keywords) or None
    ext = None
    try:
        ext = int(re.sub(r"[^\d]", "", external_score)) if external_score.strip() else None
    except ValueError:
        ext = None
    tgt = None
    try:
        tgt = int(re.sub(r"[^\d]", "", target_external_score)) if target_external_score.strip() else None
    except ValueError:
        tgt = None
    max_ats = _truthy(maximum_ats_mode) or _truthy(user_confirmed_expansion)
    conf_terms = _term_list(confirmed_terms) or None

    result_path = None
    try:
        loop = asyncio.get_event_loop()
        result = await loop.run_in_executor(
            None, repair_resume_for_api, pdf_bytes, jd_text, job_title, company,
            feedback, mk, ext, max_ats, conf_terms, tgt,
        )
        print(f"[api/repair] maximum_ats_mode={max_ats} status={result.get('status')} "
              f"after_cov={result.get('after_coverage', {}).get('pct')} "
              f"added={len(result.get('added_terms', []))}")
        if result.get("error"):
            return JSONResponse(
                {"error": result["error"], "detail": result.get("detail", "")},
                status_code=400 if result["error"] == "no_missing_keywords" else 500,
            )

        result_path = result.get("resume_path")
        docx_b64 = pdf_b64 = None
        if result_path and os.path.exists(result_path):
            with open(result_path, "rb") as f:
                docx_b64 = base64.b64encode(f.read()).decode("ascii")
            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")

        sc = result.get("scores", {})
        return JSONResponse({
            "status": result.get("status", "NEEDS_USER_INPUT"),
            "download_allowed": bool(result.get("download_allowed", False)),
            "scores": {
                "jd_match": sc.get("internal_jd_match"),
                "ats_readability": sc.get("ats_readability"),
                "independent_jd_match": sc.get("independent_jd_match"),
            },
            "external_score": result.get("external_score"),
            "target_external_score": result.get("target_external_score"),
            "maximum_ats_mode": result.get("maximum_ats_mode", max_ats),
            "before_coverage": result.get("before_coverage"),
            "after_coverage": result.get("after_coverage"),
            "added_terms": result.get("added_terms", []),
            "review_flag_terms": result.get("review_flag_terms", []),
            "unresolved_high_risk_terms": result.get("unresolved_high_risk_terms", []),
            "blocked_terms": result.get("blocked_terms", []),
            "already_present_terms": result.get("already_present_terms", []),
            "still_missing_repairable": result.get("still_missing_repairable", []),
            "coverage_report": result.get("coverage_report", {}),
            "below_target_explanation": result.get("below_target_explanation", ""),
            "vault_added": result.get("vault_added", []),
            "docx_b64": docx_b64,
            "pdf_b64": pdf_b64,
        })
    except HTTPException:
        raise
    except Exception as exc:
        return JSONResponse({"error": "repair_failed", "detail": str(exc)}, status_code=500)
    finally:
        if result_path:
            shutil.rmtree(os.path.dirname(result_path), 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."""
    # Never proxy /api/* to Streamlit — if it reaches here it's an unknown API
    # route; return a clean JSON 404 (prevents a confusing blank Streamlit page).
    if path == "api" or path.startswith("api/"):
        return JSONResponse({"error": "not_found", "path": "/" + path}, status_code=404)

    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,
            )
            # httpx already DECOMPRESSED resp.content, so we must drop the
            # upstream content-encoding/length (and hop-by-hop headers) — copying
            # them makes the browser try to gunzip plain bytes → blank page.
            _DROP = {"content-encoding", "content-length", "transfer-encoding",
                     "connection", "keep-alive"}
            clean_headers = {k: v for k, v in resp.headers.items()
                             if k.lower() not in _DROP}
            return StarletteResponse(
                content=resp.content,
                status_code=resp.status_code,
                headers=clean_headers,
                media_type=resp.headers.get("content-type"),
            )
        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")