"""
assistant.py — the capstone: a banking agentic assistant that COMPOSES the
earlier labs into one service (Knowledge Modules 00-09; the §3 reference stack).

It wires together, with no Azure required to run:
  - Lab 01 RAG (grounded, cited policy answers, security-trimmed)
  - Lab 02 LangGraph agent (balance via tools, transfer with human approval,
            refusal, durable state, least-privilege, audit)
  - Lab 03 IDP (document upload -> extract -> confidence gate -> validate)
  - Lab 06 guardrails (PII redaction + injection detection in; safety check out)

In production each lab's `Azure*` provider swaps in (Azure OpenAI, AI Search,
Document Intelligence, Content Safety) and this orchestration is unchanged.
"""
from __future__ import annotations

import os
import sys

# Make the sibling lab packages importable (every lab uses unique module names).
_LABS = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
for _d in ("lab-01-rag-azure-ai-search", "lab-02-langgraph-agent",
           "lab-03-document-intelligence", "lab-06-eval-guardrails"):
    sys.path.insert(0, os.path.join(_LABS, _d))

from data import DOCUMENTS                  # noqa: E402  (Lab 01 corpus)
from rag import RagPipeline                 # noqa: E402  (Lab 01)
from agent import BankingAgent, _intent     # noqa: E402  (Lab 02)
from agentgraph import Interrupt            # noqa: E402  (Lab 02)
from idp import IDPPipeline                 # noqa: E402  (Lab 03)
from guardrails import (detect_injection,   # noqa: E402  (Lab 06)
                        output_is_safe, redact_pii)


class Assistant:
    def __init__(self) -> None:
        self.rag = RagPipeline()
        self.rag.ingest(DOCUMENTS)
        self.agent = BankingAgent()          # carries the mock core-banking + tools
        self.idp = IDPPipeline()

    def chat(self, message: str, user: str, account_id: str, thread_id: str,
             user_groups: list[str] | None = None) -> dict:
        groups = user_groups or ["retail"]
        # --- INPUT GUARDRAILS (Lab 06) ---
        _, pii = redact_pii(message)
        injection = detect_injection(message)
        audit = [{"step": "input_guard", "pii": pii, "injection": injection}]

        intent = _intent(message)
        audit.append({"step": "route", "intent": intent})

        if intent == "policy":
            # --- grounded RAG (Lab 01) ---
            r = self.rag.ask(message, user_groups=groups)
            out = {"answer": r["answer"], "intent": intent,
                   "abstained": r["abstained"], "citations": r["citations"],
                   "pending": False, "source": "rag"}
        else:
            # --- agentic path (Lab 02): balance / transfer / refuse ---
            try:
                res = self.agent.handle(message, user, account_id, thread_id)
                out = {"answer": res["answer"], "intent": intent, "pending": False,
                       "tool_result": res.get("tool_result"),
                       "refused": res.get("refused", False), "source": "agent"}
                audit += res.get("audit", [])
            except Interrupt as it:
                out = {"answer": it.state["answer"], "intent": intent,
                       "pending": True, "thread_id": thread_id, "source": "agent"}
                audit += it.state.get("audit", [])

        # --- OUTPUT GUARDRAIL (Lab 06) ---
        need_cite = intent == "policy" and not out.get("abstained")
        safe = output_is_safe(out["answer"], require_citation=need_cite)
        out["output_safe"] = safe["ok"]
        if not safe["ok"]:                    # fail-closed: never return an unsafe answer
            out["answer"] = ("I'm unable to share that. Let me connect you with a "
                             "human agent.")
        out["audit"] = audit
        out["injection_flagged"] = injection
        return out

    def approve(self, thread_id: str, approved: bool, approver: str) -> dict:
        res = self.agent.approve(thread_id, approved, approver)
        return {"answer": res["answer"], "tool_result": res.get("tool_result"),
                "audit": res.get("audit", [])}

    def upload_document(self, raw: dict) -> dict:
        """Document Intelligence path (Lab 03): extract -> confidence gate ->
        validate -> straight-through or human review."""
        doc = self.idp.process(raw)
        return {"doc_type": doc.doc_type, "straight_through": doc.straight_through,
                "fields": doc.extracted, "review_fields": doc.review_fields,
                "validation": doc.validation}