src/core/command_router.py

from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path
from typing import Any

from scripts.orchestrator import PDFEditorOrchestrator


@dataclass(frozen=True)
class TeamMember:
    id: str
    name: str
    title: str
    agent_key: str
    icon: str
    color: str
    system_prompt: str


TEAM_MEMBERS: tuple[TeamMember, ...] = (
    TeamMember(
        id="lead",
        name="Mike",
        title="Command Lead",
        agent_key="layout_analysis",
        icon="🧭",
        color="#5b7cfa",
        system_prompt=(
            "You are a command lead for a PDF AI team. Break the user's goal into a "
            "clear execution direction for the rest of the team. Be concise."
        ),
    ),
    TeamMember(
        id="text",
        name="Iris",
        title="Text Extraction Specialist",
        agent_key="text_extraction",
        icon="📝",
        color="#2ec9b7",
        system_prompt=(
            "You decide how text should be extracted from PDFs. Focus on OCR, scanned "
            "documents, text quality, and extraction risks."
        ),
    ),
    TeamMember(
        id="tables",
        name="David",
        title="Table Analyst",
        agent_key="table_detection",
        icon="📊",
        color="#f59e0b",
        system_prompt=(
            "You analyze structured content in PDFs. Focus on tables, statements, rows, "
            "columns, and downstream data extraction needs."
        ),
    ),
    TeamMember(
        id="image",
        name="Alex",
        title="Image Quality Engineer",
        agent_key="image_processing",
        icon="🖼️",
        color="#a78bfa",
        system_prompt=(
            "You assess PDF page image quality. Focus on scan quality, blur, contrast, "
            "and whether OCR or cleanup is needed first."
        ),
    ),
    TeamMember(
        id="format",
        name="Emma",
        title="Formatting Editor",
        agent_key="format_correction",
        icon="✏️",
        color="#3ecf8e",
        system_prompt=(
            "You refine extracted PDF content. Focus on edits, formatting corrections, "
            "text insertion needs, cleanup, and output clarity."
        ),
    ),
    TeamMember(
        id="language",
        name="Sarah",
        title="Language Reviewer",
        agent_key="language_detection",
        icon="🌐",
        color="#ec4899",
        system_prompt=(
            "You detect PDF language and review multilingual requirements. Focus on OCR "
            "language choice and language-sensitive processing."
        ),
    ),
    TeamMember(
        id="quality",
        name="Bob",
        title="Quality Validator",
        agent_key="content_validation",
        icon="✅",
        color="#f87171",
        system_prompt=(
            "You validate whether the PDF workflow will produce reliable results. Focus "
            "on risks, verification, and failure prevention."
        ),
    ),
    TeamMember(
        id="export",
        name="Nora",
        title="Export Coordinator",
        agent_key="export_rendering",
        icon="📦",
        color="#06b6d4",
        system_prompt=(
            "You prepare the final delivery plan for a PDF workflow. Focus on output, "
            "artifacts, export expectations, and handoff."
        ),
    ),
)


