import asyncio
import base64
import hashlib
import hmac
import json
import logging
import os
import re
import secrets
import time
from typing import Any, Literal

import gradio as gr
import uvicorn
from fastapi import FastAPI, Header, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import HTMLResponse
from gradio_client import Client
from pydantic import BaseModel, Field


SESSION_SECRET = os.getenv("SESSION_SECRET", "")
HF_TOKEN = os.getenv("HF_TOKEN", "")
TARGET_SPACE_ID = os.getenv("TARGET_SPACE_ID", "ASQ2/ASQ")
INITIAL_PROMPT = os.getenv("ASQ_INITIAL_PROMPT", "Hello")
TARGET_CHAT_API_NAME = os.getenv("TARGET_CHAT_API_NAME", "/chat_flow")
TARGET_INITIALIZE_API_NAME = os.getenv("TARGET_INITIALIZE_API_NAME", "/initialize_flow")
SESSION_TTL_SECONDS = int(os.getenv("SESSION_TTL_SECONDS", "14400"))
MAX_BACKGROUND_CHARS = int(os.getenv("MAX_BACKGROUND_CHARS", "20000"))
PREPARE_TIMEOUT_SECONDS = int(os.getenv("PREPARE_TIMEOUT_SECONDS", "600"))
PREPARE_RETRY_SECONDS = int(os.getenv("PREPARE_RETRY_SECONDS", "15"))
PAIRING_TTL_SECONDS = int(os.getenv("PAIRING_TTL_SECONDS", "600"))

logger = logging.getLogger("asq_gateway")


class PrepareRequest(BaseModel):
    background: str = Field(min_length=20, max_length=MAX_BACKGROUND_CHARS)


class JoinRequest(BaseModel):
    code: str = Field(min_length=6, max_length=6, pattern=r"^\d{6}$")


class SuggestRequest(BaseModel):
    transcript: str = Field(min_length=1, max_length=4000)


sessions: dict[str, dict[str, Any]] = {}
interview_codes: dict[str, str] = {}
pairing_codes: dict[str, str] = {}
pairings_by_token: dict[str, dict[str, Any]] = {}
preparation_tasks: set[asyncio.Task[Any]] = set()


def _b64encode(data: bytes) -> str:
    return base64.urlsafe_b64encode(data).decode("ascii").rstrip("=")


def _b64decode(data: str) -> bytes:
    return base64.urlsafe_b64decode(data + "=" * (-len(data) % 4))


def _sign(payload: str) -> str:
    return _b64encode(
        hmac.new(SESSION_SECRET.encode("utf-8"), payload.encode("utf-8"), hashlib.sha256).digest()
    )


def _cleanup_expired() -> None:
    now = int(time.time())
    expired_ids = [sid for sid, session in sessions.items() if session["expires_at"] <= now]
    for sid in expired_ids:
        session = sessions.pop(sid, None)
        if session and session.get("interview_code"):
            interview_codes.pop(session["interview_code"], None)
    expired_pair_tokens = [
        token for token, pairing in pairings_by_token.items() if pairing["expires_at"] <= now
    ]
    for pair_token in expired_pair_tokens:
        pairing = pairings_by_token.pop(pair_token, None)
        if pairing:
            pairing_codes.pop(pairing["code"], None)


def _new_interview_code() -> str:
    _cleanup_expired()
    for _ in range(100):
        code = f"{secrets.randbelow(1_000_000):06d}"
        if code not in interview_codes:
            return code
    raise RuntimeError("Could not allocate an interview code")


def _new_pairing_code() -> str:
    _cleanup_expired()
    for _ in range(100):
        code = f"{secrets.randbelow(1_000_000):06d}"
        if code not in pairing_codes and code not in interview_codes:
            return code
    raise RuntimeError("Could not allocate a pairing code")


def create_token(session_id: str, role: Literal["owner", "glasses"], expires_at: int) -> str:
    payload = _b64encode(
        json.dumps(
            {"sid": session_id, "role": role, "exp": expires_at}, separators=(",", ":")
        ).encode("utf-8")
    )
    return f"{payload}.{_sign(payload)}"


