src/ai_agent/services/chat.py

"""Chat orchestration service.

This is the transport-agnostic version of the orchestration that used to live
inside ``ai_agent.ui.handlers.respond``. It owns:

  - validating/registering uploaded assets on the session
  - resolving asset previews into bytes for the VLM
  - resolving model/top_k/num_choices overrides from the UI
  - invoking the agent
  - turning the agent's structured output into a ``ChatTurnResult`` that
    contains both a markdown-friendly text and the raw recommendation list
  - handling pending tool approvals and demo confirmations

Phase-1 entrypoint is synchronous (``process_turn``). The FastAPI router in
phase 2 wraps this in an SSE stream, emitting recommendations, tool traces
and pending actions as discrete events.
"""

from __future__ import annotations

import logging
import os
import time
from dataclasses import dataclass, field
from datetime import datetime
from pathlib import Path
from typing import Any, Dict, List, Literal, Optional

from ai_agent.agent.agent import run_agent
from ai_agent.agent.tools.gradio_space_tool import RunExampleInput, tool_run_example
from ai_agent.agent.tools.mcp import (
    extract_downloads,
    extract_metadata,
    extract_output_field,
    extract_preview,
    get_tool,
)
from ai_agent.retriever.software_doc import SoftwareDoc
from ai_agent.utils.tags import parse_exclusions, strip_tags
from ai_agent.utils.utils import _is_affirmative

from .files import asset_paths, ingest_files
from .sessions import Session

log = logging.getLogger("services.chat")


# ---------------------------------------------------------------------------
# Result shapes
# ---------------------------------------------------------------------------
TurnStatus = Literal[
    "ok", "needs_clarification", "no_results", "error", "pending_action", "tool_executed"
]


@dataclass
class Recommendation:
    rank: int
    name: str
    accuracy: float
    why: str
    doc: Optional[Dict[str, Any]] = None
    demo_url: Optional[str] = None


@dataclass
class PendingAction:
    """A turn that ends asking the user to confirm something.

    The client surfaces an Approve / Decline control; calling
    ``approve_pending`` or ``decline_pending`` on the service resumes the
    flow.
    """

    type: Literal["demo_confirm", "tool_approval"]
    tool_name: str
    display_name: Optional[str] = None
    icon: Optional[str] = None
    image_name: Optional[str] = None
    demo_url: Optional[str] = None
    prompt: str = ""


@dataclass
class Clarification:
    question: str
    context: Optional[str] = None
    options: List[str] = field(default_factory=list)


@dataclass
class ChatTurnResult:
    status: TurnStatus
    text: str = ""
    recommendations: List[Recommendation] = field(default_factory=list)
    tool_traces: List[Dict[str, Any]] = field(default_factory=list)
    pending_action: Optional[PendingAction] = None
    clarification: Optional[Clarification] = None
    usage: Optional[Dict[str, Any]] = None
    error: Optional[str] = None
    # Extras used by Gradio rendering (preview images, downloads from a
    # tool execution turn). Empty for normal chat turns.
    images: List[str] = field(default_factory=list)
    files: List[tuple] = field(default_factory=list)


# ---------------------------------------------------------------------------
# Request shape
# ---------------------------------------------------------------------------
@dataclass
class ChatRequest:
    message: str = ""
    asset_ids: List[str] = field(default_factory=list)
    """Asset ids previously registered via ``services.files.ingest_files``."""

    new_file_paths: List[str] = field(default_factory=list)
    """Convenience for the Gradio adapter — files uploaded as part of this
    turn that haven't been pre-ingested. The chat service will ingest them
    and merge their asset_ids into ``asset_ids``."""

    model: Optional[str] = None
    """Display name from ``config.yaml#available_models`` (e.g.
    ``"openai/gpt-oss-120b [EPFL]"``). Resolved against the config map."""

    top_k: Optional[int] = None
    num_choices: Optional[int] = None


