# retoor <retoor@molodetz.nl>

from __future__ import annotations

import asyncio
import json
import sys
from datetime import datetime, timezone
from pathlib import Path

from devplacepy.config import INTERNAL_MODEL
from devplacepy.services.deepsearch.embeddings import embed_texts, local_embed
from devplacepy.services.deepsearch.store import Chunk, VectorStore
from devplacepy.utils import generate_uid

from .chunking import chunk_text
from .crawl import content_hash, crawl, search_queries, url_hash
from .enhance import plan_followup_queries, plan_queries
from .orchestrate import orchestrate
from .phases import (
    PHASE_ANALYSIS,
    PHASE_CRAWLING,
    PHASE_INDEXING,
    PHASE_LABELS,
    PHASE_PLANNING,
    PHASE_SEARCHING,
    PHASE_SYNTHESIS,
    TOTAL_PHASES,
    phase_index,
)

CONTROL_FILE = "control.json"
PAUSE_POLL_SECONDS = 1.0
EMBED_BATCH = 64
FRAME_VERSION = 1
MAX_REFINE_ROUNDS = 2
REFINE_CRAWL_DEPTH = 1

_first_frame_sent = False


def _emit(frame: dict) -> None:
    global _first_frame_sent
    if not _first_frame_sent:
        frame = {"version": FRAME_VERSION, **frame}
        _first_frame_sent = True
    sys.stdout.write(json.dumps(frame, ensure_ascii=False) + "\n")
    sys.stdout.flush()


def _stage(stage: str, message: str, phase: str) -> None:
    _emit({"type": "stage", "stage": stage, "message": message})
    _emit(
        {
            "type": "phase",
            "phase": phase,
            "index": phase_index(phase),
            "total": TOTAL_PHASES,
            "label": PHASE_LABELS.get(phase, phase),
            "message": message,
        }
    )


def _read_control(output_dir: Path) -> str:
    path = output_dir / CONTROL_FILE
    if not path.is_file():
        return "running"
    try:
        return str(json.loads(path.read_text(encoding="utf-8")).get("state", "running"))
    except (ValueError, OSError):
        return "running"


def _make_stop(output_dir: Path):
    async def should_stop() -> bool:
        while True:
            state = _read_control(output_dir)
            if state == "cancelled":
                return True
            if state == "paused":
                _emit({"type": "stage", "stage": "paused", "message": "Paused"})
                await asyncio.sleep(PAUSE_POLL_SECONDS)
                continue
            return False

    return should_stop


async def _index_chunks(
    store: VectorStore, pages: list, api_key: str, emit
) -> tuple[int, str]:
    chunks: list[Chunk] = []
    for page in pages:
        for position, text in enumerate(chunk_text(page.text)):
            chunks.append(
                Chunk(
                    uid=generate_uid(),
                    text=text,
                    url=page.url,
                    title=page.title,
                    depth=page.depth,
                    source=page.source,
                    position=position,
                )
            )
    total = len(chunks)
    if not chunks:
        emit({"type": "embed_done", "backend": "empty", "chunk_count": 0})
        return 0, "empty"
    total_batches = (total + EMBED_BATCH - 1) // EMBED_BATCH
    backend = "gateway"
    forced_local = False
    done = 0
    collected: list[list[float]] = []
    for batch_no, start in enumerate(range(0, total, EMBED_BATCH), start=1):
        batch = chunks[start : start + EMBED_BATCH]
        emit(
            {
                "type": "embed_batch",
                "batch": batch_no,
                "total_batches": total_batches,
                "backend": "local" if forced_local else backend,
                "chunks": len(batch),
                "done": done,
                "total": total,
                "message": f"Embedding batch {batch_no}/{total_batches}",
            }
        )
        if forced_local:
            result = local_embed([chunk.text for chunk in batch])
        else:
            result = await embed_texts([chunk.text for chunk in batch], api_key)
            if not result.vectors or result.backend != "gateway":
                forced_local = True
                result = local_embed([chunk.text for chunk in batch])
        backend = result.backend
        collected.extend(result.vectors)
        done += len(batch)
        emit(
            {
                "type": "embed_batch",
                "batch": batch_no,
                "total_batches": total_batches,
                "backend": backend,
                "chunks": len(batch),
                "done": done,
                "total": total,
                "message": f"Embedded batch {batch_no}/{total_batches} ({backend})",
            }
        )
    if forced_local:
        collected = local_embed([chunk.text for chunk in chunks]).vectors
        backend = "local"
    await store.add(chunks, collected)
    final_backend = "local" if forced_local else backend
    emit({"type": "embed_done", "backend": final_backend, "chunk_count": total})
    return total, final_backend


