app.py

from __future__ import annotations

import json
from typing import Any

import gradio as gr
import uvicorn

from agents.investigative_agent.orchestrator import run_investigation
from api import app as fastapi_app
from data.mock_alerts import list_alert_ids, get_alert_profile, ALERT_PROFILES


# ---------------------------------------------------------------------------
# Gradio wrapper
# ---------------------------------------------------------------------------

ALERT_CHOICES = [ALERT_PROFILES[aid]["label"] for aid in list_alert_ids()]


def _format_model_scores(model_scores: dict[str, Any]) -> str:
    """Render aggregated model scores as readable Markdown."""
    fraud = model_scores.get("fraud", {})
    kyc = model_scores.get("kyc", {})
    sanctions = model_scores.get("sanctions", {})

    fraud_score = fraud.get("fraud_score", "N/A")
    fraud_verdict = fraud.get("verdict", "N/A")
    top_features = fraud.get("top_features", [])

    kyc_score = kyc.get("anomaly_score", "N/A")
    kyc_verdict = kyc.get("verdict", "N/A")
    kyc_signals = kyc.get("flagged_signals", [])

    sanctions_found = sanctions.get("match_found", False)
    best_match = sanctions.get("best_match", {})
    sanctions_risk = best_match.get("risk_level", "CLEAR")
    matched_name = best_match.get("matched_name", "—")

    lines = [
        "### 📊 Transaction Fraud",
        f"- **Score:** `{fraud_score}`",
        f"- **Verdict:** **{fraud_verdict}**",
    ]
    if top_features:
        lines.append("- **Top Features:**")
        for feat in top_features[:5]:
            lines.append(f"  - `{feat.get('feature', '?')}` → SHAP: `{feat.get('shap_value', 0):.4f}`")

    lines.append("")
    lines.append("### 🆔 KYC Identity")
    lines.append(f"- **Anomaly Score:** `{kyc_score}`")
    lines.append(f"- **Verdict:** **{kyc_verdict}**")
    if kyc_signals:
        lines.append(f"- **Flagged Signals:** {', '.join(kyc_signals)}")

    lines.append("")
    lines.append("### 🌍 Sanctions & PEP")
    lines.append(f"- **Match Found:** `{sanctions_found}`")
    lines.append(f"- **Risk Level:** **{sanctions_risk}**")
    if sanctions_found:
        lines.append(f"- **Matched Entity:** {matched_name}")

    return "\n".join(lines)


def _format_typologies(analysis: dict[str, Any]) -> str:
    """Render typologies and confidence as Markdown."""
    confidence = analysis.get("confidence_score", 0)
    typologies = analysis.get("typologies", [])

    if confidence >= 70:
        conf_color = "🔴"
        conf_label = "HIGH"
    elif confidence >= 40:
        conf_color = "🟡"
        conf_label = "MEDIUM"
    else:
        conf_color = "🟢"
        conf_label = "LOW"

    lines = [
        f"### {conf_color} Confidence: **{confidence}/100** ({conf_label})",
        "",
        "### Identified Typologies",
    ]
    for typ in typologies:
        lines.append(f"- ⚠️ **{typ}**")
    if not typologies:
        lines.append("- ✅ No significant typologies identified")

    return "\n".join(lines)


def _run_investigation(alert_selection: str) -> tuple[Any, str, str, str, str, str]:
    """Wrapper called by the Gradio button."""
    alert_id = alert_selection.split(":")[0].strip()
    result = run_investigation(alert_id)

    profile = result.get("profile", {})
    model_scores = result.get("model_scores", {})
    analysis = result.get("analysis", {})

    # Build display-safe profile (exclude the label key)
    display_profile = {
        "transaction_fraud": profile.get("transaction_fraud", {}),
        "kyc_identity": profile.get("kyc_identity", {}),
        "sanctions_pep": profile.get("sanctions_pep", {}),
    }

    profile_json = json.dumps(display_profile, indent=2)
    scores_md = _format_model_scores(model_scores)
    typologies_md = _format_typologies(analysis)
    chain_of_thought = analysis.get("chain_of_thought", "")
    sar_draft = analysis.get("sar_draft", "")
    latency = f"⏱️ Investigation completed in **{result.get('latency_ms', 0)}ms**"

    return profile_json, scores_md, typologies_md, chain_of_thought, sar_draft, latency


# ---------------------------------------------------------------------------
# Custom CSS for premium look
# ---------------------------------------------------------------------------

