src/chat_engine.py

"""OpenAI tool-calling loop with mandatory policy gating before MCP I/O."""

from __future__ import annotations

import json
import os
import re
import time
import uuid
from dataclasses import dataclass
from typing import Any

from openai import APIError, AsyncOpenAI

from src.logger import log_event, summarize_tool_args
from src.mcp_client import MCPAuthError, MCPNotFoundError, MCPTimeoutError, MCPTransportError, MeridianMCPClient
from src.policy import Allow, Deny, RequireConfirmation, authorize, inject_sensitive_tool_args, sync_pending_order_snapshot
from src.prompts import build_system_prompt
from src.schemas import LLM_TOOL_NAMES, TOOL_INPUT_MODELS
from src.session import Message, Session

_FALLBACK_MODEL = "gpt-4o-mini"
_MAX_TOOL_ITERATIONS = 5

MCP_TOOL_WHITELIST: frozenset[str] = frozenset(TOOL_INPUT_MODELS.keys())

_CONFIRM_RE = re.compile(
    r"^(yes|yep|yeah|confirm|confirmed|place\s+(?:the\s+)?order|go\s+ahead|"
    r"do\s+it|ok(?:ay)?|sounds\s+good|please\s+do)\b",
    re.I | re.S,
)
_REJECT_RE = re.compile(r"^(no|nope|nah|cancel|stop|wait|not\s+now)\b", re.I | re.S)

_SKU_RE = re.compile(r"\b([A-Z]{2,12}-\d{2,8})\b")


@dataclass
class ToolCallRecord:
    tool_name: str
    args_summary: str
    decision: str | None = None
    latency_ms: int | None = None
    status: str | None = None


@dataclass
class TurnResult:
    assistant_text: str
    tool_calls_made: list[ToolCallRecord]
    confirmation_requested: bool
    error: str | None = None


def apply_user_confirmation_intent(session: Session, user_message: str) -> None:
    """Deterministic confirmation handling — never delegate this to the LLM."""

    text = user_message.strip()
    if not text:
        return
    pending = session.pending_confirmation
    if not pending:
        return
    if _REJECT_RE.match(text):
        session.pending_confirmation = None
        return
    if _CONFIRM_RE.match(text):
        pending.confirmed = True


def _extract_price(text: str) -> str | None:
    m = re.search(r"\$?\s*(\d+(?:\.\d{1,2})?)", text)
    return m.group(1) if m else None


def update_facts_from_mcp_text(tool_name: str, text: str, session: Session) -> None:
    """Best-effort fact extraction — MCP returns prose, not JSON."""

    if tool_name in {"search_products", "list_products"}:
        results: list[dict[str, str | None]] = []
        for line in text.splitlines():
            sku_m = _SKU_RE.search(line)
            if not sku_m:
                continue
            results.append(
                {
                    "sku": sku_m.group(1),
                    "name": line.strip()[:160],
                    "price": _extract_price(line),
                }
            )
            if len(results) >= 5:
                break
        if results:
            session.facts["last_search_results"] = results
        return

    if tool_name == "get_product":
        sku_m = _SKU_RE.search(text)
        top = text.splitlines()[0][:220] if text else ""
        session.facts["last_viewed_product"] = {
            "sku": sku_m.group(1) if sku_m else None,
            "name": top,
            "price": _extract_price(text),
        }
        return

    if tool_name == "list_orders":
        orders: list[dict[str, str | None]] = []
        uuid_re = re.compile(
            r"[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}",
            re.I,
        )
        for line in text.splitlines():
            oid_m = uuid_re.search(line)
            if not oid_m:
                continue
            status: str | None = None
            for st in ("draft", "submitted", "approved", "fulfilled", "cancelled"):
                if st in line.lower():
                    status = st
                    break
            orders.append(
                {
                    "order_id": oid_m.group(0),
                    "status": status,
                    "total": _extract_price(line),
                }
            )
            if len(orders) >= 5:
                break
        if orders:
            session.facts["last_orders"] = orders
        return

    if tool_name == "create_order":
        session.pending_confirmation = None
        uuid_re = re.compile(
            r"[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}",
            re.I,
        )
        oid_m = uuid_re.search(text)
        entry = {"order_id": oid_m.group(0) if oid_m else None, "snippet": text[:400]}
        recent = session.facts.get("recent_orders")
        buf: list[Any] = recent if isinstance(recent, list) else []
        session.facts["recent_orders"] = [entry] + buf[:4]


def _mcp_tool_to_openai(tool: Any) -> dict[str, Any]:
    return {
        "type": "function",
        "function": {
            "name": tool.name,
            "description": (tool.description or "")[:4096],
            "parameters": tool.inputSchema,
        },
    }


def _history_to_messages(session: Session) -> list[dict[str, Any]]:
    """Prior turns only — the active user utterance is appended by handle_turn."""

    msgs: list[dict[str, Any]] = []
    hist = list(session.history)
    if len(hist) < 2:
        return msgs
    for m in hist[:-1]:
        if m.role not in {"user", "assistant"}:
            continue
        msgs.append({"role": m.role, "content": m.content})
    return msgs


