mp1/pluto/server.py

# -*- coding: utf-8 -*-
"""
pluto/server.py — FastAPI server bridging pipeline <-> web UI.

Endpoints:
  POST /api/run      — start pipeline, return final JSON
  POST /api/upload   — upload files to the corpus
  GET  /api/corpus   — list corpus documents
  GET  /api/stream   — SSE stream of pipeline progress
  GET  /            — serve the frontend dashboard
"""

from __future__ import annotations

import asyncio
from functools import partial
import json
import os
import shutil
import tempfile
from uuid import uuid4
from pathlib import Path
from typing import Any

from fastapi.encoders import jsonable_encoder
from fastapi import FastAPI, File, Request, UploadFile
from fastapi.responses import HTMLResponse, JSONResponse, StreamingResponse
from fastapi.staticfiles import StaticFiles

from pluto.pipeline import PipelineRunner
from pluto.extraction_cache import ExtractionCache
from pluto.doc_index import DocIndex

app = FastAPI(title="Pluto Pipeline", version="1.0.0")

# ── State ─────────────────────────────────────────────────────────────────────

session_queues: dict[str, asyncio.Queue] = {}
session_results: dict[str, dict] = {}
session_cleanup_tasks: dict[str, asyncio.Task] = {}
SESSION_CLEANUP_DELAY_SECONDS = 300

FRONTEND_DIR = Path(__file__).parent.parent / "frontend"
CORPUS_DIR = Path(__file__).parent.parent / "corpus"
OUTPUT_DIR = Path(__file__).parent.parent / "output"

# Shared instances
_extraction_cache = ExtractionCache(str(CORPUS_DIR))
_doc_index = DocIndex(persist_path=CORPUS_DIR / ".doc_index.json")


def _docs_currently_understanding(doc_index: DocIndex) -> list[str]:
    """Return doc_ids still running background understanding."""
    return sorted(
        doc["doc_id"]
        for doc in doc_index.list_docs()
        if doc.get("processing_status") == "understanding" and not doc.get("is_processed")
    )


def _normalize_selected_doc_ids(raw_value: Any) -> list[str]:
    if not isinstance(raw_value, list):
        return []
    seen: set[str] = set()
    selected_doc_ids: list[str] = []
    for raw_doc_id in raw_value:
        doc_id = str(raw_doc_id or "").strip()
        if not doc_id or doc_id in seen:
            continue
        seen.add(doc_id)
        selected_doc_ids.append(doc_id)
    return selected_doc_ids


def _normalize_detail_level(raw_value: Any) -> str:
    return "detailed" if str(raw_value or "").strip().lower() == "detailed" else "standard"


def _processing_docs_for_scope(doc_index: DocIndex, selected_doc_ids: list[str] | None = None) -> list[str]:
    processing_docs = _docs_currently_understanding(doc_index)
    selected_doc_set = set(selected_doc_ids or [])
    if not selected_doc_set:
        return processing_docs
    return [doc_id for doc_id in processing_docs if doc_id in selected_doc_set]


def _json_safe(value: Any) -> Any:
    """Normalize Pydantic models and other rich objects into JSON-safe values."""
    return jsonable_encoder(value)


def _normalize_session_id(raw_value: Any) -> str:
    session_id = str(raw_value or "").strip()
    return session_id or str(uuid4())


def _get_session_queue(session_id: str) -> asyncio.Queue:
    cleanup_task = session_cleanup_tasks.pop(session_id, None)
    if cleanup_task:
        cleanup_task.cancel()

    queue = session_queues.get(session_id)
    if queue is None:
        queue = asyncio.Queue()
        session_queues[session_id] = queue
    return queue