# ---------------------------------------------------------------------------
# Public entrypoints
# ---------------------------------------------------------------------------
def process_turn(
    session: Session,
    request: ChatRequest,
    doc_index: Dict[str, SoftwareDoc],
) -> ChatTurnResult:
    """Run one user turn end-to-end against ``session``.

    Mutates the session in place (history, last_asset_ids, banlist,
    pending_*, tool_calls). Returns a structured result the transport layer
    can render or stream.
    """
    session.touch()

    # 1) Ingest any newly-attached files (Gradio gives us raw paths)
    if request.new_file_paths:
        result = ingest_files(session, request.new_file_paths)
        if result.validation_errors:
            issues = "\n".join(f"• {x}" for x in result.validation_errors)
            text = f"⚠️ File validation issues:\n\n{issues}"
            session.conversation_history.append(f"Assistant: {text}")
            return ChatTurnResult(status="error", text=text, error="invalid_files")
        # Merge ingested asset_ids with any client-supplied ones
        new_ids = [a.asset_id for a in result.assets]
        request.asset_ids = list(dict.fromkeys(request.asset_ids + new_ids))

    # 2) Reject empty turns (no message AND no attachments)
    has_text = bool((request.message or "").strip())
    has_attachments = bool(request.asset_ids or session.last_asset_ids)
    if not has_text and not has_attachments:
        text = "Please provide a message or upload files."
        return ChatTurnResult(status="error", text=text, error="empty_input")

    # 3) Parse banlist tags out of the message
    clean_message = strip_tags(request.message or "")
    session.banlist |= set(parse_exclusions(request.message or ""))
    session.conversation_history.append(f"User: {clean_message}")

    # 4) Demo confirmation short-circuit
    if session.pending_demo_tool and _is_affirmative(request.message):
        return _execute_pending_demo(session, request.asset_ids)

    if session.pending_demo_tool:
        # Anything that isn't affirmative cancels the pending demo
        session.pending_demo_tool = None
        session.pending_demo_url = None

    # 5) Resolve attachment paths (default to last upload if none provided)
    if request.asset_ids:
        # Anything explicitly attached becomes the "active" set
        effective_paths, attached_assets = asset_paths(session, request.asset_ids)
        session.last_asset_ids = [a.asset_id for a in attached_assets]
    else:
        effective_paths = session.last_asset_paths()
        attached_assets = [session.assets[a] for a in session.last_asset_ids if a in session.assets]

    # Images are optional. When the user hasn't uploaded anything we run the
    # agent in text-only mode — retrieval still works on the prompt.

    # 6) Find the latest preview asset (used as the VLM image)
    preview_asset = session.last_preview()
    image_bytes: Optional[bytes] = None
    if preview_asset and preview_asset.preview_path:
        try:
            preview_path = Path(preview_asset.preview_path)
            if preview_path.exists():
                image_bytes = preview_path.read_bytes()
        except Exception as e:
            log.warning("Failed to read preview bytes: %r", e)

    image_metadata = preview_asset.metadata_text if preview_asset else None

    # 7) Resolve model config from the display name (if provided by the UI)
    model_name, base_url_override, api_key_env = _resolve_model_choice(request.model)

    # 8) Run the agent
    log.info(
        "Running agent: task=%r, attachments=%d, excluded=%d, model=%s",
        clean_message,
        len(effective_paths),
        len(session.banlist),
        request.model,
    )

    try:
        agent_result = run_agent(
            clean_message,
            image_paths=effective_paths,
            image_bytes=image_bytes,
            excluded=list(session.banlist),
            conversation_history=session.conversation_history,
            model=model_name,
            base_url=base_url_override if request.model else None,
            api_key_env=api_key_env,
            top_k=request.top_k,
            num_choices=request.num_choices,
            image_metadata=image_metadata,
        )
    except ValueError as e:
        return _format_config_error(e, session)
    except Exception as e:
        return _format_runtime_error(e, session)

    return _shape_agent_result(session, agent_result, doc_index, effective_paths)


