# retoor <retoor@molodetz.nl>

import json
import logging
from typing import Annotated

from fastapi import Depends, APIRouter, Request, WebSocket, WebSocketDisconnect
from fastapi.responses import HTMLResponse, JSONResponse, Response

from devplacepy import database
from devplacepy.config import DEEPSEARCH_DIR
from devplacepy.models import DeepsearchChatForm, DeepsearchRunForm
from devplacepy.responses import respond
from devplacepy.schemas import DeepsearchJobOut, DeepsearchSessionOut
from devplacepy.seo import base_seo_context, site_url, web_application_schema, website_schema
from devplacepy.services.deepsearch.chat import DeepsearchChat
from devplacepy.services.deepsearch.export import to_json, to_markdown, to_pdf
from devplacepy.services.jobs import queue
from devplacepy.services.jobs.deepsearch.progress import hub
from devplacepy.services.jobs.deepsearch.service import DeepsearchService
from devplacepy.services.manager import service_manager
from devplacepy.templating import templates
from devplacepy.utils import generate_uid, get_current_user, is_admin, not_found, track_action
from devplacepy.dependencies import json_or_form

from ._shared import owner_for

logger = logging.getLogger(__name__)
router = APIRouter()

ACTIVE_STATES = (queue.PENDING, queue.RUNNING)
TERMINAL_STATES = (queue.DONE, queue.FAILED)
CONTROL_STATES = ("running", "paused", "cancelled")
TOUCH_EXTEND_SECONDS = 3600
_service = DeepsearchService()

def _collection_for(uid: str) -> str:
    return _service.collection_name(uid)

def _control_path(uid: str):
    return DEEPSEARCH_DIR / uid / "control.json"

def _set_control(uid: str, state: str) -> None:
    path = _control_path(uid)
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(json.dumps({"state": state}), encoding="utf-8")

def _owns(request: Request, session: dict) -> bool:
    owner_kind, owner_id = owner_for(request)
    return (
        session.get("owner_kind") == owner_kind
        and session.get("owner_id") == owner_id
    )

def _job_payload(job: dict) -> dict:
    result = job.get("result", {})
    payload = job.get("payload", {})
    done = job.get("status") == queue.DONE
    uid = job.get("uid", "")
    return {
        "uid": uid,
        "kind": job.get("kind", ""),
        "status": job.get("status", ""),
        "query": payload.get("query"),
        "depth": int(payload.get("depth") or 0),
        "max_pages": int(payload.get("max_pages") or 0),
        "ws_url": f"/tools/deepsearch/{uid}/ws",
        "chat_ws_url": f"/tools/deepsearch/{uid}/chat" if done else None,
        "session_url": f"/tools/deepsearch/{uid}/session" if done else None,
        "score": result.get("score") if done else None,
        "confidence": result.get("confidence") if done else None,
        "source_diversity": result.get("source_diversity") if done else None,
        "page_count": int(result.get("page_count") or 0),
        "chunk_count": int(result.get("chunk_count") or 0),
        "error": job.get("error") or None,
        "created_at": job.get("created_at"),
        "completed_at": job.get("completed_at") or None,
    }

@router.get("", response_class=HTMLResponse)
async def deepsearch_page(request: Request):
    user = get_current_user(request)
    base = site_url(request)
    description = (
        "Run a multi-agent deep web research job: it plans queries, crawls and indexes "
        "sources, then synthesises a cited report with confidence scoring and gap "
        "analysis, plus a grounded chat over the results."
    )
    seo_ctx = base_seo_context(
        request,
        title="DeepSearch",
        description=description,
        breadcrumbs=[
            {"name": "Home", "url": "/feed"},
            {"name": "Tools", "url": "/tools"},
            {"name": "DeepSearch", "url": "/tools/deepsearch"},
        ],
        schemas=[
            website_schema(base),
            web_application_schema("DeepSearch", description, "/tools/deepsearch", base),
        ],
    )
    return templates.TemplateResponse(
        request,
        "tools/deepsearch.html",
        {**seo_ctx, "request": request, "user": user},
    )

@router.post("/run")
async def deepsearch_run(request: Request, data: Annotated[DeepsearchRunForm, Depends(json_or_form(DeepsearchRunForm))]):
    owner_kind, owner_id = owner_for(request)
    from devplacepy.services.audit import record as audit

    active = [
        job
        for job in queue.list_jobs(kind="deepsearch", owner=(owner_kind, owner_id))
        if job.get("status") in ACTIVE_STATES
    ]
    if active:
        audit.record(
            request,
            "deepsearch.run.request",
            result="denied",
            summary=f"DeepSearch denied for {data.query}: already running",
            metadata={"query": data.query, "reason": "active_job", "uid": active[0]["uid"]},
            links=[audit.job(active[0]["uid"])],
        )
        return JSONResponse(
            {
                "error": {
                    "status": 429,
                    "message": "You already have a DeepSearch running. Wait for it to finish.",
                    "uid": active[0]["uid"],
                }
            },
            status_code=429,
        )
    user = get_current_user(request)
    api_key = (
        user.get("api_key") if user else ""
    ) or database.internal_gateway_key()
    uid = generate_uid()
    payload = {
        "query": data.query,
        "depth": data.depth,
        "max_pages": data.max_pages,
        "api_key": api_key,
        "collection": _collection_for(uid),
    }
    database.create_deepsearch_session(
        uid, owner_kind, owner_id, data.query, data.depth, data.max_pages, _collection_for(uid)
    )
    _enqueue(uid, payload, owner_kind, owner_id, data.query)
    audit.record(
        request,
        "deepsearch.run.request",
        summary=f"requested DeepSearch for {data.query}",
        metadata={"query": data.query, "depth": data.depth, "max_pages": data.max_pages},
        links=[audit.job(uid)],
    )
    if owner_kind == "user":
        track_action(owner_id, "deepsearch")
    return JSONResponse(
        {
            "uid": uid,
            "status_url": f"/tools/deepsearch/{uid}",
            "ws_url": f"/tools/deepsearch/{uid}/ws",
        }
    )

