# retoor <retoor@molodetz.nl>

from __future__ import annotations

import asyncio
import json
import logging
import shutil
import sys
from pathlib import Path

from devplacepy.config import BASE_DIR, DEEPSEARCH_DIR, INTERNAL_MODEL
from devplacepy.services.base import ConfigField
from devplacepy.services.deepsearch.store import VectorStore
from devplacepy.services.jobs.base import JobService

from .progress import hub

logger = logging.getLogger(__name__)

WORKER_MODULE = "devplacepy.services.jobs.deepsearch.worker"
STREAM_LIMIT = 16 * 1024 * 1024


class DeepsearchService(JobService):
    kind = "deepsearch"
    title = "DeepSearch"
    description = (
        "Runs a multi-agent web research job: it plans queries, crawls and indexes "
        "sources into a per-session vector collection, then synthesises a cited report "
        "with confidence scoring and source diversity, streaming live "
        "progress over a websocket."
    )

    def __init__(self):
        super().__init__(name="deepsearch", interval_seconds=2)
        self.config_fields = list(self.config_fields) + [
            ConfigField(
                "deepsearch_model",
                "AI model",
                type="str",
                default="",
                help="Model sent to the internal AI gateway for query planning and report synthesis. Blank uses the gateway default model.",
                group="DeepSearch",
            ),
        ]

    def session_dir(self, uid: str) -> Path:
        return DEEPSEARCH_DIR / uid

    def collection_name(self, uid: str) -> str:
        return f"ds_{uid.replace('-', '')}"

    async def process(self, job: dict) -> dict:
        from devplacepy import database
        from devplacepy.services.audit import record as audit

        uid = job["uid"]
        payload = dict(job.get("payload", {}))
        query = payload.get("query", "")
        output_dir = self.session_dir(uid)
        await asyncio.to_thread(output_dir.mkdir, parents=True, exist_ok=True)
        await asyncio.to_thread(
            (output_dir / "control.json").write_text,
            json.dumps({"state": "running"}),
            encoding="utf-8",
        )
        payload["collection"] = self.collection_name(uid)
        payload["model"] = self.get_config()["deepsearch_model"] or INTERNAL_MODEL
        payload["cached_hashes"] = self._cached_hashes(database)
        payload_path = output_dir / "payload.json"
        await asyncio.to_thread(payload_path.write_text, json.dumps(payload), encoding="utf-8")

        actor_kind = job.get("owner_kind") or "system"
        actor_uid = job.get("owner_id") if job.get("owner_kind") == "user" else None
        database.update_deepsearch_session(uid, {"status": "running"})

        try:
            summary = await self._run_worker(
                uid, payload_path, output_dir, actor_kind, job.get("owner_id") or ""
            )
        except Exception as exc:
            hub.publish(uid, {"type": "failed", "message": str(exc)[:300]})
            database.update_deepsearch_session(uid, {"status": "failed"})
            audit.record_system(
                "deepsearch.run.failed",
                actor_kind=actor_kind,
                actor_uid=actor_uid,
                result="failure",
                summary=f"DeepSearch for {query} failed",
                metadata={"query": query, "error": str(exc)[:200]},
                links=[audit.job(uid)],
            )
            await asyncio.to_thread(shutil.rmtree, output_dir, ignore_errors=True)
            raise

        report = await self._load_report(output_dir)
        await self._persist_cache(database, output_dir)
        from datetime import datetime, timezone

        database.update_deepsearch_session(
            uid,
            {
                "status": "done",
                "score": int(summary.get("score") or 0),
                "confidence": float(summary.get("confidence") or 0.0),
                "source_diversity": float(summary.get("source_diversity") or 0.0),
                "page_count": int(summary.get("page_count") or 0),
                "chunk_count": int(summary.get("chunk_count") or 0),
                "summary": (report.get("summary") or "")[:4000],
                "completed_at": datetime.now(timezone.utc).isoformat(),
            },
        )
        hub.publish(
            uid,
            {
                "type": "done",
                "score": summary.get("score"),
                "confidence": summary.get("confidence"),
                "source_diversity": summary.get("source_diversity"),
                "page_count": summary.get("page_count"),
                "chunk_count": summary.get("chunk_count"),
                "session_url": f"/tools/deepsearch/{uid}/session",
            },
        )
        hub.clear(uid)
        audit.record_system(
            "deepsearch.run.complete",
            actor_kind=actor_kind,
            actor_uid=actor_uid,
            summary=f"DeepSearch for {query} scored {summary.get('score')}",
            metadata={
                "query": query,
                "score": summary.get("score"),
                "page_count": summary.get("page_count"),
            },
            links=[audit.job(uid)],
        )
        return {
            "query": query,
            "score": summary.get("score", 0),
            "confidence": summary.get("confidence", 0.0),
            "source_diversity": summary.get("source_diversity", 0.0),
            "page_count": summary.get("page_count", 0),
            "chunk_count": summary.get("chunk_count", 0),
            "report": report,
            "collection": self.collection_name(uid),
            "session_url": f"/tools/deepsearch/{uid}/session",
            "bytes_in": 0,
            "bytes_out": len(json.dumps(report)) if report else 0,
            "item_count": summary.get("page_count", 0),
        }

    def _cached_hashes(self, database) -> list[str]:
        if "deepsearch_url_cache" not in database.db.tables:
            return []
        return [
            row.get("url_hash", "")
            for row in database.get_table("deepsearch_url_cache").find()
            if row.get("url_hash")
        ]

    async def _persist_cache(self, database, output_dir: Path) -> None:
        path = output_dir / "url_cache.json"
        if not path.is_file():
            return
        try:
            text = await asyncio.to_thread(path.read_text, encoding="utf-8")
            entries = json.loads(text)
        except (ValueError, OSError):
            return
        for entry in entries:
            database.upsert_deepsearch_url_cache(
                entry.get("url_hash", ""),
                entry.get("url", ""),
                entry.get("title", ""),
                entry.get("content_hash", ""),
                int(entry.get("status") or 0),
                int(entry.get("byte_size") or 0),
            )

    def _ledger_rsearch(self, owner_kind: str, owner_id: str, frame: dict) -> None:
        from devplacepy.services.openai_gateway.usage import record_rsearch_call

        success = bool(frame.get("success"))
        record_rsearch_call(
            owner_kind,
            owner_id,
            frame.get("endpoint") or "/search",
            success,
            200 if success else 0,
        )

    async def _run_worker(
        self,
        uid: str,
        payload_path: Path,
        output_dir: Path,
        owner_kind: str = "system",
        owner_id: str = "",
    ) -> dict:
        proc = await asyncio.create_subprocess_exec(
            sys.executable,
            "-m",
            WORKER_MODULE,
            str(payload_path),
            str(output_dir),
            cwd=str(BASE_DIR),
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE,
            limit=STREAM_LIMIT,
        )
        summary: dict = {}
        worker_error = ""
        err = ""
        stderr_task = asyncio.create_task(proc.stderr.read())
        try:
            while True:
                line = await proc.stdout.readline()
                if not line:
                    break
                try:
                    frame = json.loads(line.decode("utf-8", "replace"))
                except (ValueError, TypeError):
                    continue
                hub.publish(uid, frame)
                if frame.get("type") == "report_ready":
                    summary = frame
                elif frame.get("type") == "rsearch":
                    self._ledger_rsearch(owner_kind, owner_id, frame)
                elif frame.get("type") == "error":
                    worker_error = frame.get("message", "worker error")
            err = (await stderr_task).decode("utf-8", "replace")
            await proc.wait()
        finally:
            if proc.returncode is None:
                try:
                    proc.kill()
                except ProcessLookupError:
                    pass
                await proc.wait()
            if not stderr_task.done():
                stderr_task.cancel()
        if proc.returncode != 0 or not summary:
            raise RuntimeError(
                worker_error or err[:500] or f"deepsearch worker exited {proc.returncode}"
            )
        return summary

    async def _load_report(self, output_dir: Path) -> dict:
        report_path = output_dir / "report.json"
        if not report_path.is_file():
            return {}
        try:
            text = await asyncio.to_thread(report_path.read_text, encoding="utf-8")
            return json.loads(text)
        except (ValueError, OSError):
            return {}

    def cleanup(self, job: dict) -> None:
        uid = job["uid"]
        hub.clear(uid)
        VectorStore(self.collection_name(uid)).drop()
        shutil.rmtree(self.session_dir(uid), ignore_errors=True)