def verify_token(token: str, roles: set[str] | None = None) -> tuple[dict[str, Any], str]:
    if not SESSION_SECRET:
        raise HTTPException(status_code=503, detail="Gateway is not configured")
    try:
        payload, signature = token.split(".", 1)
        if not hmac.compare_digest(signature, _sign(payload)):
            raise ValueError("bad signature")
        data = json.loads(_b64decode(payload))
        if int(data["exp"]) <= int(time.time()):
            raise ValueError("expired")
        role = str(data["role"])
        session = sessions.get(str(data["sid"]))
        if not session:
            raise ValueError("unknown session")
        if roles is not None and role not in roles:
            raise HTTPException(status_code=403, detail="This session cannot perform that action")
        return session, role
    except HTTPException:
        raise
    except (ValueError, KeyError, json.JSONDecodeError):
        raise HTTPException(status_code=401, detail="Invalid or expired session")


def _bearer(authorization: str | None, roles: set[str] | None = None) -> tuple[dict[str, Any], str]:
    if not authorization or not authorization.startswith("Bearer "):
        raise HTTPException(status_code=401, detail="Missing session token")
    return verify_token(authorization[7:], roles)


def _extract_turn_count(value: Any) -> int:
    if isinstance(value, (int, float)):
        return int(value)
    match = re.search(r"\d+", str(value or ""))
    return int(match.group(0)) if match else 0


def _message_text(message: Any) -> str:
    if isinstance(message, str):
        return message.strip()
    if isinstance(message, dict):
        content = message.get("content", message.get("text", ""))
        return _message_text(content)
    if isinstance(message, (list, tuple)):
        parts = [_message_text(item) for item in message]
        return "\n".join(part for part in parts if part).strip()
    return str(message or "").strip()


def _first_complete_sentence(text: str) -> str:
    cleaned = text.strip()
    sentence_end = re.search(r"[.!?。！？]", cleaned)
    if sentence_end:
        return cleaned[: sentence_end.end()].strip()
    return cleaned.splitlines()[0].strip() if cleaned else ""


def extract_initial_state(result: Any) -> tuple[str, list[Any], int]:
    # /initialize_flow returns:
    # turn markdown, interview history, hypotheses, summary, strategy, status.
    if not isinstance(result, (list, tuple)) or len(result) < 6:
        raise RuntimeError("Unexpected initialization response from the ASQ Space")
    hypotheses = str(result[2] or "").strip()
    history = result[1] if isinstance(result[1], list) else []
    turns = _extract_turn_count(result[0])
    if not hypotheses:
        raise RuntimeError("ASQ returned no initial hypotheses")
    return hypotheses, history, turns


def extract_question(result: Any) -> tuple[str, list[Any], int, str | None]:
    # /chat_flow returns:
    # history, cleared input, turn markdown, hypotheses, summary, strategy, status.
    if not isinstance(result, (list, tuple)) or len(result) < 7:
        raise RuntimeError("Unexpected response from the ASQ Space")
    history = result[0]
    turns = _extract_turn_count(result[2])
    hypotheses = str(result[3]).strip() if result[3] is not None else None
    if not isinstance(history, list) or not history:
        raise RuntimeError("ASQ returned no conversation history")
    assistant_messages = [
        message
        for message in history
        if isinstance(message, dict) and message.get("role") == "assistant"
    ]
    question = _first_complete_sentence(
        _message_text((assistant_messages or history)[-1])
    )
    if not question:
        raise RuntimeError("ASQ returned an empty question")
    return question, history, turns, hypotheses


def initialize_asq(background: str) -> tuple[str, list[Any], int]:
    if not HF_TOKEN:
        raise RuntimeError("HF_TOKEN is not configured")
    client = Client(TARGET_SPACE_ID, hf_token=HF_TOKEN, verbose=False)
    result = client.predict(background, api_name=TARGET_INITIALIZE_API_NAME)
    return extract_initial_state(result)


def call_asq(transcript: str, session: dict[str, Any]) -> tuple[str, list[Any], int, str | None]:
    if not HF_TOKEN:
        raise RuntimeError("HF_TOKEN is not configured")
    client = Client(TARGET_SPACE_ID, hf_token=HF_TOKEN, verbose=False)
    result = client.predict(
        transcript,
        session["history"],
        session["background"],
        api_name=TARGET_CHAT_API_NAME,
    )
    return extract_question(result)