class ChatEngine:
    def __init__(
        self,
        mcp_client: MeridianMCPClient,
        session: Session,
        openai_client: AsyncOpenAI,
        *,
        session_id: str | None = None,
        model: str | None = None,
    ) -> None:
        self._mcp = mcp_client
        self._session = session
        self._openai = openai_client
        self._session_id = session_id or str(uuid.uuid4())
        self._model = model or os.environ.get("OPENAI_MODEL", "").strip() or _FALLBACK_MODEL
        self._tools_cache: list[dict[str, Any]] | None = None

    @property
    def session_id(self) -> str:
        return self._session_id

    async def _openai_tools(self) -> list[dict[str, Any]]:
        if self._tools_cache is None:
            remote = await self._mcp.list_tools()
            self._tools_cache = [_mcp_tool_to_openai(t) for t in remote if t.name in LLM_TOOL_NAMES]
        return self._tools_cache

    async def handle_turn(self, user_message: str) -> TurnResult:
        apply_user_confirmation_intent(self._session, user_message)

        self._session.add_message(Message(role="user", content=user_message))

        hist = list(self._session.history)
        if not hist or hist[-1].role != "user":
            return TurnResult(
                "I'm missing internal chat state. Please restart the session.",
                [],
                False,
                error="history_invariant_broken",
            )
        prior_user_text = hist[-1].content
        api_messages: list[dict[str, Any]] = _history_to_messages(self._session) + [
            {"role": "user", "content": prior_user_text}
        ]

        tool_records: list[ToolCallRecord] = []
        confirmation_requested = False
        tools = await self._openai_tools()

        for iteration in range(_MAX_TOOL_ITERATIONS):
            system_prompt = build_system_prompt(self._session)
            t_llm = time.perf_counter()
            try:
                completion = await self._openai.chat.completions.create(
                    model=self._model,
                    max_tokens=1024,
                    messages=[{"role": "system", "content": system_prompt}, *api_messages],
                    tools=tools,
                    tool_choice="auto",
                )
            except APIError as e:
                log_event(
                    "error",
                    session_id=self._session_id,
                    tool=None,
                    decision=None,
                    deny_reason=str(e),
                    latency_ms=int((time.perf_counter() - t_llm) * 1000),
                    status="error",
                    args_summary=None,
                )
                msg = "I'm experiencing a technical issue. Please try again."
                self._session.add_message(Message(role="assistant", content=msg))
                return TurnResult(msg, tool_records, confirmation_requested, error=str(e))

            choice = completion.choices[0]
            finish = choice.finish_reason
            msg = choice.message

            log_event(
                "llm_call",
                session_id=self._session_id,
                tool=None,
                decision=None,
                deny_reason=None,
                latency_ms=int((time.perf_counter() - t_llm) * 1000),
                status="success",
                args_summary=f"finish_reason={finish};iteration={iteration + 1}",
            )

            tool_calls = msg.tool_calls or []
            if finish != "tool_calls" or not tool_calls:
                text = (msg.content or "").strip() or (
                    "I'm not sure how to help with that — could you rephrase?"
                )
                self._session.add_message(Message(role="assistant", content=text))
                return TurnResult(text, tool_records, confirmation_requested, None)

            assistant_msg: dict[str, Any] = {
                "role": "assistant",
                "content": msg.content or "",
                "tool_calls": [
                    {
                        "id": tc.id,
                        "type": "function",
                        "function": {
                            "name": tc.function.name,
                            "arguments": tc.function.arguments or "{}",
                        },
                    }
                    for tc in tool_calls
                ],
            }
            api_messages.append(assistant_msg)

            for tc in tool_calls:
                tool_name = tc.function.name
                tool_id = tc.id
                raw_args = tc.function.arguments or "{}"
                try:
                    parsed = json.loads(raw_args) if raw_args.strip() else {}
                except json.JSONDecodeError:
                    parsed = {}
                if not isinstance(parsed, dict):
                    summary = summarize_tool_args(tool_name, {})
                    tool_records.append(
                        ToolCallRecord(
                            tool_name,
                            summary,
                            decision=None,
                            latency_ms=None,
                            status="bad_args",
                        )
                    )
                    log_event(
                        "error",
                        session_id=self._session_id,
                        tool=tool_name,
                        decision=None,
                        deny_reason="invalid_tool_arguments_json",
                        latency_ms=None,
                        status="error",
                        args_summary=summary,
                    )
                    api_messages.append(
                        {
                            "role": "tool",
                            "tool_call_id": tool_id,
                            "content": "Tool arguments must be a JSON object; please fix and retry.",
                        }
                    )
                    continue

                args = dict(parsed)
                summary = summarize_tool_args(tool_name, args)
                t0 = time.perf_counter()

                if tool_name not in MCP_TOOL_WHITELIST:
                    lat = int((time.perf_counter() - t0) * 1000)
                    tool_records.append(
                        ToolCallRecord(tool_name, summary, decision="deny", latency_ms=lat, status="blocked")
                    )
                    log_event(
                        "policy_decision",
                        session_id=self._session_id,
                        tool=tool_name,
                        decision="deny",
                        deny_reason="Unknown tool",
                        latency_ms=lat,
                        status=None,
                        args_summary=summary,
                    )
                    api_messages.append(
                        {
                            "role": "tool",
                            "tool_call_id": tool_id,
                            "content": "Unknown tool — pick from the registered Meridian MCP tools only.",
                        }
                    )
                    continue

                decision = authorize(tool_name, args, self._session)
                if isinstance(decision, Deny):
                    lat = int((time.perf_counter() - t0) * 1000)
                    tool_records.append(
                        ToolCallRecord(
                            tool_name,
                            summary,
                            decision="deny",
                            latency_ms=lat,
                            status="denied",
                        )
                    )
                    log_event(
                        "policy_decision",
                        session_id=self._session_id,
                        tool=tool_name,
                        decision="deny",
                        deny_reason=decision.reason,
                        latency_ms=lat,
                        status=None,
                        args_summary=summary,
                    )
                    api_messages.append(
                        {"role": "tool", "tool_call_id": tool_id, "content": decision.reason}
                    )
                    continue

                if isinstance(decision, RequireConfirmation):
                    lat = int((time.perf_counter() - t0) * 1000)
                    confirmation_requested = True
                    tool_records.append(
                        ToolCallRecord(
                            tool_name,
                            summary,
                            decision="require_confirmation",
                            latency_ms=lat,
                            status="pending",
                        )
                    )
                    log_event(
                        "policy_decision",
                        session_id=self._session_id,
                        tool=tool_name,
                        decision="require_confirmation",
                        deny_reason=None,
                        latency_ms=lat,
                        status=None,
                        args_summary=summary,
                    )
                    if tool_name == "create_order" and self._session.pending_confirmation is None:
                        canonical = inject_sensitive_tool_args(
                            tool_name,
                            args,
                            self._session,
                        )
                        sync_pending_order_snapshot(self._session, canonical)
                    api_messages.append(
                        {"role": "tool", "tool_call_id": tool_id, "content": decision.prompt}
                    )
                    continue

                merged = inject_sensitive_tool_args(tool_name, args, self._session)
                log_event(
                    "policy_decision",
                    session_id=self._session_id,
                    tool=tool_name,
                    decision="allow",
                    deny_reason=None,
                    latency_ms=int((time.perf_counter() - t0) * 1000),
                    status=None,
                    args_summary=summary,
                )

                mcp_t0 = time.perf_counter()
                try:
                    payload = await self._mcp.call_tool(tool_name, merged)
                    mcp_lat = int((time.perf_counter() - mcp_t0) * 1000)
                    update_facts_from_mcp_text(tool_name, payload, self._session)
                    tool_records.append(
                        ToolCallRecord(
                            tool_name,
                            summarize_tool_args(tool_name, merged),
                            decision="allow",
                            latency_ms=mcp_lat,
                            status="success",
                        )
                    )
                    log_event(
                        "tool_call",
                        session_id=self._session_id,
                        tool=tool_name,
                        decision="allow",
                        deny_reason=None,
                        latency_ms=mcp_lat,
                        status="success",
                        args_summary=summarize_tool_args(tool_name, merged),
                    )
                    api_messages.append(
                        {"role": "tool", "tool_call_id": tool_id, "content": payload}
                    )
                except (MCPTransportError, MCPAuthError, MCPNotFoundError, MCPTimeoutError) as e:
                    mcp_lat = int((time.perf_counter() - mcp_t0) * 1000)
                    tool_records.append(
                        ToolCallRecord(
                            tool_name,
                            summarize_tool_args(tool_name, merged),
                            decision="allow",
                            latency_ms=mcp_lat,
                            status="error",
                        )
                    )
                    log_event(
                        "error",
                        session_id=self._session_id,
                        tool=tool_name,
                        decision="allow",
                        deny_reason=str(e),
                        latency_ms=mcp_lat,
                        status="error",
                        args_summary=summarize_tool_args(tool_name, merged),
                    )
                    api_messages.append(
                        {
                            "role": "tool",
                            "tool_call_id": tool_id,
                            "content": "I couldn't reach that information right now. Please try again shortly.",
                        }
                    )

        fallback = "I'm having trouble completing that request. Could you rephrase?"
        self._session.add_message(Message(role="assistant", content=fallback))
        log_event(
            "error",
            session_id=self._session_id,
            tool=None,
            decision=None,
            deny_reason="tool_iteration_cap",
            latency_ms=None,
            status="error",
            args_summary=None,
        )
        return TurnResult(fallback, tool_records, confirmation_requested, error="tool_iteration_cap")