def _enqueue(uid: str, payload: dict, owner_kind: str, owner_id: str, query: str) -> None:
    from datetime import datetime, timezone

    now = datetime.now(timezone.utc).isoformat()
    database.get_table("jobs").insert(
        {
            "uid": uid,
            "kind": "deepsearch",
            "status": queue.PENDING,
            "owner_kind": owner_kind,
            "owner_id": owner_id,
            "preferred_name": f"DeepSearch: {query}"[:64],
            "payload": json.dumps(payload),
            "result": "",
            "error": "",
            "retry_count": 0,
            "created_at": now,
            "started_at": "",
            "completed_at": "",
            "updated_at": now,
            "duration_ms": 0,
            "last_accessed_at": "",
            "expires_at": "",
            "bytes_in": 0,
            "bytes_out": 0,
            "item_count": 0,
        }
    )

@router.get("/{uid}")
async def deepsearch_status(request: Request, uid: str):
    job = queue.get_job(uid)
    if not job or job.get("kind") != "deepsearch":
        raise not_found("Research not found")
    return JSONResponse(
        DeepsearchJobOut.model_validate(_job_payload(job)).model_dump(mode="json")
    )

def _report_from_disk(uid: str) -> dict:
    path = DEEPSEARCH_DIR / uid / "report.json"
    try:
        return json.loads(path.read_text(encoding="utf-8"))
    except (ValueError, OSError):
        return {}

def _report_for(uid: str, job: dict) -> dict:
    if job.get("status") == queue.DONE:
        report = job.get("result", {}).get("report", {})
        if report:
            return report
    if job.get("status") == queue.FAILED:
        return {}
    return _report_from_disk(uid)

def _session_context(request: Request, uid: str, job: dict, session: dict) -> dict:
    report = _report_for(uid, job)
    user = get_current_user(request)
    viewer_is_admin = is_admin(user)
    done = bool(report) or job.get("status") == queue.DONE
    return {
        "uid": uid,
        "status": queue.DONE if done else job.get("status", ""),
        "query": report.get("query") or session.get("query"),
        "depth": int(session.get("depth") or 0),
        "max_pages": int(session.get("max_pages") or 0),
        "score": report.get("score"),
        "confidence": report.get("confidence"),
        "source_diversity": report.get("source_diversity"),
        "synthesis": report.get("synthesis", ""),
        "page_count": report.get("page_count", 0),
        "chunk_count": report.get("chunk_count", 0),
        "summary": report.get("summary", ""),
        "sources": report.get("sources", []),
        "findings": report.get("findings", []),
        "timeline": report.get("timeline", []),
        "follow_up_questions": report.get("follow_up_questions", []),
        "chat_ws_url": f"/tools/deepsearch/{uid}/chat" if done else None,
        "export_md_url": f"/tools/deepsearch/{uid}/export.md" if done else None,
        "export_json_url": f"/tools/deepsearch/{uid}/export.json" if done else None,
        "export_pdf_url": f"/tools/deepsearch/{uid}/export.pdf" if done else None,
        "viewer_is_admin": viewer_is_admin,
        "viewer_owns": _owns(request, session),
        "created_at": session.get("created_at"),
        "completed_at": session.get("completed_at") or None,
        "request": request,
        "user": user,
        "meta_robots": "noindex,nofollow",
    }

@router.get("/{uid}/session")
async def deepsearch_session(request: Request, uid: str):
    job = queue.get_job(uid)
    session = database.get_deepsearch_session(uid)
    if not job or job.get("kind") != "deepsearch" or not session:
        raise not_found("Research not found")
    queue.touch_job(uid, TOUCH_EXTEND_SECONDS)
    context = _session_context(request, uid, job, session)
    return respond(request, "tools/deepsearch_session.html", context, model=DeepsearchSessionOut)

@router.post("/{uid}/pause")
async def deepsearch_pause(request: Request, uid: str):
    return _control(request, uid, "paused")

@router.post("/{uid}/resume")
async def deepsearch_resume(request: Request, uid: str):
    return _control(request, uid, "running")