def approve_pending(session: Session) -> ChatTurnResult:
    """Resume a turn that ended with a ``tool_approval`` pending action.

    Calls the registered tool with the previously-captured parameters, then
    clears the pending state.
    """
    tool_name = session.pending_tool_approval
    if not tool_name:
        return ChatTurnResult(
            status="error",
            text="There is no pending tool approval to confirm.",
            error="no_pending_action",
        )
    params = dict(session.pending_tool_params)
    return _execute_registered_tool(session, tool_name, params)


def decline_pending(session: Session) -> ChatTurnResult:
    """Decline both pending demo and pending tool approval."""
    session.pending_demo_tool = None
    session.pending_demo_url = None
    session.pending_tool_approval = None
    session.pending_tool_params = {}
    text = "👍 Got it — I won't run that. Tell me what to try instead."
    session.conversation_history.append(f"Assistant: {text}")
    return ChatTurnResult(status="ok", text=text)


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------
def _resolve_model_choice(display_name: Optional[str]):
    if not display_name:
        return None, None, None
    # Lazy import to avoid the Gradio dependency footprint when running from
    # the FastAPI backend.
    try:
        from ai_agent.ui.components import get_model_config

        cfg = get_model_config(display_name)
    except Exception as e:
        log.warning("Could not resolve model %r: %r", display_name, e)
        return None, None, None
    return (
        cfg.get("name"),
        cfg.get("base_url"),
        cfg.get("api_key_env", "OPENAI_API_KEY"),
    )


def _format_config_error(exc: ValueError, session: Session) -> ChatTurnResult:
    msg = str(exc)
    log.error("Configuration error: %s", msg)
    text = f"⚠️ **Configuration Error**\n\n{msg}\n\n"
    if "EPFL_API_KEY" in msg:
        text += (
            "💡 **Tip:** EPFL models require VPN connection and `EPFL_API_KEY` in "
            "your `.env` file. Try selecting an OpenAI model instead."
        )
    elif "OPENAI_API_KEY" in msg:
        text += "💡 **Tip:** Set `OPENAI_API_KEY` in your `.env` file to use OpenAI models."
    session.conversation_history.append(f"Assistant: {text}")
    return ChatTurnResult(status="error", text=text, error=msg)


def _format_runtime_error(exc: Exception, session: Session) -> ChatTurnResult:
    msg = str(exc)
    log.error("Agent execution error: %s", msg, exc_info=True)
    text = f"❌ **Error**\n\n{msg}\n\n"
    if "key_model_access_denied" in msg or "key not allowed" in msg.lower():
        text += "💡 **Tip:** This API key doesn't have access to this model.\n\n"
    elif "ConnectError" in msg or "Connection" in msg:
        text += "💡 **Tip:** Connection failed. If using EPFL models, ensure you're on EPFL VPN."
    session.conversation_history.append(f"Assistant: {text}")
    return ChatTurnResult(status="error", text=text, error=msg)