def _schedule_session_cleanup(session_id: str, queue: asyncio.Queue) -> None:
    cleanup_task = session_cleanup_tasks.pop(session_id, None)
    if cleanup_task:
        cleanup_task.cancel()

    async def cleanup_later() -> None:
        try:
            await asyncio.sleep(SESSION_CLEANUP_DELAY_SECONDS)
            if session_queues.get(session_id) is queue:
                session_queues.pop(session_id, None)
                session_results.pop(session_id, None)
        except asyncio.CancelledError:
            pass
        finally:
            if session_cleanup_tasks.get(session_id) is task:
                session_cleanup_tasks.pop(session_id, None)

    task = asyncio.create_task(cleanup_later())
    session_cleanup_tasks[session_id] = task


def _session_doc_id(selected_doc_ids: list[str], result_data: dict | None = None) -> str:
    if selected_doc_ids:
        return selected_doc_ids[0]
    trace = (result_data or {}).get("trace_summary", {})
    docs_opened = trace.get("docs_opened", []) if isinstance(trace, dict) else []
    if docs_opened:
        return str(docs_opened[0])
    return "corpus"


def _schedule_session_compression(session_id: str) -> None:
    result_data = session_results.get(session_id)
    if not result_data:
        return

    doc_id = str(result_data.get("doc_id") or "corpus")

    async def compress_later() -> None:
        from pluto.session_memory import compress_session

        await asyncio.to_thread(compress_session, session_id, doc_id, result_data, CORPUS_DIR)

    asyncio.create_task(compress_later())


# ── Startup: re-index existing corpus files ─────────────────────────────────

@app.on_event("startup")
async def startup_reindex():
    """On server start, index any corpus files not already in DocIndex."""
    import logging
    from pluto.ingest import ingest_file, _split_into_chunks, _classify_and_tag_chunks
    from pluto.doc_index import ChunkMeta

    logger = logging.getLogger("pluto")
    CORPUS_DIR.mkdir(parents=True, exist_ok=True)

    for md_file in sorted(CORPUS_DIR.glob("*.md")):
        doc_id = md_file.stem
        if _doc_index.has_doc(doc_id):
            continue  # Already indexed (loaded from disk)

        logger.info(f"Re-indexing existing corpus file: {doc_id}")
        try:
            content = md_file.read_text(encoding="utf-8", errors="replace")
            chunks = _split_into_chunks(content)
            chunk_meta_list = _classify_and_tag_chunks(chunks)
            meta_objects = [
                ChunkMeta(
                    chunk_id=m["chunk_id"],
                    chunk_type=m["chunk_type"],
                    mode=m["mode"],
                    header=m["header"],
                )
                for m in chunk_meta_list
            ]
            _doc_index.register_doc(
                doc_id=doc_id,
                filename=md_file.name,
                chunks=chunks,
                chunk_meta=meta_objects,
            )
            _doc_index.set_overview(
                doc_id,
                "Preloaded corpus document re-indexed at startup; no generated overview is available yet.",
            )
        except Exception as e:
            logger.warning(f"Failed to re-index {doc_id}: {e}")

    logger.info(f"DocIndex ready: {len(_doc_index.list_docs())} documents indexed")


# ── Serve frontend ────────────────────────────────────────────────────────────

@app.get("/", response_class=HTMLResponse)
async def index():
    html_path = FRONTEND_DIR / "index.html"
    return html_path.read_text(encoding="utf-8")



# ── API routes ────────────────────────────────────────────────────────────────