@router.post("/{uid}/cancel")
async def deepsearch_cancel(request: Request, uid: str):
    return _control(request, uid, "cancelled")

def _control(request: Request, uid: str, state: str):
    session = database.get_deepsearch_session(uid)
    job = queue.get_job(uid)
    if not job or job.get("kind") != "deepsearch" or not session:
        raise not_found("Research not found")
    if not _owns(request, session):
        return JSONResponse(
            {"error": {"status": 403, "message": "Not your research session."}},
            status_code=403,
        )
    _set_control(uid, state)
    return JSONResponse({"ok": True, "state": state})

def _export_report(uid: str) -> dict | None:
    job = queue.get_job(uid)
    if not job or job.get("kind") != "deepsearch" or job.get("status") == queue.FAILED:
        return None
    report = _report_for(uid, job)
    if not report:
        return None
    queue.touch_job(uid, TOUCH_EXTEND_SECONDS)
    return report

@router.get("/{uid}/export.md")
async def deepsearch_export_md(request: Request, uid: str):
    report = _export_report(uid)
    if report is None:
        raise not_found("Report not ready")
    return Response(content=to_markdown(report), media_type="text/markdown")

@router.get("/{uid}/export.json")
async def deepsearch_export_json(request: Request, uid: str):
    report = _export_report(uid)
    if report is None:
        raise not_found("Report not ready")
    return Response(content=to_json(report), media_type="application/json")

@router.get("/{uid}/export.pdf")
async def deepsearch_export_pdf(request: Request, uid: str):
    report = _export_report(uid)
    if report is None:
        raise not_found("Report not ready")
    return Response(content=to_pdf(report), media_type="application/pdf")

@router.websocket("/{uid}/ws")
async def deepsearch_ws(websocket: WebSocket, uid: str):
    await websocket.accept()
    if not service_manager.owns_lock():
        await websocket.close(code=4013)
        return
    job = queue.get_job(uid)
    if not job or job.get("kind") != "deepsearch":
        await websocket.close(code=1008)
        return
    listener = hub.subscribe(uid)
    try:
        for frame in hub.snapshot(uid):
            await websocket.send_json(frame)
        current = queue.get_job(uid)
        if current and current.get("status") in TERMINAL_STATES:
            done = current.get("status") == queue.DONE
            await websocket.send_json(
                {
                    "type": "done" if done else "failed",
                    "status": current.get("status"),
                    "session_url": f"/tools/deepsearch/{uid}/session",
                    "error": current.get("error") or None,
                }
            )
            return
        while True:
            frame = await listener.get()
            await websocket.send_json(frame)
            if frame.get("type") in ("done", "failed"):
                break
    except WebSocketDisconnect:
        pass
    except Exception:
        logger.exception("DeepSearch websocket loop failed for %s", uid)
    finally:
        hub.unsubscribe(uid, listener)

@router.websocket("/{uid}/chat")
async def deepsearch_chat_ws(websocket: WebSocket, uid: str):
    await websocket.accept()
    if not service_manager.owns_lock():
        await websocket.close(code=4013)
        return
    job = queue.get_job(uid)
    session = database.get_deepsearch_session(uid)
    ready = job and (
        job.get("status") == queue.DONE or (session or {}).get("status") == "done"
    )
    if not job or job.get("kind") != "deepsearch" or not ready or not session:
        await websocket.close(code=1008)
        return
    user = get_current_user(websocket)
    api_key = (
        user.get("api_key") if user else ""
    ) or database.internal_gateway_key()
    chat = DeepsearchChat(_collection_for(uid), api_key)
    history: list[dict] = []
    for row in database.get_deepsearch_messages(uid):
        history.append({"role": row.get("role"), "content": row.get("content")})
    await websocket.send_json({"type": "ready", "history": history})
    from devplacepy.services.audit import record as audit

    try:
        while True:
            raw = await websocket.receive_text()
            try:
                frame = json.loads(raw)
            except (ValueError, TypeError):
                continue
            if frame.get("type") != "input":
                continue
            try:
                form = DeepsearchChatForm.model_validate({"message": frame.get("text", "")})
            except Exception:
                await websocket.send_json({"type": "error", "text": "Empty message."})
                continue
            question = form.message
            await websocket.send_json({"type": "status", "text": "Searching the research collection"})
            database.add_deepsearch_message(generate_uid(), uid, "user", question)
            history.append({"role": "user", "content": question})
            answer = await chat.answer(question, history)
            database.add_deepsearch_message(
                generate_uid(),
                uid,
                "assistant",
                answer.text,
                json.dumps(answer.citations),
            )
            history.append({"role": "assistant", "content": answer.text})
            await websocket.send_json(
                {"type": "reply", "text": answer.text, "citations": answer.citations}
            )
            audit.record(
                websocket,
                "deepsearch.chat",
                summary=f"DeepSearch chat on session {uid[:8]}",
                metadata={"session": uid},
                links=[audit.job(uid)],
            )
    except WebSocketDisconnect:
        pass
    except Exception:
        logger.exception("DeepSearch chat loop failed for %s", uid)
