src/ragforge/workspace.py

from __future__ import annotations

import hashlib
import json
import shutil
import threading
import time
import uuid
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Callable

from .chunking import chunk_documents
from .config import get_settings
from .context_budget import corpus_scale_label
from .corpus import build_source_profiles, corpus_manifest, profile_chunks
from .llm import GeminiGateway
from .json_utils import to_jsonable
from .loaders import DocumentLoader
from .retrieval import HybridRetriever
from .schemas import CorpusSummary, Document, SearchHit, SourceProfile
from .sql_agent import SQLWorkspace


SERVER_BOOT_ID = uuid.uuid4().hex[:12]


class Workspace:
    def __init__(self, session_id: str):
        self.session_id = session_id
        self.created_at = time.time()
        self.last_access = self.created_at
        self.version = 0
        self.settings = get_settings()
        self.dir = self.settings.data_dir / session_id
        self.dir.mkdir(parents=True, exist_ok=True)
        self.documents: list[Document] = []
        self.sources: list[str] = []
        self.chunks = []
        self.retriever = HybridRetriever("chunks")
        self.source_profiles: dict[str, SourceProfile] = {}
        self.source_retriever = HybridRetriever("source_profiles")
        self.sql = SQLWorkspace()
        self.history: list[dict[str, str]] = []
        self.ingested_hashes: set[str] = set()
        self.evaluation_reports: dict[str, dict[str, Any]] = {}
        self.lock = threading.RLock()

    def touch(self) -> None:
        self.last_access = time.time()

    def ingest(
        self,
        paths: list[Path],
        ocr: bool = False,
        semantic_chunking: bool = False,
        api_key: str | None = None,
        model: str | None = None,
        progress_callback: Callable[[float, str], None] | None = None,
    ) -> CorpusSummary:
        with self.lock:
            self.touch()
            notify = progress_callback or (lambda _progress, _message: None)
            notify(0.02, "Preparing corpus inputs")
            gateway = GeminiGateway(api_key, model) if ocr else None
            loader = DocumentLoader(gateway)
            expanded = loader.expand_inputs(paths, self.dir)
            notify(0.08, f"Found {len(expanded)} supported file(s)")
            new_docs: list[Document] = []
            total_files = max(1, len(expanded))
            for file_idx, path in enumerate(expanded, start=1):
                notify(0.10 + 0.42 * (file_idx - 1) / total_files, f"Parsing {path.name} ({file_idx}/{len(expanded)})")
                digest = hashlib.sha256(path.read_bytes()).hexdigest()
                if digest in self.ingested_hashes:
                    continue
                docs, tables = loader.load(path)
                new_docs.extend(docs)
                for name, df in tables:
                    self.sql.add_dataframe(name, df)
                self.ingested_hashes.add(digest)

            notify(0.54, "Chunking documents")
            self.documents.extend(new_docs)
            self.sources = sorted(set(self.sources + [d.source for d in new_docs]))
            self.chunks = chunk_documents(self.documents, semantic=semantic_chunking)
            notify(0.66, f"Building hybrid chunk index ({len(self.chunks)} chunks)")
            self.retriever.index(self.chunks)

            # Build a second, source-level representation. This powers semantic
            # source selection and corpus overviews without allowing a 48-page
            # PDF to swamp a five-file corpus simply because it produced more
            # chunks than the other sources.
            notify(0.82, "Building source profiles")
            self.source_profiles = build_source_profiles(self.documents, self.chunks)
            notify(0.92, f"Building source index ({len(self.source_profiles)} sources)")
            self.source_retriever.index(profile_chunks(self.source_profiles))
            self.version += 1
            notify(1.0, "Corpus ready")
            return self.summary()


    @property
    def is_empty(self) -> bool:
        return not self.chunks and not self.sql.tables

    def stats(self) -> dict[str, object]:
        return {
            "session_id": self.session_id,
            "version": self.version,
            "documents": len(self.documents),
            "chunks": len(self.chunks),
            "source_profiles": len(self.source_profiles),
            "sources": len(self.sources),
            "tables": len(self.sql.tables),
            "table_names": list(self.sql.tables),
            "saved_evaluations": sorted(self.evaluation_reports),
            "status": "empty" if self.is_empty else "ready",
            "corpus_scale": corpus_scale_label(len(self.chunks), len(self.sources)),
            "server_boot_id": SERVER_BOOT_ID,
        }

    def health_snapshot(self) -> dict[str, object]:
        """Operational snapshot for the Architecture/API runtime view."""
        now = time.time()
        chunk_chars = sum(len(getattr(chunk, "text", "") or "") for chunk in self.chunks)
        vector_bytes = int(getattr(getattr(self.retriever, "_vectors", None), "nbytes", 0) or 0)
        source_vector_bytes = int(getattr(getattr(self.source_retriever, "_vectors", None), "nbytes", 0) or 0)
        text_bytes = chunk_chars
        estimated_index_mb = (vector_bytes + source_vector_bytes + text_bytes) / (1024 * 1024)
        max_chunks = max(1, int(self.settings.max_chunks_per_session))
        utilization = len(self.chunks) / max_chunks
        return {
            **self.stats(),
            "session_age_minutes": round((now - self.created_at) / 60.0, 1),
            "idle_minutes": round((now - self.last_access) / 60.0, 1),
            "session_ttl_minutes": int(self.settings.session_ttl_minutes),
            "max_chunks_per_session": max_chunks,
            "chunk_capacity_utilization": round(utilization, 3),
            "capacity_status": "warning" if utilization >= 0.80 else "ok",
            "chunk_text_chars": chunk_chars,
            "vector_bytes": vector_bytes,
            "source_vector_bytes": source_vector_bytes,
            "estimated_index_memory_mb": round(estimated_index_mb, 2),
            "chunk_index_ready": bool(getattr(self.retriever, "_ready", False)),
            "source_index_ready": bool(getattr(self.source_retriever, "_ready", False)),
            "evaluation_history_runs": len(self.evaluation_history_inventory()),
            "adaptive_policy": {
                "adaptive_top_k": True,
                "focused_context_budget": "2-5 chunks based on scale/confidence",
                "focused_evidence_compression": True,
                "small_corpus_reranker_skip": len(self.chunks) < 250 and len(self.sources) < 10,
            },
        }

    @property
    def evaluation_dir(self) -> Path:
        path = self.dir / "evaluations"
        path.mkdir(parents=True, exist_ok=True)
        return path

    @property
    def evaluation_history_dir(self) -> Path:
        path = self.evaluation_dir / "history"
        path.mkdir(parents=True, exist_ok=True)
        return path

    def save_evaluation(
        self,
        level: str,
        report: dict[str, Any],
        *,
        model: str,
        benchmark_version: str,
    ) -> dict[str, Any]:
        """Persist the latest evaluation for one depth within this workspace.

        Evaluation reports are deliberately separate from the response cache.
        They survive browser refreshes while the Hugging Face container is
        alive, but remain ephemeral with the rest of the workspace storage.
        """
        with self.lock:
            self.touch()
            saved = to_jsonable(report)
            saved["evaluation_cache"] = {
                "level": level,
                "model": model,
                "benchmark_version": benchmark_version,
                "workspace_version": self.version,
                "saved_at": datetime.now(timezone.utc).isoformat(),
                "run_id": uuid.uuid4().hex[:12],
                "server_boot_id": SERVER_BOOT_ID,
            }
            self.evaluation_reports[level] = saved
            try:
                target = self.evaluation_dir / f"{level.lower()}.json"
                target.write_text(json.dumps(saved, indent=2, ensure_ascii=False), encoding="utf-8")
                stamp = saved["evaluation_cache"]["saved_at"].replace(":", "-").replace("+", "_")
                archive = self.evaluation_history_dir / f"{stamp}_{level.lower()}.json"
                archive.write_text(json.dumps(saved, indent=2, ensure_ascii=False), encoding="utf-8")
            except Exception:
                # In-memory history is still useful even if persistence fails.
                pass
            return saved

    def get_evaluation(
        self,
        level: str,
        *,
        model: str | None = None,
        benchmark_version: str | None = None,
        require_current_corpus: bool = True,
    ) -> dict[str, Any] | None:
        with self.lock:
            self.touch()
            report = self.evaluation_reports.get(level)
            if report is None:
                path = self.evaluation_dir / f"{level.lower()}.json"
                if path.exists():
                    try:
                        report = to_jsonable(json.loads(path.read_text(encoding="utf-8")))
                        self.evaluation_reports[level] = report
                    except Exception:
                        report = None
            if not report:
                return None
            report = to_jsonable(report)
            self.evaluation_reports[level] = report
            meta = report.get("evaluation_cache", {})
            if require_current_corpus and int(meta.get("workspace_version", -1)) != int(self.version):
                return None
            if model and meta.get("model") != model:
                return None
            if benchmark_version and meta.get("benchmark_version") != benchmark_version:
                return None
            return report

    def evaluation_inventory(self) -> list[dict[str, Any]]:
        rows: list[dict[str, Any]] = []
        for level in ("Quick", "Standard", "Deep"):
            report = self.get_evaluation(level, require_current_corpus=False)
            if not report:
                continue
            meta = report.get("evaluation_cache", {})
            summary = report.get("summary", {})
            rows.append(
                {
                    "level": level,
                    "grade": summary.get("quality_grade", "-"),
                    "score": summary.get("deterministic_quality_score"),
                    "model": meta.get("model", "-"),
                    "benchmark": meta.get("benchmark_version", "-"),
                    "workspace_version": meta.get("workspace_version"),
                    "current_corpus": int(meta.get("workspace_version", -1)) == int(self.version),
                    "saved_at": meta.get("saved_at", ""),
                    "run_id": meta.get("run_id", ""),
                    "server_boot_id": meta.get("server_boot_id", ""),
                    "current_server": meta.get("server_boot_id") in {None, "", SERVER_BOOT_ID},
                }
            )
        return rows

    def evaluation_history_inventory(self, limit: int = 50) -> list[dict[str, Any]]:
        rows: list[dict[str, Any]] = []
        try:
            paths = sorted(self.evaluation_history_dir.glob("*.json"))[-max(1, int(limit)) :]
        except Exception:
            paths = []
        previous_by_level: dict[str, dict[str, Any]] = {}
        for path in paths:
            try:
                report = to_jsonable(json.loads(path.read_text(encoding="utf-8")))
                meta = report.get("evaluation_cache", {})
                summary = report.get("summary", {})
                level = str(meta.get("level", summary.get("evaluation_level", "-")))
                row = {
                    "saved_at": meta.get("saved_at", ""),
                    "level": level,
                    "benchmark": meta.get("benchmark_version", summary.get("benchmark_version", "-")),
                    "model": meta.get("model", "-"),
                    "workspace_version": meta.get("workspace_version"),
                    "grade": summary.get("quality_grade", "-"),
                    "score": summary.get("deterministic_quality_score"),
                    "citation_coverage": summary.get("citation_coverage"),
                    "hard_mode_pass": summary.get("hard_mode_pass_rate"),
                    "p50_ms": summary.get("latency_p50_ms"),
                    "gemini_requests": summary.get("gemini_requests"),
                    "run_id": meta.get("run_id", ""),
                    "server_boot_id": meta.get("server_boot_id", ""),
                }
                prev = previous_by_level.get(level)
                if prev:
                    try:
                        row["delta_score"] = round(float(row.get("score") or 0) - float(prev.get("score") or 0), 3)
                        row["delta_p50_ms"] = round(float(row.get("p50_ms") or 0) - float(prev.get("p50_ms") or 0), 1)
                    except Exception:
                        row["delta_score"] = None
                        row["delta_p50_ms"] = None
                else:
                    row["delta_score"] = None
                    row["delta_p50_ms"] = None
                previous_by_level[level] = row
                rows.append(row)
            except Exception:
                continue
        return rows

    def manifest(self, max_chars: int = 9000, include_excerpts: bool = True) -> str:
        base = corpus_manifest(
            self.source_profiles,
            list(self.sql.tables),
            max_chars=max_chars,
            include_excerpts=include_excerpts,
        )
        schema = self.sql.schema_text().strip()
        if schema:
            base += "\nStructured table schemas:\n" + schema
        return base[:max_chars]

    def select_source_hits(self, query: str, limit: int = 5) -> list[SearchHit]:
        if not self.source_profiles:
            return []
        return self.source_retriever.search(
            query,
            top_k=min(max(1, limit), len(self.source_profiles)),
            use_reranker=False,
        )

    def select_sources(self, query: str, limit: int = 5) -> list[str]:
        hits = self.select_source_hits(query, limit)
        sources: list[str] = []
        for hit in hits:
            if hit.chunk.source not in sources:
                sources.append(hit.chunk.source)
        return sources

    def global_evidence(self, query: str, top_k: int, use_reranker: bool = True) -> list[SearchHit]:
        """Stable source-balanced evidence for corpus overview/synthesis tasks.

        Source profiles decide *which sources* matter. Evidence then comes from
        deterministic representative original chunks, rather than asking an
        abstract overview query to choose an arbitrary page from every source.
        """
        if not self.source_profiles:
            return []
        source_limit = min(max(1, top_k), len(self.source_profiles))
        source_hits = self.select_source_hits(query, source_limit)
        if not source_hits:
            return []

        selected: list[SearchHit] = []
        round_idx = 0
        while len(selected) < top_k:
            added = False
            for source_hit in source_hits:
                profile = self.source_profiles.get(source_hit.chunk.source)
                if not profile or round_idx >= len(profile.representative_chunk_ids):
                    continue
                chunk_id = profile.representative_chunk_ids[round_idx]
                chunk = self.retriever.chunk_by_id.get(chunk_id)
                if not chunk:
                    continue
                decay = 1.0 / (1.0 + 0.15 * round_idx)
                selected.append(
                    SearchHit(
                        chunk=chunk,
                        score=float(source_hit.score) * decay,
                        dense_score=(float(source_hit.dense_score) * decay if source_hit.dense_score is not None else None),
                        sparse_score=(float(source_hit.sparse_score) * decay if source_hit.sparse_score is not None else None),
                    )
                )
                added = True
                if len(selected) >= top_k:
                    break
            if not added:
                break
            round_idx += 1

        # Reranking can improve ordering, but keep the source-balanced selection
        # itself intact. Raw reranker logits are never used as confidence.
        if use_reranker and selected:
            try:
                from .retrieval import ModelRegistry

                scores = list(ModelRegistry.reranker().rerank(query, [h.chunk.text for h in selected]))
                for hit, score in zip(selected, scores):
                    hit.rerank_score = float(score)
            except Exception:
                pass
        return selected[:top_k]

    def hierarchical_evidence(
        self,
        query: str,
        top_k: int,
        use_reranker: bool = True,
        source_limit: int = 4,
        diversify: bool = False,
    ) -> tuple[list[SearchHit], list[str]]:
        """Retrieve source profiles first, then chunks only from selected sources."""
        selected_sources = self.select_sources(query, min(source_limit, max(1, len(self.source_profiles))))
        hits = self.retriever.search(
            query,
            top_k=max(top_k * 2, 8),
            use_reranker=use_reranker,
            allowed_sources=selected_sources,
        )
        if diversify:
            # Re-run a source-balanced evidence pass for synthesis/comparison so
            # multiple selected sources are represented when relevant.
            diverse = self.retriever.source_balanced_search(
                query,
                top_k=top_k,
                sources=selected_sources,
                per_source=1,
                use_reranker=use_reranker,
            )
            if diverse:
                hits = diverse
        return hits[:top_k], selected_sources

    def reset(self) -> None:
        try:
            shutil.rmtree(self.dir, ignore_errors=True)
        finally:
            self.__init__(self.session_id)

    def summary(self) -> CorpusSummary:
        return CorpusSummary(
            session_id=self.session_id,
            documents=len(self.documents),
            chunks=len(self.chunks),
            tables=list(self.sql.tables),
            sources=list(self.sources),
            source_profiles=len(self.source_profiles),
        )