@app.post("/api/run")
async def run_pipeline(request: Request):
    """Run the full pipeline for a user query."""
    body = await request.json()
    query = body.get("query", "")
    corpus_dir = body.get("corpus_dir", str(CORPUS_DIR))
    selected_doc_ids = _normalize_selected_doc_ids(body.get("selected_doc_ids"))
    detail_level = _normalize_detail_level(body.get("detail_level"))
    session_id = _normalize_session_id(body.get("session_id"))
    query_timestamp = body.get("query_timestamp")
    prev_query = body.get("prev_query", "")
    prev_query_timestamp = body.get("prev_query_timestamp")
    prev_session_id = str(body.get("prev_session_id") or "").strip()
    progress_queue = _get_session_queue(session_id)
    doc_id = _session_doc_id(selected_doc_ids)
    prior_session_context = []
    if selected_doc_ids:
        from pluto.session_memory import list_session_context
        prior_session_context = list_session_context(doc_id, CORPUS_DIR)

    if not query:
        return JSONResponse({"error": "No query provided", "session_id": session_id}, status_code=400)

    _capture_behavioral_signals(
        query=query,
        query_timestamp=query_timestamp,
        prev_query=prev_query,
        prev_query_timestamp=prev_query_timestamp,
        prev_session_id=prev_session_id,
        fallback_session_id=session_id,
    )

    processing_docs = _processing_docs_for_scope(_doc_index, selected_doc_ids)
    if processing_docs:
        return JSONResponse(
            {
                "error": "Please wait for document understanding to finish before running a query.",
                "processing_docs": processing_docs,
                "session_id": session_id,
            },
            status_code=409,
            headers={"Cache-Control": "no-store"},
        )

    # Reset queue for this run (drain any leftover events without replacing the object)
    while not progress_queue.empty():
        try:
            progress_queue.get_nowait()
        except asyncio.QueueEmpty:
            break

    def progress_callback(stage: str, data: dict):
        progress_queue.put_nowait(_json_safe({"stage": stage, **data}))

    # Run pipeline in a thread to avoid blocking
    loop = asyncio.get_event_loop()
    runner = PipelineRunner(
        corpus_dir=corpus_dir, output_dir=str(OUTPUT_DIR),
        doc_index=_doc_index,
        prior_session_context=prior_session_context,
    )
    runner.on_progress(progress_callback)

    try:
        result = await loop.run_in_executor(
            None,
            partial(
                runner.run,
                query,
                selected_doc_ids=selected_doc_ids,
                detail_level=detail_level,
            ),
        )
        session_results[session_id] = result.model_dump()

        # Include cache stats in the response
        cache_stats = runner.cache.stats()
        session_results[session_id]["cache_hits"] = cache_stats["hits"]
        session_results[session_id]["cache_misses"] = cache_stats["misses"]
        session_results[session_id]["session_id"] = session_id
        session_results[session_id]["query"] = query
        session_results[session_id]["doc_id"] = _session_doc_id(selected_doc_ids, session_results[session_id])

        # Signal completion
        await progress_queue.put({"stage": "done", "status": "complete", "session_id": session_id})

        return JSONResponse(session_results[session_id])

    except Exception as e:
        import traceback
        err_msg = str(e)
        traceback.print_exc()

        # Always signal error to SSE stream
        try:
            await progress_queue.put(
                {"stage": "error", "status": "failed", "detail": err_msg, "session_id": session_id}
            )
        except Exception:
            pass

        # ALWAYS return valid JSON — never let FastAPI return HTML 500
        return JSONResponse(
            {"error": f"Pipeline error: {err_msg}", "session_id": session_id},
            status_code=200  # Return 200 so browser can parse the JSON body
        )


@app.get("/api/stream")
async def stream_progress(session_id: str):
    """SSE stream of pipeline progress events."""
    progress_queue = _get_session_queue(session_id)

    async def event_generator():
        # Send one event immediately so EventSource opens before the POST
        # starts producing pipeline progress events.
        yield f"data: {json.dumps({'stage': 'connected', 'session_id': session_id})}\n\n"

        # Wait for events from the pipeline — keep connection open
        try:
            while True:
                try:
                    event = await asyncio.wait_for(progress_queue.get(), timeout=120.0)
                    yield f"data: {json.dumps(_json_safe(event))}\n\n"
                    if event.get("stage") in ("done", "error"):
                        if event.get("stage") == "done":
                            _schedule_session_compression(session_id)
                        break
                except asyncio.TimeoutError:
                    yield f"data: {json.dumps({'stage': 'heartbeat', 'session_id': session_id})}\n\n"
        finally:
            _schedule_session_cleanup(session_id, progress_queue)

    return StreamingResponse(
        event_generator(),
        media_type="text/event-stream",
        headers={
            "Cache-Control": "no-cache",
            "Connection": "keep-alive",
            "X-Accel-Buffering": "no",
        },
    )


