# 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.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_queries
from .orchestrate import orchestrate

CONTROL_FILE = "control.json"
PAUSE_POLL_SECONDS = 1.0
EMBED_BATCH = 64


def _emit(frame: dict) -> None:
    sys.stdout.write(json.dumps(frame, ensure_ascii=False) + "\n")
    sys.stdout.flush()


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
) -> 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,
                )
            )
    if not chunks:
        return 0, "empty"
    backend = "gateway"
    for start in range(0, len(chunks), EMBED_BATCH):
        batch = chunks[start : start + EMBED_BATCH]
        result = await embed_texts([chunk.text for chunk in batch], api_key)
        if not result.vectors:
            result = local_embed([chunk.text for chunk in batch])
        backend = result.backend
        store.add(batch, result.vectors)
    return len(chunks), 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", "")
    collection = payload.get("collection", "")
    cached_hashes = set(payload.get("cached_hashes", []))
    should_stop = _make_stop(output_dir)

    _emit({"type": "stage", "stage": "planning", "message": "Planning research queries"})
    queries = await plan_queries(query, api_key)
    _emit({"type": "queries", "queries": queries})

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

    _emit({"type": "stage", "stage": "crawling", "message": "Crawling sources"})
    outcome = await crawl(
        candidates,
        max_pages,
        _emit,
        lambda url: url_hash(url) in cached_hashes,
        should_stop,
    )

    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)
    _emit({"type": "stage", "stage": "indexing", "message": "Indexing content"})
    chunk_count, embed_backend = await _index_chunks(store, outcome.pages, api_key)

    _emit({"type": "stage", "stage": "analysis", "message": "Running research agents"})
    result = await orchestrate(query, outcome.pages, api_key, _emit)

    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,
        "gaps": result.gaps,
        "sources": sources,
        "score": result.score,
        "confidence": result.confidence,
        "source_diversity": result.source_diversity,
        "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"],
            "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))