async def prepare_session(session_id: str, background: str) -> None:
    session = sessions.get(session_id)
    if not session:
        return
    deadline = time.monotonic() + PREPARE_TIMEOUT_SECONDS
    attempt = 0
    while True:
        attempt += 1
        try:
            hypotheses, history, turns = await asyncio.to_thread(initialize_asq, background)
            break
        except Exception as error:
            remaining = deadline - time.monotonic()
            logger.warning("ASQ preparation attempt %s failed: %s", attempt, type(error).__name__)
            if remaining <= 0:
                if session_id in sessions:
                    session["status"] = "error"
                    session["error"] = "Unable to start the model. Please try again."
                return
            await asyncio.sleep(min(PREPARE_RETRY_SECONDS, remaining))
    if session_id in sessions:
        session.update(
            status="ready",
            hypotheses=hypotheses,
            history=history,
            turns=turns,
            error=None,
        )


def session_view(session: dict[str, Any], role: str) -> dict[str, Any]:
    response: dict[str, Any] = {
        "status": session["status"],
        "ready": session["status"] == "ready",
        "expiresAt": session["expires_at"],
        "question": session.get("current_question", ""),
        "questionRevision": session.get("question_revision", 0),
        "turns": session.get("turns", 0),
    }
    if session.get("error"):
        response["error"] = session["error"]
    if role == "owner":
        response["interviewCode"] = session.get("interview_code")
        response["hypotheses"] = session.get("hypotheses", "")
    return response


api = FastAPI(title="ASQ Rokid Gateway", version="0.5.0")
api.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["GET", "POST", "OPTIONS"],
    allow_headers=["Authorization", "Content-Type"],
)


@api.get("/", response_class=HTMLResponse)
def mobile_app() -> str:
    return MOBILE_APP_HTML


@api.get("/health")
def health() -> dict[str, Any]:
    return {
        "ok": True,
        "configured": bool(SESSION_SECRET and HF_TOKEN),
        "target": TARGET_SPACE_ID,
        "version": "0.5.0",
    }


@api.post("/auth")
def create_owner_session() -> dict[str, Any]:
    if not SESSION_SECRET:
        raise HTTPException(status_code=503, detail="Gateway is not configured")
    _cleanup_expired()
    session_id = secrets.token_urlsafe(18)
    expires_at = int(time.time()) + SESSION_TTL_SECONDS
    sessions[session_id] = {
        "background": "",
        "hypotheses": "",
        "history": [],
        "turns": 0,
        "current_question": "",
        "question_revision": 0,
        "status": "draft",
        "error": None,
        "interview_code": None,
        "expires_at": expires_at,
    }
    return {
        "token": create_token(session_id, "owner", expires_at),
        "expiresIn": SESSION_TTL_SECONDS,
    }


@api.post("/prepare", status_code=202)
async def prepare(
    request: PrepareRequest,
    authorization: str | None = Header(default=None),
) -> dict[str, Any]:
    session, _ = _bearer(authorization, {"owner"})
    if session["status"] == "starting":
        return session_view(session, "owner")
    old_code = session.get("interview_code")
    if old_code:
        interview_codes.pop(old_code, None)
    code = _new_interview_code()
    session_id = next(sid for sid, value in sessions.items() if value is session)
    interview_codes[code] = session_id
    session.update(
        background=request.background.strip(),
        hypotheses="",
        history=[],
        turns=0,
        current_question="",
        question_revision=0,
        status="starting",
        error=None,
        interview_code=code,
    )
    task = asyncio.create_task(prepare_session(session_id, session["background"]))
    preparation_tasks.add(task)
    task.add_done_callback(preparation_tasks.discard)
    return session_view(session, "owner")


@api.post("/pair")
def create_pairing() -> dict[str, Any]:
    _cleanup_expired()
    code = _new_pairing_code()
    pair_token = secrets.token_urlsafe(24)
    expires_at = int(time.time()) + PAIRING_TTL_SECONDS
    pairings_by_token[pair_token] = {
        "code": code,
        "session_id": None,
        "expires_at": expires_at,
    }
    pairing_codes[code] = pair_token
    return {
        "pairingCode": code,
        "pairingToken": pair_token,
        "expiresIn": PAIRING_TTL_SECONDS,
    }