@app.get("/api/result")
async def get_result(session_id: str):
    """Return the latest pipeline result for a session."""
    result = session_results.get(session_id)
    if result:
        return JSONResponse(result)
    return JSONResponse({"error": "No result yet", "session_id": session_id}, status_code=404)


@app.get("/api/session-context/{doc_id}")
async def get_session_context(doc_id: str):
    """Return recent compressed session context for a document."""
    from pluto.session_memory import list_session_context

    sessions = list_session_context(doc_id, CORPUS_DIR, limit=10)
    return JSONResponse({"doc_id": doc_id, "sessions": sessions}, headers={"Cache-Control": "no-store"})


@app.post("/api/compare")
async def benchmark_compare(request: Request):
    """Run benchmark: Pluto vs Single Model Baseline."""
    from benchmark.compare import ComparisonRunner
    
    body = await request.json()
    query = body.get("query", "")
    selected_doc_ids = _normalize_selected_doc_ids(body.get("selected_doc_ids"))
    detail_level = _normalize_detail_level(body.get("detail_level"))
    
    if not query:
        return JSONResponse({"error": "No query provided"}, status_code=400)

    processing_docs = _processing_docs_for_scope(_doc_index, selected_doc_ids)
    if processing_docs:
        return JSONResponse(
            {
                "error": "Please wait for document understanding to finish before running the benchmark.",
                "processing_docs": processing_docs,
            },
            status_code=409,
            headers={"Cache-Control": "no-store"},
        )

    try:
        runner = ComparisonRunner(str(CORPUS_DIR), doc_index=_doc_index)
        results = runner.compare(
            query,
            selected_doc_ids=selected_doc_ids,
            detail_level=detail_level,
        )
        return JSONResponse(results, headers={"Cache-Control": "no-store"})
    except Exception as e:
        return JSONResponse(
            {"error": f"Benchmark error: {e}"},
            status_code=200,
            headers={"Cache-Control": "no-store"},
        )


def _capture_behavioral_signals(
    query: str,
    query_timestamp: Any,
    prev_query: str,
    prev_query_timestamp: Any,
    prev_session_id: str,
    fallback_session_id: str,
) -> None:
    from pluto.signal_logger import check_prior_reference, check_rephrase, log_signal, query_hash

    referenced_session_id = prev_session_id or fallback_session_id

    if prev_query and prev_query_timestamp is not None and query_timestamp is not None:
        try:
            delta_seconds = (float(query_timestamp) - float(prev_query_timestamp)) / 1000.0
        except (TypeError, ValueError):
            delta_seconds = -1
        if check_rephrase(query, prev_query, delta_seconds):
            log_signal(referenced_session_id, query_hash(prev_query), "rephrase_fail")

    if check_prior_reference(query):
        log_signal(referenced_session_id, query_hash(query), "prior_reference")


# ── File upload ───────────────────────────────────────────────────────────────

ALLOWED_EXTENSIONS = {".pdf", ".docx", ".doc", ".txt", ".md", ".markdown"}