def _shape_agent_result(
    session: Session,
    agent_result,
    doc_index: Dict[str, SoftwareDoc],
    effective_paths: List[str],
) -> ChatTurnResult:
    """Translate an ``AgentToolSelection`` into a ``ChatTurnResult``."""
    legacy = agent_result.to_legacy_dict()

    tool_traces = legacy.get("tool_calls", []) or []
    if tool_traces:
        session.tool_calls.extend(tool_traces)

    usage = legacy.get("usage")
    usage_payload = None
    if usage:
        usage_payload = {
            "total": usage.get("total_tokens", 0),
            "input": usage.get("input_tokens", 0),
            "output": usage.get("output_tokens", 0),
        }

    status = legacy["conversation"]["status"]
    if status == "needs_clarification":
        question = legacy["conversation"]["question"]
        context = legacy["conversation"].get("context")
        options = legacy["conversation"].get("options", []) or []

        text = f"ℹ️ **I need more information:**\n\n{question}\n\n"
        if options:
            text += "**Options:**\n" + "\n".join(f"- {o}" for o in options) + "\n\n"
        if context:
            text += f"_{context}_"
        session.conversation_history.append(f"Assistant: {text}")
        return ChatTurnResult(
            status="needs_clarification",
            text=text,
            tool_traces=tool_traces,
            usage=usage_payload,
            clarification=Clarification(
                question=question, context=context, options=options
            ),
        )

    choices = legacy.get("choices") or []
    if not choices:
        reason = legacy.get("reason") or ""
        explanation = legacy.get("explanation") or ""
        parts = ["❌ **No suitable tools found.**\n"]
        if reason:
            parts.append(f"**Reason:** `{reason}`\n")
        if explanation:
            parts.append(explanation)
        text = "\n".join(parts)
        session.conversation_history.append(f"Assistant: {text}")
        return ChatTurnResult(
            status="no_results",
            text=text,
            tool_traces=tool_traces,
            usage=usage_payload,
        )

    # Recommendations path
    session.last_choices = {c["name"]: c for c in choices}
    for c in choices:
        if c.get("name"):
            session.banlist.add(c["name"])

    recommendations: List[Recommendation] = []
    for i, c in enumerate(choices, 1):
        doc = doc_index.get(c["name"])
        recommendations.append(
            Recommendation(
                rank=i,
                name=c["name"],
                accuracy=float(c.get("accuracy", 0.0)),
                why=c.get("why", ""),
                doc=doc.model_dump(mode="python") if doc is not None else None,
                demo_url=c.get("demo_link"),
            )
        )

    top = choices[0]
    text_parts = [
        f"✅ **I recommend {top['name']}** ({top.get('accuracy', 0):.1f}% match)\n",
        f"_{top.get('why', '')}_\n",
    ]
    text = "\n".join(text_parts)

    # Decide whether the top tool implies a pending action (registry-driven)
    pending_action: Optional[PendingAction] = None
    top_name = top["name"]
    demo_url = top.get("demo_link") or ""
    tool_config = get_tool(top_name)

    if tool_config and tool_config.requires_approval:
        image_path = effective_paths[0] if effective_paths else None
        session.pending_tool_approval = tool_config.name
        session.pending_tool_params = {
            "image_path": image_path,
            "description": f"Recommended by agent: {top.get('why', '')}",
        }
        pending_action = PendingAction(
            type="tool_approval",
            tool_name=tool_config.name,
            display_name=tool_config.display_name,
            icon=tool_config.icon,
            image_name=os.path.basename(image_path) if image_path else None,
            demo_url=demo_url or None,
            prompt=f"Run {tool_config.display_name} on your image?",
        )
    elif demo_url:
        session.pending_demo_tool = top_name
        session.pending_demo_url = demo_url
        pending_action = PendingAction(
            type="demo_confirm",
            tool_name=top_name,
            demo_url=demo_url,
            prompt=f"Would you like me to run the demo for {top_name}?",
        )

    session.conversation_history.append(f"Assistant: {text}")
    return ChatTurnResult(
        status="pending_action" if pending_action else "ok",
        text=text,
        recommendations=recommendations,
        tool_traces=tool_traces,
        pending_action=pending_action,
        usage=usage_payload,
    )