@api.get("/pair/status")
def pairing_status(
    x_pairing_token: str | None = Header(default=None, alias="X-Pairing-Token"),
) -> dict[str, Any]:
    _cleanup_expired()
    pairing = pairings_by_token.get(x_pairing_token or "")
    if not pairing:
        raise HTTPException(status_code=401, detail="Invalid or expired pairing")
    session_id = pairing.get("session_id")
    if not session_id:
        return {"paired": False, "expiresAt": pairing["expires_at"]}
    session = sessions.get(session_id)
    if not session:
        raise HTTPException(status_code=410, detail="Interview session expired")
    return {
        "paired": True,
        "token": create_token(session_id, "glasses", session["expires_at"]),
        **session_view(session, "glasses"),
    }


@api.post("/pair/claim")
def claim_pairing(
    request: JoinRequest,
    authorization: str | None = Header(default=None),
) -> dict[str, bool]:
    session, _ = _bearer(authorization, {"owner"})
    _cleanup_expired()
    pair_token = pairing_codes.get(request.code)
    pairing = pairings_by_token.get(pair_token or "")
    if not pairing or not hmac.compare_digest(request.code, pairing.get("code") or ""):
        raise HTTPException(status_code=401, detail="Incorrect or expired glasses pairing code")
    session_id = next(sid for sid, value in sessions.items() if value is session)
    pairing["session_id"] = session_id
    return {"paired": True}


@api.post("/join")
def join(request: JoinRequest) -> dict[str, Any]:
    _cleanup_expired()
    session_id = interview_codes.get(request.code)
    session = sessions.get(session_id or "")
    if not session or not hmac.compare_digest(request.code, session.get("interview_code") or ""):
        raise HTTPException(status_code=401, detail="Incorrect or expired interview code")
    return {
        "token": create_token(session_id, "glasses", session["expires_at"]),
        **session_view(session, "glasses"),
    }


@api.get("/session")
def get_session(authorization: str | None = Header(default=None)) -> dict[str, Any]:
    session, role = _bearer(authorization)
    return session_view(session, role)


@api.post("/start")
async def start(authorization: str | None = Header(default=None)) -> dict[str, Any]:
    session, role = _bearer(authorization, {"owner", "glasses"})
    if session["status"] != "ready":
        raise HTTPException(status_code=409, detail="Interview is not ready")
    if session.get("current_question"):
        return session_view(session, role)
    try:
        question, history, turns, hypotheses = await asyncio.to_thread(
            call_asq, INITIAL_PROMPT, session
        )
    except Exception as error:
        raise HTTPException(status_code=502, detail=f"ASQ inference failed: {error}")
    session["history"] = history
    session["turns"] = turns
    session["current_question"] = question
    session["question_revision"] = session.get("question_revision", 0) + 1
    if hypotheses:
        session["hypotheses"] = hypotheses
    return session_view(session, role)


@api.post("/suggest")
async def suggest(
    request: SuggestRequest,
    authorization: str | None = Header(default=None),
) -> dict[str, str]:
    session, _ = _bearer(authorization, {"owner", "glasses"})
    if session["status"] != "ready":
        raise HTTPException(status_code=409, detail="Interview is not ready")
    try:
        question, history, turns, hypotheses = await asyncio.to_thread(
            call_asq, request.transcript.strip(), session
        )
    except Exception as error:
        raise HTTPException(status_code=502, detail=f"ASQ inference failed: {error}")
    session["history"] = history
    session["turns"] = turns
    session["current_question"] = question
    session["question_revision"] = session.get("question_revision", 0) + 1
    if hypotheses:
        session["hypotheses"] = hypotheses
    return {"question": question}