class WorkspaceRegistry:
    def __init__(self):
        self.settings = get_settings()
        self._items: dict[str, Workspace] = {}
        self._lock = threading.RLock()

    def create(self) -> Workspace:
        with self._lock:
            self.cleanup()
            session_id = uuid.uuid4().hex
            ws = Workspace(session_id)
            self._items[session_id] = ws
            return ws

    def get(self, session_id: str | None) -> Workspace:
        """UI-friendly lookup: return an existing workspace or create a fresh one."""
        with self._lock:
            self.cleanup()
            if session_id and session_id in self._items:
                ws = self._items[session_id]
                ws.touch()
                return ws
            return self.create()

    def contains(self, session_id: str | None) -> bool:
        with self._lock:
            self.cleanup()
            return bool(session_id and session_id in self._items)

    def require(self, session_id: str) -> Workspace:
        """API lookup: never silently replace a missing/expired client session id."""
        with self._lock:
            self.cleanup()
            ws = self._items.get(session_id)
            if ws is None:
                raise KeyError("Unknown or expired session_id; create a new session first")
            ws.touch()
            return ws

    def delete(self, session_id: str) -> None:
        with self._lock:
            ws = self._items.pop(session_id, None)
            if ws:
                shutil.rmtree(ws.dir, ignore_errors=True)

    def cleanup(self) -> None:
        cutoff = time.time() - self.settings.session_ttl_minutes * 60
        stale = [sid for sid, ws in self._items.items() if ws.last_access < cutoff]
        for sid in stale:
            ws = self._items.pop(sid)
            shutil.rmtree(ws.dir, ignore_errors=True)


registry = WorkspaceRegistry()