def _execute_pending_demo(session: Session, attached_ids: List[str]) -> ChatTurnResult:
    """Generic-demo flow (no registry entry, just a runnable demo URL)."""
    tool_name = session.pending_demo_tool
    demo_url = session.pending_demo_url
    log.info("User confirmed demo run for %s", tool_name)

    candidate_paths: List[str]
    if attached_ids:
        candidate_paths, _ = asset_paths(session, attached_ids)
    else:
        candidate_paths = session.last_asset_paths()

    session.pending_demo_tool = None
    session.pending_demo_url = None

    if not candidate_paths:
        text = "⚠️ No files available. Please upload an image first."
        session.conversation_history.append(f"Assistant: {text}")
        return ChatTurnResult(status="error", text=text, error="no_attachments")

    # Prefer TIFF if any
    pick = next(
        (
            p
            for p in candidate_paths
            if os.path.splitext(p)[1].lower() in (".tif", ".tiff")
        ),
        candidate_paths[0],
    )

    text = f"🚀 Running demo for **{tool_name}**...\n\n"
    images: List[str] = []
    files: List[tuple] = []
    try:
        demo_result = tool_run_example(
            RunExampleInput(
                tool_name=tool_name,
                image_path=pick,
                endpoint_url=demo_url or None,
            )
        )
        if demo_result.ran and (demo_result.result_preview or demo_result.result_image):
            preview_path = demo_result.result_preview or demo_result.result_image
            text += "✅ Demo completed!\n\n"
            images.append(preview_path)
            if demo_result.result_origin:
                files.append((demo_result.result_origin, "Download result"))
        else:
            note = demo_result.notes or "No output image returned"
            text += f"ℹ️ Demo ran but {note}"

        session.tool_calls.append(
            {
                "tool": "run_example",
                "tool_name": tool_name,
                "ran": demo_result.ran,
                "endpoint_url": demo_result.endpoint_url,
                "timestamp": datetime.now().isoformat(),
            }
        )
    except Exception as e:
        log.exception("Demo execution failed")
        text += f"❌ Error: {e}"

    session.conversation_history.append(f"Assistant: {text}")
    return ChatTurnResult(
        status="tool_executed", text=text, images=images, files=files
    )


def _execute_registered_tool(
    session: Session, tool_name: str, params: Dict[str, Any]
) -> ChatTurnResult:
    """Execute a registered tool that gated on user approval."""
    tool_config = get_tool(tool_name)
    if not tool_config:
        text = f"❌ Error: Unknown tool '{tool_name}'"
        session.pending_tool_approval = None
        session.pending_tool_params = {}
        session.conversation_history.append(f"Assistant: {text}")
        return ChatTurnResult(status="error", text=text, error="unknown_tool")

    started = time.time()
    text = f"{tool_config.icon} Running {tool_config.display_name}...\n\n"
    images: List[str] = []
    files: List[tuple] = []
    try:
        # Backfill missing image path from last upload
        if "image_path" in params and not params["image_path"]:
            paths = session.last_asset_paths()
            if paths:
                params["image_path"] = paths[0]

        input_obj = tool_config.input_model(**params)
        result = tool_config.executor(input_obj)

        success = extract_output_field(result, tool_config.success_field)
        error = extract_output_field(result, tool_config.error_field)
        compute_time_seconds = (
            extract_output_field(result, tool_config.compute_time_field) or 0.0
        )
        notes = extract_output_field(result, tool_config.notes_field)

        session.tool_calls.append(
            {
                "tool": tool_name,
                "success": success,
                "compute_time_seconds": compute_time_seconds,
                "error": error,
                "timestamp": datetime.now().isoformat(),
                **params,
            }
        )

        if success:
            text += f"✅ {tool_config.display_name} completed!\n\n"
            preview_path = extract_preview(result, tool_name)
            if preview_path and os.path.exists(preview_path):
                images.append(preview_path)
            for dp in extract_downloads(result, tool_name):
                if os.path.exists(dp):
                    files.append((dp, f"Download {tool_config.display_name} result"))
            metadata = extract_metadata(result, tool_name)
            if metadata:
                text += f"_{metadata}_\n\n"
            if notes:
                text += f"_{notes}_\n\n"
        else:
            text += f"❌ {tool_config.display_name} failed.\n\n"
            if error:
                text += f"**Error:** {error}\n\n"
    except Exception as e:
        log.exception("Tool %s execution failed", tool_name)
        text += f"❌ Error: {e}\n\n"

    session.pending_tool_approval = None
    session.pending_tool_params = {}
    elapsed = time.time() - started
    log.info("Tool %s finished in %.2fs", tool_name, elapsed)
    session.conversation_history.append(f"Assistant: {text}")
    return ChatTurnResult(
        status="tool_executed", text=text, images=images, files=files
    )


__all__ = [
    "ChatRequest",
    "ChatTurnResult",
    "Clarification",
    "PendingAction",
    "Recommendation",
    "approve_pending",
    "decline_pending",
    "process_turn",
]