MOBILE_APP_HTML = """<!doctype html>
<html lang="en"><head><meta charset="utf-8"><meta name="viewport" content="width=device-width,initial-scale=1">
<title>ASQ Interview Setup</title>
<style>
:root{font-family:Inter,system-ui,sans-serif;color:#16201d;background:#eef5f1}*{box-sizing:border-box}
body{margin:0;padding:24px 16px}.card{max-width:760px;margin:auto;background:#fff;border-radius:20px;padding:24px;box-shadow:0 12px 36px #173b2b1a}
h1{margin:0 0 6px;font-size:1.55rem}.muted{color:#65716d;margin:0 0 24px}.step{border-top:1px solid #e4ebe7;padding-top:20px;margin-top:20px}
label{display:block;font-weight:650;margin-bottom:8px}input,textarea,button{font:inherit;width:100%;border-radius:12px}input,textarea{border:1px solid #b9c9c1;padding:13px}textarea{min-height:220px;resize:vertical}
button{border:0;background:#087a54;color:#fff;padding:14px;font-weight:700;margin-top:12px}button:disabled{background:#9cada5}.hidden{display:none}.status{padding:14px;border-radius:12px;background:#edf7f2;margin-top:14px;white-space:pre-wrap}
.hypo{white-space:pre-wrap;line-height:1.5;background:#f7faf8;padding:14px;border-radius:12px;max-height:320px;overflow:auto}
</style></head><body><main class="card">
<h1>ASQ Interview Setup</h1><p class="muted">Open ASQ on the Rokid glasses first. The glasses will display a six-digit pairing code; enter that code on this phone.</p>
<section id="prepare"><label for="pairingCode">Code shown on the glasses</label><input id="pairingCode" inputmode="numeric" maxlength="6" placeholder="Six digits"><label for="background" style="margin-top:16px">Case and interview background</label><textarea id="background" placeholder="Enter the case facts, context, existing evidence, interview goals, and any constraints..."></textarea><button id="prepareBtn" disabled>Connect glasses and prepare interview</button></section>
<section id="progress" class="step hidden"><div id="status" class="status">Connecting the glasses…</div><div id="hypoWrap" class="hidden"><h2>Initial hypotheses</h2><div id="hypotheses" class="hypo"></div></div></section>
</main><script>
let token="", timer=null; const $=id=>document.getElementById(id);
async function api(path, options={}){const headers={"Content-Type":"application/json",...(options.headers||{})};if(token)headers.Authorization=`Bearer ${token}`;const r=await fetch(path,{...options,headers});let data={};try{data=await r.json()}catch{}if(!r.ok)throw new Error(data.detail||`Request failed (${r.status})`);return data}
function errorMessage(e){$("status").textContent=e.message;$("progress").classList.remove("hidden")}
async function begin(){try{const data=await api("/auth",{method:"POST"});token=data.token;$("prepareBtn").disabled=false}catch(e){errorMessage(e)}}
$("prepareBtn").onclick=async()=>{const pairingCode=$("pairingCode").value.trim();const background=$("background").value.trim();if(!/^\d{6}$/.test(pairingCode)){errorMessage("Enter the six-digit code shown on the glasses.");return}if(background.length<20){errorMessage("Please provide at least 20 characters of background.");return}$("prepareBtn").disabled=true;$("progress").classList.remove("hidden");$("status").textContent="Connecting the glasses and generating initial hypotheses…";try{await api("/pair/claim",{method:"POST",body:JSON.stringify({code:pairingCode})});const data=await api("/prepare",{method:"POST",body:JSON.stringify({background})});show(data);timer=setInterval(poll,3000)}catch(e){errorMessage(e);$("prepareBtn").disabled=false}};
async function poll(){try{show(await api("/session"))}catch(e){clearInterval(timer);errorMessage(e)}}
function show(data){if(data.status==="ready"){clearInterval(timer);$("status").textContent="Ready. The glasses are connected and can begin the interview.";$("hypotheses").textContent=data.hypotheses||"No hypotheses returned.";$("hypoWrap").classList.remove("hidden")}else if(data.status==="error"){clearInterval(timer);$("status").textContent=data.error||"Model startup failed.";$("prepareBtn").disabled=false}else{$("status").textContent="Glasses connected. Starting the model and generating initial hypotheses…"}}
begin();
</script></body></html>"""


with gr.Blocks(title="ASQ Rokid Gateway") as status_page:
    gr.Markdown(
        "# ASQ Rokid Gateway\n"
        "This service provides phone preparation and Rokid interview pairing APIs.\n\n"
        "Open the site root on a phone to prepare an interview. Use `GET /health` for service status."
    )


api = gr.mount_gradio_app(api, status_page, path="/ui")


if __name__ == "__main__":
    uvicorn.run(api, host="0.0.0.0", port=7860)