class CommandRouter:
    """Atoms-like AI team planner for PDF workflows."""

    def __init__(self, project_root: Path | None = None, provider: str = "lmstudio"):
        self.project_root = Path(project_root) if project_root else Path.cwd()
        self.orchestrator = PDFEditorOrchestrator(provider=provider)

    def describe_team(self) -> list[dict[str, Any]]:
        mapping = self.orchestrator.describe_agent_models()
        members = []
        for member in TEAM_MEMBERS:
            llm = mapping.get(member.agent_key, {})
            members.append(
                {
                    "id": member.id,
                    "name": member.name,
                    "title": member.title,
                    "agent_key": member.agent_key,
                    "icon": member.icon,
                    "color": member.color,
                    "requested_provider": llm.get("requested_provider", ""),
                    "provider": llm.get("provider", ""),
                    "model": llm.get("model", ""),
                    "base_url": llm.get("base_url", ""),
                }
            )
        return members

    def plan_command(self, payload: dict[str, Any]) -> dict[str, Any]:
        brief = (payload.get("brief") or "").strip()
        if not brief:
            return {"success": False, "error": "Command brief is required."}

        selected_file = str(payload.get("input_pdf") or "").strip()
        output_dir = str(payload.get("output_dir") or "").strip()
        recommendation = self._build_recommendation(brief)
        team = self.describe_team()
        insights = [self._member_plan(member, brief, selected_file) for member in TEAM_MEMBERS]

        return {
            "success": True,
            "brief": brief,
            "input_pdf": selected_file,
            "output_dir": output_dir,
            "team": team,
            "insights": insights,
            "recommendation": recommendation,
            "available_pipelines": self.orchestrator.list_available_pipelines(),
        }

    def _member_plan(
        self, member: TeamMember, brief: str, selected_file: str
    ) -> dict[str, Any]:
        llm = self.orchestrator.agent_llms[member.agent_key]
        prompt = (
            f"User command:\n{brief}\n\n"
            f"Input PDF: {selected_file or 'not specified'}\n\n"
            "Respond in exactly 3 short bullet points:\n"
            "- Focus: your main responsibility\n"
            "- Action: what you would do next\n"
            "- Handoff: what the next role should receive"
        )
        raw = llm.complete(
            prompt,
            system=member.system_prompt,
            max_tokens=220,
            temperature=0.4,
        )
        focus, action, handoff = self._parse_member_response(raw, member)
        return {
            "member_id": member.id,
            "name": member.name,
            "title": member.title,
            "agent_key": member.agent_key,
            "focus": focus,
            "action": action,
            "handoff": handoff,
            "raw": raw,
        }

    def _parse_member_response(
        self, raw: str, member: TeamMember
    ) -> tuple[str, str, str]:
        if not raw or raw.startswith("Error:"):
            return self._fallback_member_response(member)

        lines = []
        for line in raw.splitlines():
            cleaned = line.strip().lstrip("-").strip()
            if cleaned:
                lines.append(cleaned)

        if len(lines) < 3:
            return self._fallback_member_response(member)

        return lines[0], lines[1], lines[2]

    def _fallback_member_response(self, member: TeamMember) -> tuple[str, str, str]:
        fallback = {
            "layout_analysis": (
                "Define the overall PDF processing direction.",
                "Decide whether the command needs OCR, editing, annotation, or validation first.",
                "Pass a prioritized flow to the rest of the team.",
            ),
            "text_extraction": (
                "Check whether text must be extracted from scanned pages.",
                "Recommend OCR when the command depends on searchable text.",
                "Pass extracted-text requirements to editing and validation.",
            ),
            "table_detection": (
                "Check whether the command targets statement or table structure.",
                "Highlight whether row/column preservation matters.",
                "Pass structured-data risks to the editor and validator.",
            ),
            "image_processing": (
                "Check scan quality and OCR readiness.",
                "Recommend cleanup when blur, noise, or low contrast will block OCR.",
                "Pass image-quality constraints to OCR and validation.",
            ),
            "format_correction": (
                "Plan text cleanup and edit operations.",
                "Translate the user command into concrete text or region edits.",
                "Pass cleaned content requirements to export.",
            ),
            "language_detection": (
                "Choose the right language handling path.",
                "Recommend OCR language and multilingual handling.",
                "Pass language settings to OCR and quality review.",
            ),
            "content_validation": (
                "Define what must be checked before delivery.",
                "Flag accuracy, truncation, and corruption risks.",
                "Pass final QA criteria to export.",
            ),
            "export_rendering": (
                "Decide final output expectations.",
                "Prepare final PDF, manifest, and delivery artifacts.",
                "Return the deliverables to the user.",
            ),
        }
        return fallback.get(
            member.agent_key,
            (
                "Own this part of the command.",
                "Produce the next useful step.",
                "Hand it to the next role.",
            ),
        )

    def _build_recommendation(self, brief: str) -> dict[str, Any]:
        text = brief.lower()
        wants_ocr = any(
            token in text
            for token in ["ocr", "scan", "scanned", "searchable", "extract text", "text layer"]
        )
        wants_edit = any(
            token in text
            for token in [
                "edit",
                "replace",
                "insert",
                "delete page",
                "reorder",
                "rotate",
                "fix text",
                "cleanup",
            ]
        )
        wants_annotation = any(
            token in text
            for token in ["highlight", "comment", "annotate", "redact", "mark", "rectangle"]
        )

        recommended_agents = ["inspector"]
        if wants_ocr:
            recommended_agents.append("ocr")
        if wants_edit:
            recommended_agents.append("editor")
        if wants_annotation:
            recommended_agents.append("annotator")
        recommended_agents.extend(["validator", "exporter"])

        deduped_agents = []
        for agent in recommended_agents:
            if agent not in deduped_agents:
                deduped_agents.append(agent)

        if wants_edit and wants_annotation:
            pipeline_name = "full_document"
        elif wants_edit:
            pipeline_name = "text_edit"
        elif wants_annotation:
            pipeline_name = "annotation_review"
        elif wants_ocr:
            pipeline_name = "ocr_preflight"
        else:
            pipeline_name = "full_document"

        reasons = []
        if wants_ocr:
            reasons.append("Command viitab OCR-ile või skännitud PDF-ile.")
        if wants_edit:
            reasons.append("Command sisaldab muutmist või sisu parandamist.")
        if wants_annotation:
            reasons.append("Command sisaldab märkimist, kommentaare või redaktsiooni.")
        if not reasons:
            reasons.append("Vaikimisi soovitan täielikku ülevaate + valideerimise töövoogu.")

        return {
            "pipeline_name": pipeline_name,
            "recommended_agents": deduped_agents,
            "reasons": reasons,
        }