# retoor <retoor@molodetz.nl>

from __future__ import annotations

import json
import logging
import re
from dataclasses import dataclass, field
from urllib.parse import urlparse

from devplacepy import stealth
from devplacepy.config import INTERNAL_GATEWAY_URL, INTERNAL_MODEL

logger = logging.getLogger(__name__)

AGENT_TIMEOUT_SECONDS = 120.0
SUMMARY_MAX_TOKENS = 900
CRITIC_MAX_TOKENS = 600
LINKER_MAX_TOKENS = 600
MAX_CONTEXT_CHARS = 11000
SCORE_MAX = 100
CONFIDENCE_BASELINE = 0.35
DIVERSITY_PAGES_PER_DOMAIN = 2.0


@dataclass
class Orchestration:
    summary: str = ""
    findings: list[dict] = field(default_factory=list)
    gaps: list[str] = field(default_factory=list)
    confidence: float = 0.0
    source_diversity: float = 0.0
    score: int = 0


def _domain(url: str) -> str:
    try:
        return urlparse(url).netloc.lower()
    except ValueError:
        return ""


def source_diversity(pages: list) -> float:
    if not pages:
        return 0.0
    domains = {_domain(page.url) for page in pages if getattr(page, "url", "")}
    domains.discard("")
    if not domains:
        return 0.0
    ratio = len(domains) / max(1.0, len(pages) / DIVERSITY_PAGES_PER_DOMAIN)
    return round(min(1.0, ratio), 3)


def _build_context(pages: list) -> str:
    blocks: list[str] = []
    used = 0
    for index, page in enumerate(pages, start=1):
        snippet = (page.text or "")[:1600]
        block = f"[{index}] {page.title} ({page.url})\n{snippet}"
        if used + len(block) > MAX_CONTEXT_CHARS and blocks:
            break
        used += len(block)
        blocks.append(block)
    return "\n\n".join(blocks)


async def _complete(messages: list[dict], api_key: str, max_tokens: int) -> str:
    payload = {
        "model": INTERNAL_MODEL,
        "messages": messages,
        "max_tokens": max_tokens,
        "temperature": 0.2,
    }
    headers = {"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"}
    async with stealth.stealth_async_client(timeout=AGENT_TIMEOUT_SECONDS) as client:
        response = await client.post(INTERNAL_GATEWAY_URL, json=payload, headers=headers)
    if response.status_code >= 400:
        raise RuntimeError(f"agent gateway returned {response.status_code}")
    data = response.json()
    return (data.get("choices") or [{}])[0].get("message", {}).get("content") or ""


def _parse_json(text: str) -> dict:
    match = re.search(r"\{.*\}", text, re.DOTALL)
    if not match:
        return {}
    try:
        return json.loads(match.group())
    except (ValueError, TypeError):
        return {}


SUMMARIZER_PROMPT = (
    "You are a research summarizer. Using ONLY the numbered SOURCES, write a JSON object "
    "with keys: 'summary' (a grounded markdown summary answering the question) and "
    "'findings' (an array of objects, each with 'title', 'detail', 'confidence' between 0 "
    "and 1, and 'citations' an array of source numbers). Cite only the provided sources. "
    "Return ONLY the JSON object."
)
CRITIC_PROMPT = (
    "You are a research critic. Given a QUESTION, a draft SUMMARY and FINDINGS, identify "
    "what is missing, contradictory, or weakly supported. Return ONLY a JSON object with "
    "key 'gaps': an array of short strings describing open questions or weak spots."
)
LINKER_PROMPT = (
    "You are a research linker. Given FINDINGS and the SOURCES, refine the confidence of "
    "each finding based on how many independent sources support it. Return ONLY a JSON "
    "object with key 'confidence': a number between 0 and 1 estimating overall answer "
    "confidence given source agreement and coverage."
)


def _heuristic(question: str, pages: list) -> Orchestration:
    diversity = source_diversity(pages)
    findings = []
    for page in pages[:5]:
        findings.append(
            {
                "title": page.title[:120] or page.url,
                "detail": (page.text or "")[:400],
                "confidence": round(min(0.6, CONFIDENCE_BASELINE + diversity / 4), 3),
                "citations": [page.url],
            }
        )
    summary = (
        f"Collected {len(pages)} sources for '{question}'. Automatic synthesis was "
        "unavailable, so the top findings are listed verbatim from the gathered sources."
    )
    confidence = round(min(0.6, CONFIDENCE_BASELINE + diversity / 3), 3)
    score = int(min(SCORE_MAX, len(pages) * 6 + diversity * 30))
    return Orchestration(
        summary=summary,
        findings=findings,
        gaps=["Automatic critique was unavailable for this run."],
        confidence=confidence,
        source_diversity=diversity,
        score=score,
    )


async def orchestrate(question: str, pages: list, api_key: str, emit) -> Orchestration:
    diversity = source_diversity(pages)
    if not pages:
        return Orchestration(gaps=["No sources were gathered."], source_diversity=0.0)
    context = _build_context(pages)
    try:
        emit({"type": "agent", "agent": "summarizer", "message": "Synthesising findings"})
        summary_raw = await _complete(
            [
                {"role": "system", "content": SUMMARIZER_PROMPT},
                {
                    "role": "user",
                    "content": f"QUESTION: {question}\n\nSOURCES:\n{context}",
                },
            ],
            api_key,
            SUMMARY_MAX_TOKENS,
        )
        parsed = _parse_json(summary_raw)
        summary = str(parsed.get("summary", "")).strip()
        findings = [f for f in (parsed.get("findings") or []) if isinstance(f, dict)]
        if not summary and not findings:
            return _heuristic(question, pages)

        emit({"type": "agent", "agent": "critic", "message": "Reviewing for gaps"})
        gaps_raw = await _complete(
            [
                {"role": "system", "content": CRITIC_PROMPT},
                {
                    "role": "user",
                    "content": (
                        f"QUESTION: {question}\n\nSUMMARY: {summary}\n\n"
                        f"FINDINGS: {json.dumps(findings)[:4000]}"
                    ),
                },
            ],
            api_key,
            CRITIC_MAX_TOKENS,
        )
        gaps = [str(g).strip() for g in (_parse_json(gaps_raw).get("gaps") or []) if str(g).strip()]

        emit({"type": "agent", "agent": "linker", "message": "Scoring confidence"})
        link_raw = await _complete(
            [
                {"role": "system", "content": LINKER_PROMPT},
                {
                    "role": "user",
                    "content": (
                        f"FINDINGS: {json.dumps(findings)[:4000]}\n\nSOURCES:\n{context[:4000]}"
                    ),
                },
            ],
            api_key,
            LINKER_MAX_TOKENS,
        )
        try:
            confidence = float(_parse_json(link_raw).get("confidence", 0.0))
        except (TypeError, ValueError):
            confidence = 0.0
        confidence = round(max(CONFIDENCE_BASELINE, min(1.0, confidence)), 3)
        coverage = min(1.0, len(pages) / 10.0)
        score = int(
            min(SCORE_MAX, (confidence * 0.5 + diversity * 0.3 + coverage * 0.2) * SCORE_MAX)
        )
        return Orchestration(
            summary=summary,
            findings=findings,
            gaps=gaps,
            confidence=confidence,
            source_diversity=diversity,
            score=score,
        )
    except Exception as exc:
        logger.warning("deepsearch orchestration failed, using heuristic: %s", exc)
        return _heuristic(question, pages)