@app.post("/api/upload")
async def upload_files(files: list[UploadFile] = File(...)):
    """Upload one or more files to the corpus."""
    from pluto.ingest import ingest_file

    results = []
    errors = []

    for file in files:
        ext = Path(file.filename or "").suffix.lower()
        if ext not in ALLOWED_EXTENSIONS:
            errors.append({"filename": file.filename, "error": f"Unsupported type: {ext}"})
            continue

        # Save to temp, then ingest
        tmp_dir = tempfile.mkdtemp()
        try:
            tmp_path = Path(tmp_dir) / (file.filename or "upload")
            with open(tmp_path, "wb") as f:
                content = await file.read()
                f.write(content)

            info = ingest_file(tmp_path, str(CORPUS_DIR), doc_index=_doc_index)

            # ── Phase A: understand in BACKGROUND (don't block upload) ──
            doc_id = info["doc_id"]
            if not _doc_index.is_processed(doc_id):
                _doc_index.mark_processing(doc_id)
                import threading
                def _bg_understand(did):
                    try:
                        from pluto.stages.understand import run_understand
                        from pluto.tracer import Tracer
                        tracer = Tracer()
                        print(f"  [SERVER] Starting background Phase A for {did}...")
                        run_understand(did, _doc_index, tracer)
                        from pluto.doc_summary import generate_doc_summary
                        generate_doc_summary(did, CORPUS_DIR)
                        print(f"  [SERVER] Background Phase A COMPLETE for {did}")
                    except BaseException as e:
                        import traceback
                        print(f"  [CRITICAL] Background Phase A failed for {did}: {e}")
                        _doc_index.mark_failed(did, str(e))
                        traceback.print_exc()
                        # Ensure we don't leave the UI in a "loading" state if possible
                        # (though DocIndex handles state, a crash might bypass set_overview)
                threading.Thread(target=_bg_understand, args=(doc_id,), daemon=True).start()
                info["understanding"] = "in_progress"
            else:
                info["understanding"] = "complete"

            results.append(info)
        except Exception as e:
            errors.append({"filename": file.filename, "error": str(e)})
        finally:
            shutil.rmtree(tmp_dir, ignore_errors=True)

    return JSONResponse({
        "uploaded": results,
        "errors": errors,
        "corpus_size": len(list(CORPUS_DIR.glob("*.md"))),
    })


@app.get("/api/doc-status/{doc_id}")
async def doc_status(doc_id: str):
    """Check if a document has been fully understood (Phase A complete)."""
    if not _doc_index.has_doc(doc_id):
        return JSONResponse(
            {"doc_id": doc_id, "status": "not_found"},
            status_code=404,
            headers={"Cache-Control": "no-store"},
        )
    status = _doc_index.get_effective_status(doc_id)
    return JSONResponse({
        "doc_id": doc_id,
        "status": status,
        "has_overview": bool(_doc_index.get_overview(doc_id)),
        "chunk_count": _doc_index.get_chunk_count(doc_id),
        "error": _doc_index.get_last_error(doc_id),
    }, headers={"Cache-Control": "no-store"})


@app.get("/api/cache/stats")
async def cache_stats():
    """Return extraction cache statistics."""
    return JSONResponse(_extraction_cache.stats())


@app.get("/api/corpus")
async def list_corpus():
    """List all documents in the corpus."""
    CORPUS_DIR.mkdir(parents=True, exist_ok=True)
    docs = []
    for f in sorted(CORPUS_DIR.glob("*.md")):
        doc_id = f.stem
        has_doc = _doc_index.has_doc(doc_id)
        display_name = _doc_index.get_filename(doc_id) if has_doc else ""
        docs.append({
            "doc_id": doc_id,
            "filename": display_name or f.name,
            "stored_filename": f.name,
            "size": f.stat().st_size,
            "chunk_count": _doc_index.get_chunk_count(doc_id) if has_doc else 0,
            "processing_status": _doc_index.get_effective_status(doc_id) if has_doc else "not_found",
            "is_processed": _doc_index.is_processed(doc_id) if has_doc else False,
        })
    return JSONResponse({"documents": docs, "total": len(docs)}, headers={"Cache-Control": "no-store"})


@app.delete("/api/corpus/{doc_id}")
async def delete_corpus_doc(doc_id: str):
    """Delete a document from the corpus."""
    target = CORPUS_DIR / f"{doc_id}.md"
    if target.exists():
        target.unlink()
        # Remove from doc index
        _doc_index.remove_doc(doc_id)
        # Invalidate extraction cache for this doc
        removed = _extraction_cache.invalidate_doc(doc_id)
        _extraction_cache.save()
        return JSONResponse({"deleted": doc_id, "cache_entries_cleared": removed})
    return JSONResponse({"error": f"Document {doc_id} not found"}, status_code=404)


# ── Static file mount (AFTER all API routes to prevent shadowing) ─────────────
if FRONTEND_DIR.exists():
    app.mount("/static", StaticFiles(directory=str(FRONTEND_DIR)), name="static")

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)