CUSTOM_CSS = """
/* Global overrides for premium feel */
.gradio-container {
    max-width: 1400px !important;
    margin: 0 auto !important;
}

/* Header styling */
#dashboard-header {
    background: linear-gradient(135deg, #0f172a 0%, #1e293b 50%, #0f172a 100%);
    border-radius: 16px;
    padding: 28px 36px;
    margin-bottom: 16px;
    border: 1px solid rgba(99, 102, 241, 0.25);
    box-shadow: 0 4px 24px rgba(0, 0, 0, 0.15);
}
#dashboard-header h1 {
    color: #f8fafc !important;
    font-size: 1.6rem !important;
    margin: 0 0 4px 0 !important;
    letter-spacing: -0.02em;
}
#dashboard-header p {
    color: #94a3b8 !important;
    font-size: 0.95rem !important;
    margin: 0 !important;
}

/* Control row */
#control-row {
    background: rgba(30, 41, 59, 0.5);
    backdrop-filter: blur(12px);
    border-radius: 12px;
    padding: 16px 20px;
    margin-bottom: 16px;
    border: 1px solid rgba(99, 102, 241, 0.15);
}

/* Column panels */
.panel-left, .panel-right {
    background: rgba(15, 23, 42, 0.35);
    backdrop-filter: blur(8px);
    border-radius: 12px;
    padding: 20px;
    border: 1px solid rgba(99, 102, 241, 0.12);
    min-height: 600px;
}

/* Section labels */
.section-header {
    background: linear-gradient(90deg, rgba(99, 102, 241, 0.12), transparent);
    border-left: 3px solid #6366f1;
    padding: 8px 14px;
    border-radius: 0 8px 8px 0;
    margin-bottom: 12px;
    font-weight: 600;
    color: #c7d2fe;
}

/* SAR draft textbox */
#sar-draft textarea {
    border: 2px solid rgba(251, 191, 36, 0.3) !important;
    border-radius: 8px !important;
    background: rgba(251, 191, 36, 0.03) !important;
}
#sar-draft textarea:focus {
    border-color: rgba(251, 191, 36, 0.6) !important;
    box-shadow: 0 0 0 3px rgba(251, 191, 36, 0.1) !important;
}

/* Run button pulse */
#run-btn {
    background: linear-gradient(135deg, #6366f1, #8b5cf6) !important;
    border: none !important;
    font-weight: 600 !important;
    letter-spacing: 0.02em !important;
    transition: all 0.3s ease !important;
    box-shadow: 0 2px 12px rgba(99, 102, 241, 0.3) !important;
}
#run-btn:hover {
    transform: translateY(-1px) !important;
    box-shadow: 0 4px 20px rgba(99, 102, 241, 0.45) !important;
}

/* Latency bar */
#latency-bar {
    text-align: center;
    padding: 8px;
}
"""


# ---------------------------------------------------------------------------
# Gradio Blocks
# ---------------------------------------------------------------------------

with gr.Blocks(
    title="AML-intelligence-suite - Investigator Dashboard",
    theme=gr.themes.Soft(),
    css=CUSTOM_CSS,
) as gradio_app:

    # --- Header ---
    gr.HTML(
        """
        <div id="dashboard-header">
            <h1>🔍 AML-intelligence-suite</h1>
            <p>Investigator Dashboard — Agentic Workflow for AML Case Management</p>
        </div>
        """,
    )

    # --- Control row ---
    with gr.Row(elem_id="control-row"):
        alert_dropdown = gr.Dropdown(
            choices=ALERT_CHOICES,
            value=ALERT_CHOICES[0] if ALERT_CHOICES else None,
            label="Select Alert Case",
            interactive=True,
            scale=3,
        )
        run_btn = gr.Button(
            "🔍 Run Investigation",
            variant="primary",
            scale=1,
            elem_id="run-btn",
        )

    # --- Latency bar ---
    latency_display = gr.Markdown("", elem_id="latency-bar")

    # --- Main two-column layout ---
    with gr.Row(equal_height=False):

        # ===== LEFT COLUMN: Data Aggregation View =====
        with gr.Column(scale=1, elem_classes=["panel-left"]):
            gr.HTML('<div class="section-header">📁 Data Aggregation View — Source Truth</div>')

            gr.Markdown("#### Alert Profile (Raw Inputs)")
            profile_display = gr.Code(
                label="Alert Profile",
                language="json",
                interactive=False,
                lines=18,
            )

            gr.Markdown("#### Model Scores")
            scores_display = gr.Markdown(
                value="*Select an alert case and run an investigation to see model scores.*",
            )

        # ===== RIGHT COLUMN: Agentic Workflow View =====
        with gr.Column(scale=1, elem_classes=["panel-right"]):
            gr.HTML('<div class="section-header">🤖 Agentic Workflow View — AI Analysis</div>')

            gr.Markdown("#### Typologies & Confidence")
            typologies_display = gr.Markdown(
                value="*Awaiting investigation...*",
            )

            gr.Markdown("#### Chain of Thought")
            chain_display = gr.Textbox(
                label="Agent Reasoning",
                interactive=False,
                lines=12,
                max_lines=20,
            )

            gr.Markdown("#### 📝 SAR Draft — *Review and edit before final sign-off*")
            sar_display = gr.Textbox(
                label="SAR Draft (Editable)",
                interactive=True,
                lines=14,
                max_lines=25,
                elem_id="sar-draft",
            )

    # --- Wire button ---
    run_btn.click(
        fn=_run_investigation,
        inputs=[alert_dropdown],
        outputs=[
            profile_display,
            scores_display,
            typologies_display,
            chain_display,
            sar_display,
            latency_display,
        ],
    )


# ---------------------------------------------------------------------------
# Mount & serve
# ---------------------------------------------------------------------------

import uvicorn
from fastapi import Request
from starlette.responses import RedirectResponse

# Since Gradio swallows root mounts when mounted at "/", 
# we mount Gradio to "/ui" and redirect the root to "/ui".
# This leaves the FastAPI router (and our /mcp sub-mount) perfectly intact!
app = fastapi_app
app = gr.mount_gradio_app(app, gradio_app, path="/ui", root_path="/ui")

@app.get("/")
def root_redirect(request: Request):
    return RedirectResponse(url="/ui/")

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