async def _run(payload: dict, output_dir: Path) -> dict:
    query = payload.get("query", "")
    max_pages = int(payload.get("max_pages", 12))
    depth = int(payload.get("depth", 2))
    api_key = payload.get("api_key", "")
    model = payload.get("model") or INTERNAL_MODEL
    collection = payload.get("collection", "")
    cached_hashes = set(payload.get("cached_hashes", []))
    should_stop = _make_stop(output_dir)

    _stage("planning", "Planning research queries", PHASE_PLANNING)
    queries = await plan_queries(query, api_key, _emit, model=model)
    _emit({"type": "queries", "queries": queries})

    _stage("searching", "Searching the web", PHASE_SEARCHING)
    candidates = await search_queries(queries, _emit)
    _emit({"type": "candidates", "count": len(candidates)})

    _stage("crawling", "Crawling sources", PHASE_CRAWLING)
    outcome = await crawl(
        candidates,
        max_pages,
        _emit,
        lambda url: url_hash(url) in cached_hashes,
        should_stop,
        query=query,
        depth=depth,
    )

    round_no = 0
    while (
        len(outcome.pages) < max_pages
        and round_no < MAX_REFINE_ROUNDS
        and not await should_stop()
    ):
        covered_titles = [page.title for page in outcome.pages if page.title]
        followups = await plan_followup_queries(
            query, covered_titles, api_key, _emit, model=model
        )
        if not followups:
            break
        round_no += 1
        _emit(
            {
                "type": "substep",
                "phase": PHASE_SEARCHING,
                "message": f"Refinement round {round_no}: filling gaps with {len(followups)} more queries",
            }
        )
        _emit({"type": "queries", "queries": followups, "round": round_no})
        more_candidates = await search_queries(followups, _emit)
        existing_urls = {page.url for page in outcome.pages}
        more_candidates = [c for c in more_candidates if c["url"] not in existing_urls]
        if not more_candidates:
            break
        _emit({"type": "candidates", "count": len(more_candidates), "round": round_no})
        extra = await crawl(
            more_candidates,
            max_pages - len(outcome.pages),
            _emit,
            lambda url: url_hash(url) in cached_hashes,
            should_stop,
            query=query,
            depth=REFINE_CRAWL_DEPTH,
            seen_hashes=outcome.seen_hashes,
        )
        if not extra.pages:
            break
        outcome.pages.extend(extra.pages)

    new_cache = [
        {
            "url_hash": url_hash(page.url),
            "url": page.url,
            "title": page.title,
            "content_hash": content_hash(page.text),
            "status": page.status,
            "byte_size": len(page.text),
        }
        for page in outcome.pages
    ]

    store = VectorStore(collection)
    _stage("indexing", "Indexing content", PHASE_INDEXING)
    chunk_count, embed_backend = await _index_chunks(
        store, outcome.pages, api_key, _emit
    )

    _stage("analysis", "Running research agents", PHASE_ANALYSIS)
    result = await orchestrate(
        query, outcome.pages, api_key, _emit, store=store, queries=queries, model=model
    )

    _stage("synthesis", "Compiling cited report", PHASE_SYNTHESIS)

    sources = [
        {"url": page.url, "title": page.title, "source": page.source}
        for page in outcome.pages
    ]
    report = {
        "query": query,
        "depth": depth,
        "max_pages": max_pages,
        "generated_at": datetime.now(timezone.utc).isoformat(),
        "summary": result.summary,
        "findings": result.findings,
        "sources": sources,
        "score": result.score,
        "confidence": result.confidence,
        "source_diversity": result.source_diversity,
        "synthesis": result.synthesis,
        "follow_up_questions": result.follow_up_questions,
        "page_count": len(outcome.pages),
        "chunk_count": chunk_count,
        "embed_backend": embed_backend,
        "collection": collection,
    }
    (output_dir / "report.json").write_text(
        json.dumps(report, ensure_ascii=False), encoding="utf-8"
    )
    (output_dir / "url_cache.json").write_text(
        json.dumps(new_cache, ensure_ascii=False), encoding="utf-8"
    )
    _emit(
        {
            "type": "report_ready",
            "score": report["score"],
            "confidence": report["confidence"],
            "source_diversity": report["source_diversity"],
            "synthesis": report["synthesis"],
            "page_count": report["page_count"],
            "chunk_count": report["chunk_count"],
        }
    )
    return report


def main(argv: list) -> int:
    if len(argv) != 3:
        sys.stderr.write("usage: deepsearch.worker <payload_json> <output_dir>\n")
        return 2
    payload = json.loads(Path(argv[1]).read_text(encoding="utf-8"))
    output_dir = Path(argv[2])
    output_dir.mkdir(parents=True, exist_ok=True)
    try:
        asyncio.run(_run(payload, output_dir))
    except Exception as exc:
        _emit({"type": "error", "message": str(exc)[:500]})
        return 1
    return 0


if __name__ == "__main__":
    raise SystemExit(main(sys.argv))
