app.py

"""
RegTech BR — Comparison Space
==============================
Side-by-side comparison of:
  A) RAG + Claude Sonnet (existing pipeline)
  B) Fine-tuned Mixtral-8x7B + LoRA adapter (AutoScientist)
Requirements (requirements.txt):
  gradio>=4.44.0
  faiss-cpu
  sentence-transformers
  transformers>=4.40.0
  peft>=0.10.0
  bitsandbytes>=0.43.0
  torch>=2.1.0
  accelerate>=0.27.0
  numpy
  requests
HuggingFace Space setup:
  - Hardware: ZeroGPU (PRO account) or A10G
  - Secrets: ANTHROPIC_API_KEY
  - Files: chunks_meta.jsonl, embeddings.npy, faiss_index.bin (RAG index)
  - The LoRA adapter is loaded from HuggingFace Hub:
    Fernandosr85/regtech-br-legal-adapter
"""

import os
import sys
import json
import re
import types
import html
import warnings
import unicodedata
from pathlib import Path

# ---------------------------------------------------------------------
# Python 3.13 compatibility
# ---------------------------------------------------------------------
# Python 3.13 removed the stdlib audioop module. Gradio imports pydub,
# and pydub tries to import audioop/pyaudioop during startup. This shim
# must run BEFORE importing gradio.
try:
    import audioop as _audioop
except Exception:
    try:
        import pyaudioop as _audioop
    except Exception:
        _audioop = types.ModuleType("audioop")
        _audioop.error = Exception
        sys.modules.setdefault("audioop", _audioop)
        sys.modules.setdefault("pyaudioop", _audioop)
    else:
        sys.modules.setdefault("audioop", _audioop)
        sys.modules.setdefault("pyaudioop", _audioop)
else:
    sys.modules.setdefault("audioop", _audioop)
    sys.modules.setdefault("pyaudioop", _audioop)

import gradio, fastapi, starlette, jinja2
print("VERSIONS:",
      "gradio", gradio.__version__,
      "fastapi", fastapi.__version__,
      "starlette", starlette.__version__,
      "jinja2", jinja2.__version__,
      flush=True)

import numpy as np
import requests

# CUDA memory behavior for single-GPU Hugging Face Spaces.
# Set before importing torch so PyTorch reads the allocator configuration.
os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True,max_split_size_mb:64")

import torch

warnings.filterwarnings("ignore", category=DeprecationWarning)

# audioop/pyaudioop compatibility is handled before importing gradio.
import gradio as gr

# ── ZeroGPU decorator (no-op when not on HF Spaces ZeroGPU) ──────────
try:
    import spaces
    HAS_ZEROGPU = True
except ImportError:
    HAS_ZEROGPU = False
    class spaces:
        @staticmethod
        def GPU(fn=None, duration=120):
            if fn is not None:
                return fn
            def decorator(f):
                return f
            return decorator

# ── Paths ─────────────────────────────────────────────────────────────
INDEX_DIR = Path(".")
ADAPTER_REPO = "Fernandosr85/regtech-br-legal-adapter"
ADAPTER_FILENAME = "finetune-artifact-adapter_brazilian_crypto.tgz"
# Use a pre-quantized 4-bit Mixtral repo by default.
# Loading "mistralai/Mixtral-8x7B-Instruct-v0.1" and quantizing on the fly can exceed 24GB VRAM
# during the loading/staging phase. The adapter was trained against Mixtral-Instruct; this
# pre-quantized repo should preserve the same module structure for PEFT adapter loading.
BASE_MODEL_ID = os.environ.get(
    "BASE_MODEL_ID",
    "ybelkada/Mixtral-8x7B-Instruct-v0.1-bnb-4bit",
)
ORIGINAL_BASE_MODEL_ID = "mistralai/Mixtral-8x7B-Instruct-v0.1"
USE_PREQUANTIZED_BASE = os.environ.get("USE_PREQUANTIZED_BASE", "1").strip().lower() in {"1", "true", "yes", "on"}
def _read_int_env(name: str, default: int) -> int:
    """Read an integer environment variable with a safe fallback."""
    try:
        return int(os.environ.get(name, str(default)))
    except Exception:
        return default


# ZeroGPU quota is charged by the requested duration.
# Keep the default small so Model A can be tested freely and Model B does not
# request 180s by accident. After quota reset, set:
#   ZEROGPU_DURATION_SECONDS=120
#   ZEROGPU_MAX_DURATION_SECONDS=120
# in the Space variables if you want a longer Model B run.
_REQUESTED_GPU_DURATION_SECONDS = _read_int_env("ZEROGPU_DURATION_SECONDS", 20)
_MAX_GPU_DURATION_SECONDS = _read_int_env("ZEROGPU_MAX_DURATION_SECONDS", 20)
GPU_DURATION_SECONDS = max(10, min(_REQUESTED_GPU_DURATION_SECONDS, _MAX_GPU_DURATION_SECONDS))
FT_MAX_NEW_TOKENS = max(80, min(_read_int_env("FT_MAX_NEW_TOKENS", 700), 900))

def _read_float_env(name: str, default: float) -> float:
    """Read a float environment variable with a safe fallback."""
    try:
        return float(os.environ.get(name, str(default)))
    except Exception:
        return default

FT_TEMPERATURE = max(0.0, min(_read_float_env("FT_TEMPERATURE", 0.05), 1.0))

print(
    "ZERO GPU CONFIG:",
    f"requested_env={_REQUESTED_GPU_DURATION_SECONDS}s",
    f"max_env={_MAX_GPU_DURATION_SECONDS}s",
    f"effective_duration={GPU_DURATION_SECONDS}s",
    f"ft_max_new_tokens={FT_MAX_NEW_TOKENS}",
    f"ft_temperature={FT_TEMPERATURE}",
    flush=True,
)

# Model B is a large Mixtral-8x7B LoRA path. A very short ZeroGPU request
# avoids quota errors, but it is not enough to load/generate reliably.
# Therefore, the app skips Model B before calling ZeroGPU unless the effective
# requested window is large enough. This prevents repeated "GPU task aborted".
MODEL_B_MIN_DURATION_SECONDS = _read_int_env("MODEL_B_MIN_DURATION_SECONDS", 120)
MODEL_B_ENABLED = os.environ.get("ENABLE_MODEL_B_GPU", "1").strip().lower() in {"1", "true", "yes", "on"}

print(
    "MODEL B POLICY:",
    f"enabled={MODEL_B_ENABLED}",
    f"min_required_duration={MODEL_B_MIN_DURATION_SECONDS}s",
    f"will_call_gpu={MODEL_B_ENABLED and GPU_DURATION_SECONDS >= MODEL_B_MIN_DURATION_SECONDS}",
    flush=True,
)
print(
    "MODEL B BASE:",
    f"BASE_MODEL_ID={BASE_MODEL_ID}",
    f"USE_PREQUANTIZED_BASE={USE_PREQUANTIZED_BASE}",
    flush=True,
)

print("Loading RAG index...", flush=True)

required = [
    INDEX_DIR / "chunks_meta.jsonl",
    INDEX_DIR / "embeddings.npy",
    INDEX_DIR / "faiss_index.bin",
]
missing = [str(p) for p in required if not p.exists()]
if missing:
    raise FileNotFoundError("Missing index files:\n" + "\n".join(missing))

import faiss
from sentence_transformers import SentenceTransformer

CHUNKS: list[dict] = []
with open(INDEX_DIR / "chunks_meta.jsonl", encoding="utf-8") as f:
    for line in f:
        line = line.strip()
        if line:
            CHUNKS.append(json.loads(line))

EMBEDDINGS = np.load(INDEX_DIR / "embeddings.npy")
RAG_INDEX = faiss.read_index(str(INDEX_DIR / "faiss_index.bin"))
EMBED_MODEL = SentenceTransformer("paraphrase-multilingual-MiniLM-L12-v2")

print(f"RAG index ready — {len(CHUNKS)} chunks", flush=True)

# ── Lazy-load fine-tuned model (only when ZeroGPU is allocated) ───────
_ft_model = None
_ft_tokenizer = None
_ft_load_error: str | None = None


def _safe_extract_tar(tar, target_dir: Path) -> None:
    """Safely extract a tar archive, rejecting path traversal entries."""
    target_root = target_dir.resolve()
    for member in tar.getmembers():
        member_path = (target_root / member.name).resolve()
        if not str(member_path).startswith(str(target_root)):
            raise RuntimeError(f"Unsafe archive member path: {member.name}")
    tar.extractall(target_root)


def _extract_adapter_archive(archive_path: str | Path, extract_dir: Path) -> None:
    """Extract the LoRA artifact even when the extension is misleading.
    Adaption artifacts may be named .tgz but not actually be gzip-compressed.
    This helper tries tar auto-detection, zstd decompression, and system tar.
    """
    import shutil
    import subprocess
    import tarfile

    archive_path = Path(archive_path)
    if extract_dir.exists():
        shutil.rmtree(extract_dir)
    extract_dir.mkdir(parents=True, exist_ok=True)

    with open(archive_path, "rb") as f:
        magic = f.read(8)

    size_mb = archive_path.stat().st_size / (1024 * 1024)
    print(
        f"Adapter artifact: {archive_path} | size={size_mb:.2f} MB | magic={magic.hex()}",
        flush=True,
    )

    errors: list[str] = []

    # 1) Python tar auto-detection: gzip, bzip2, xz, or plain tar.
    try:
        with tarfile.open(archive_path, "r:*") as tar:
            _safe_extract_tar(tar, extract_dir)
        print("Adapter extracted with tarfile.open(..., 'r:*').", flush=True)
        return
    except Exception as exc:
        errors.append(f"tarfile r:* failed: {type(exc).__name__}: {exc}")

    # 2) zstd-compressed tar. Magic bytes: 28 b5 2f fd.
    if magic.startswith(b"\x28\xb5\x2f\xfd"):
        try:
            import zstandard as zstd

            tar_path = extract_dir.parent / "adapter_decompressed.tar"
            with open(archive_path, "rb") as fin, open(tar_path, "wb") as fout:
                dctx = zstd.ZstdDecompressor()
                dctx.copy_stream(fin, fout)

            with tarfile.open(tar_path, "r:") as tar:
                _safe_extract_tar(tar, extract_dir)

            print("Adapter extracted after zstd decompression.", flush=True)
            return
        except Exception as exc:
            errors.append(f"zstd decompress failed: {type(exc).__name__}: {exc}")

    # 3) GNU tar fallback. Some HF images include tar with auto-compression support.
    try:
        completed = subprocess.run(
            ["tar", "-xaf", str(archive_path), "-C", str(extract_dir)],
            check=True,
            capture_output=True,
            text=True,
        )
        if completed.stdout:
            print(completed.stdout, flush=True)
        if completed.stderr:
            print(completed.stderr, flush=True)
        print("Adapter extracted with system tar -xaf.", flush=True)
        return
    except Exception as exc:
        stderr = getattr(exc, "stderr", "")
        errors.append(f"system tar -xaf failed: {type(exc).__name__}: {exc} {stderr}")

    raise RuntimeError("Could not extract adapter artifact.\n" + "\n".join(errors))


def _find_peft_adapter_dir(extract_dir: Path) -> Path:
    """Find the directory that contains adapter_config.json for PEFT."""
    candidates = []
    for config_path in extract_dir.rglob("adapter_config.json"):
        parent = config_path.parent
        has_weights = any(
            (parent / name).exists()
            for name in [
                "adapter_model.safetensors",
                "adapter_model.bin",
                "pytorch_model.bin",
            ]
        )
        candidates.append((parent, has_weights))

    if candidates:
        # Prefer directories that contain both adapter_config.json and adapter weights.
        candidates.sort(key=lambda x: (not x[1], len(str(x[0]))))
        chosen = candidates[0][0]
        print(f"PEFT adapter directory detected: {chosen}", flush=True)
        print(
            "Adapter directory files: "
            + ", ".join(sorted(p.name for p in chosen.iterdir())[:20]),
            flush=True,
        )
        return chosen

    preview = []
    for p in list(extract_dir.rglob("*"))[:80]:
        preview.append(str(p.relative_to(extract_dir)))
    raise FileNotFoundError(
        "adapter_config.json was not found after extraction. "
        "Extracted tree preview:\n" + "\n".join(preview)
    )


def _load_finetuned():
    """Load Mixtral + LoRA adapter. Called once inside a ZeroGPU context."""
    global _ft_model, _ft_tokenizer, _ft_load_error
    if _ft_model is not None:
        return True
    if _ft_load_error is not None:
        return False

    try:
        from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
        from peft import PeftModel
        from huggingface_hub import hf_hub_download

        print("Downloading LoRA adapter from HuggingFace Hub...", flush=True)

        adapter_dir = Path("/tmp/regtech_br_adapter")
        adapter_dir.mkdir(parents=True, exist_ok=True)

        artifact_path = hf_hub_download(
            repo_id=ADAPTER_REPO,
            filename=ADAPTER_FILENAME,
            repo_type="model",
            local_dir=str(adapter_dir),
        )

        extract_dir = adapter_dir / "extracted"
        _extract_adapter_archive(artifact_path, extract_dir)
        peft_adapter_dir = _find_peft_adapter_dir(extract_dir)

        print(
            "Adapter extracted. Loading base model...",
            f"BASE_MODEL_ID={BASE_MODEL_ID}",
            f"USE_PREQUANTIZED_BASE={USE_PREQUANTIZED_BASE}",
            flush=True,
        )

        if torch.cuda.is_available():
            torch.cuda.empty_cache()
            gpu_name = torch.cuda.get_device_name(0)
            total_gb = torch.cuda.get_device_properties(0).total_memory / (1024 ** 3)
            print(f"CUDA device: {gpu_name} | VRAM={total_gb:.2f} GiB", flush=True)

        _ft_tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL_ID, use_fast=True)
        if _ft_tokenizer.pad_token_id is None:
            _ft_tokenizer.pad_token = _ft_tokenizer.eos_token

        if USE_PREQUANTIZED_BASE:
            # Important for 24GB GPUs such as L4:
            # do not pass a new BitsAndBytesConfig here. The model repo already
            # contains its quantization config. Forcing device_map={"": 0} avoids
            # Accelerate dispatching layers to CPU/disk, which triggers:
            # "Some modules are dispatched on the CPU or the disk..."
            base = AutoModelForCausalLM.from_pretrained(
                BASE_MODEL_ID,
                device_map={"": 0},
                torch_dtype=torch.float16,
                low_cpu_mem_usage=True,
            )
        else:
            # Fallback path: quantize original Mixtral on the fly.
            # This usually needs more than 24GB during loading and is not recommended
            # for L4/T4 Spaces. Keep it available only for larger GPUs.
            bnb_config = BitsAndBytesConfig(
                load_in_4bit=True,
                bnb_4bit_quant_type="nf4",
                bnb_4bit_compute_dtype=torch.float16,
                bnb_4bit_use_double_quant=True,
            )
            base = AutoModelForCausalLM.from_pretrained(
                ORIGINAL_BASE_MODEL_ID,
                quantization_config=bnb_config,
                device_map={"": 0},
                torch_dtype=torch.float16,
                low_cpu_mem_usage=True,
            )

        print("Loading LoRA adapter weights...", flush=True)
        _ft_model = PeftModel.from_pretrained(
            base,
            str(peft_adapter_dir),
            is_trainable=False,
        )
        _ft_model.eval()

        print("Fine-tuned model ready.", flush=True)
        return True

    except Exception as exc:
        _ft_load_error = f"{type(exc).__name__}: {exc}"
        print(f"Fine-tuned model load error: {_ft_load_error}", flush=True)
        return False


# ── Shared utilities ──────────────────────────────────────────────────

def normalize(text: str) -> str:
    text = unicodedata.normalize("NFD", text or "")
    text = "".join(c for c in text if unicodedata.category(c) != "Mn")
    return re.sub(r"\s+", " ", re.sub(r"[^a-z0-9]+", " ", text.lower())).strip()


def detect_language(text: str) -> str:
    pt_markers = ["nossa", "nosso", "nao", "para", "dos", "das", "uma", "com",
                  "sao", "que", "por", "mas", "como", "qual", "ativos", "clientes",
                  "empresa", "plataforma", "banco", "central", "criptoativos"]
    return "pt" if sum(1 for m in pt_markers if m in normalize(text)) >= 2 else "en"


def esc(v) -> str:
    return html.escape("" if v is None else str(v))


def as_list(value) -> list[str]:
    if value is None:
        return []
    if isinstance(value, list):
        out = []
        for item in value:
            if isinstance(item, dict):
                out.append("; ".join(f"{k}: {v}" for k, v in item.items() if v))
            else:
                s = str(item).strip()
                if s:
                    out.append(s)
        return list(dict.fromkeys(out))
    return [str(value).strip()] if str(value).strip() else []



def extract_json_object(raw: str) -> str:
    """Extract the first balanced JSON object from a model response.

    The fine-tuned Mixtral adapter may return:
      ```json
      {...}
      ```
    or a short preamble before the object. This function removes markdown
    fences anywhere near the beginning and then extracts the first balanced
    JSON object from the first "{". It does not require the response to start
    exactly with a brace.
    """
    text = (raw or "").strip()

    # Normalize curly quotes and common escaped fence artifacts.
    text = text.replace("“", '"').replace("”", '"').replace("’", "'")
    text = text.replace("\\`\\`\\`", "```")

    # Remove markdown code fences even if there is whitespace or a short preamble.
    text = re.sub(r"^\s*```(?:json|JSON)?\s*", "", text).strip()
    text = re.sub(r"\s*```\s*$", "", text).strip()

    # Remove common model preambles before JSON.
    text = re.sub(
        r"^\s*(here(?:'| i)?s the json|here is the json|response|json output|json)\s*[:\-]?\s*",
        "",
        text,
        flags=re.IGNORECASE,
    ).strip()

    # If a fence still appears before the first brace, drop everything up to it.
    fence_pos = text.find("```")
    brace_pos = text.find("{")
    if fence_pos != -1 and (brace_pos == -1 or fence_pos < brace_pos):
        text = text[fence_pos + 3:].strip()
        text = re.sub(r"^(?:json|JSON)\s*", "", text).strip()

    start = text.find("{")
    if start < 0:
        return text

    in_string = False
    escaped = False
    depth = 0

    for pos in range(start, len(text)):
        ch = text[pos]

        if escaped:
            escaped = False
            continue

        if ch == "\\":
            escaped = True
            continue

        if ch == '"':
            in_string = not in_string
            continue

        if in_string:
            continue

        if ch == "{":
            depth += 1
        elif ch == "}":
            depth -= 1
            if depth == 0:
                return text[start:pos + 1]

    # Truncated JSON: return from the first brace onward. The caller will
    # attempt repair/fallback and expose raw output in the panel.
    return text[start:]

def load_model_json(raw: str) -> dict:
    """Load model JSON with light repair for common generation artifacts."""
    import ast

    candidates = []
    clean = (raw or "").strip()
    extracted = extract_json_object(clean)
    candidates.extend([clean, extracted])

    repaired = extracted
    repaired = repaired.replace("“", '"').replace("”", '"').replace("’", "'")
    repaired = re.sub(r",\s*([}\]])", r"\1", repaired)
    candidates.append(repaired)

    for candidate in candidates:
        candidate = (candidate or "").strip()
        if not candidate:
            continue
        try:
            parsed = json.loads(candidate)
            if isinstance(parsed, dict):
                return parsed
        except Exception:
            pass

    # Python-dict-like fallback: {'risk_level': 'HIGH', ...}
    try:
        parsed = ast.literal_eval(repaired)
        if isinstance(parsed, dict):
            return parsed
    except Exception:
        pass

    raise json.JSONDecodeError("Could not parse a JSON object from model output", raw or "", 0)


def fallback_unstructured_adapter_report(raw: str) -> dict:
    """Return a valid report when the adapter output is not valid JSON.

    This prevents the Model B panel from collapsing into a generic
    "JSON parse failed" message and surfaces the raw text for evaluation.
    """
    excerpt = re.sub(r"\s+", " ", (raw or "").strip())[:900]
    if not excerpt:
        excerpt = "The fine-tuned adapter returned an empty response."

    return {
        "risk_level": "UNCLEAR",
        "compliance_status": "REQUIRES_REVIEW",
        "applicable_regulations": ["See retrieved RAG context"],
        "relevant_articles": ["See retrieved RAG context"],
        "finding": (
            "The fine-tuned adapter responded, but its output was incomplete or not valid JSON after markdown cleanup. "
            "Raw adapter output excerpt: " + excerpt
        ),
        "corrective_action": (
            "Review the raw adapter output and compare it with Model A. "
            "For production use, keep the JSON repair layer enabled or fine-tune with stricter JSON-only examples."
        ),
        "confidence": "LOW",
        "authority": "mixed",
        "_raw_adapter_output": raw or "",
        "_parse_warning": "Adapter output was not valid JSON; fallback report generated.",
    }


# ── RAG retrieval ─────────────────────────────────────────────────────

AUTHORITY_KW = {
    "BCB": [
        "banco central", "bcb", "bacen", "psav", "psaav",
        "prestadora de servicos de ativos virtuais", "prestador de servicos de ativos virtuais",
        "ativos virtuais", "ativo virtual", "criptoativos", "criptoativo",
        "autorizacao", "autorizar", "authorization", "authorisation", "licensed", "license",
        "segregacao", "segregacao patrimonial", "segregation", "segregate", "segregated",
        "asset segregation", "client assets", "customer assets", "customer funds", "own funds",
        "company funds", "corporate treasury", "commingle", "commingled", "custody", "custodial",
        "wallet", "proof of reserves", "reserves", "exchange", "brokerage", "broker",
        "circular", "resolucao bcb", "instrucao normativa bcb",
    ],
    "CVM": [
        "cvm", "comissao de valores", "valores mobiliarios", "valor mobiliario",
        "token", "tokens", "security token", "oferta publica", "public offering",
        "dividendos", "dividends", "receita", "revenue share", "direito de voto", "voting rights",
        "investidor", "investors", "captacao", "fundraising", "rwa",
    ],
    "COAF": [
        "coaf", "pep", "pessoa exposta politicamente", "politically exposed",
        "kyc", "conheca seu cliente", "identificacao", "customer identification",
        "anonimo", "anonymous", "aml", "cft", "pld", "ftp", "lavagem", "terrorismo",
    ],
}

SOURCE_KW = {
    "lei_14478": ["lei 14 478", "lei 14478", "marco legal", "crypto legal framework", "virtual assets law"],
    "decreto_11563": ["decreto 11 563", "decreto 11563", "banco central", "central bank"],
    "bcb_circular_3978": ["pld", "ftp", "lavagem", "kyc", "anonimo", "anonymous", "aml", "cft"],
    "bcb_in701": [
        "segregacao", "segregacao patrimonial", "segregation", "segregate", "segregated",
        "client assets", "customer assets", "customer funds", "own funds", "company funds",
        "custody", "custodial", "proof of reserves", "certificacao tecnica", "technical certification",
    ],
    "bcb_res548": [
        "autorizacao", "authorization", "authorisation", "psav", "vasp", "licensed", "license",
        "prestadora", "service provider", "central bank", "banco central",
    ],
    "cvm": ["cvm", "valores mobiliarios", "token", "tokens", "dividendos", "dividends", "voting", "public offering"],
    "coaf": ["coaf", "pep", "pessoa exposta politicamente", "politically exposed"],
}


def detect_route(query: str) -> dict:
    q = normalize(query)
    route = {"authority_filters": [], "source_id_contains": [], "query_expansion": []}

    hits = {a: sum(1 for k in kws if normalize(k) in q) for a, kws in AUTHORITY_KW.items()}
    max_hits = max(hits.values()) if hits else 0
    if max_hits:
        route["authority_filters"] = [a for a, h in hits.items() if h == max_hits or h >= 2]

    for key, kws in SOURCE_KW.items():
        if any(normalize(k) in q for k in kws):
            route["source_id_contains"].append(key)

    segregation_terms = [
        "segregacao", "segregacao patrimonial", "segregation", "segregate", "segregated",
        "client assets", "customer assets", "customer funds", "own funds", "company funds",
        "corporate treasury", "commingle", "commingled", "custody", "custodial",
        "proof of reserves", "wallet", "brokerage", "exchange",
    ]
    if any(normalize(t) in q for t in segregation_terms):
        route["authority_filters"].append("BCB")
        route["source_id_contains"].extend(["bcb_in701", "bcb_res548", "lei_14478", "decreto_11563"])
        route["query_expansion"].extend([
            "segregação patrimonial de ativos virtuais",
            "ativos de titularidade de clientes e usuários",
            "ativos de titularidade da instituição",
            "asset segregation for virtual asset service providers",
            "client assets separated from own funds",
            "custody controls and proof of reserves",
            "PSAV technical certification and custody requirements",
        ])

    authorization_terms = [
        "autorizacao", "authorization", "authorisation", "psav", "vasp",
        "banco central", "central bank", "licensed", "license", "without authorization",
    ]
    if any(normalize(t) in q for t in authorization_terms):
        route["authority_filters"].append("BCB")
        route["source_id_contains"].extend(["bcb_res548", "bcb_in701", "lei_14478", "decreto_11563"])
        route["query_expansion"].extend([
            "autorização para prestadora de serviços de ativos virtuais",
            "Banco Central do Brasil",
            "PSAV authorization requirements",
            "virtual asset service provider authorization",
        ])

    aml_terms = ["kyc", "anonimo", "anonymous", "identificacao", "customer identification", "lavagem", "aml", "cft", "pld", "ftp"]
    if any(normalize(t) in q for t in aml_terms):
        route["authority_filters"].extend(["BCB", "COAF"])
        route["source_id_contains"].extend(["bcb_circular_3978", "coaf"])
        route["query_expansion"].extend([
            "identificação e qualificação de clientes",
            "cliente anônimo",
            "prevenção à lavagem de dinheiro e financiamento do terrorismo",
            "AML CFT customer due diligence",
        ])

    for key in route:
        if isinstance(route[key], list):
            route[key] = list(dict.fromkeys(route[key]))
    return route


def _chunk_text_for_matching(chunk: dict) -> str:
    tags = chunk.get("tags", [])
    if not isinstance(tags, list):
        tags = [tags]
    return normalize(" ".join([
        str(chunk.get("source_id", "")),
        str(chunk.get("source_label", "")),
        str(chunk.get("authority", "")),
        str(chunk.get("article_hint", "")),
        str(chunk.get("normative_reference_hint", "")),
        " ".join(str(t) for t in tags),
        str(chunk.get("text", "")),
    ]))


def _semantic_score(raw_score: float) -> float:
    """Convert FAISS raw scores into a higher-is-better score.
    Some indexes are inner-product indexes and some are L2-distance indexes.
    The previous code assumed all scores were similarities and discarded
    scores below 0.20 before route boosts. That can return zero chunks for
    short English prompts such as "does not segregate client assets".
    """
    metric = getattr(RAG_INDEX, "metric_type", None)
    try:
        is_l2 = metric == faiss.METRIC_L2
    except Exception:
        is_l2 = False
    if is_l2:
        return 1.0 / (1.0 + max(float(raw_score), 0.0))
    return float(raw_score)


def _lexical_boost(query_norm: str, chunk_norm: str) -> float:
    terms = [t for t in query_norm.split() if len(t) >= 4]
    if not terms:
        return 0.0
    unique_terms = list(dict.fromkeys(terms))
    hits = sum(1 for t in unique_terms if t in chunk_norm)
    return min(0.18, hits / max(6, len(unique_terms)) * 0.18)


def _route_boost(chunk: dict, route: dict) -> float:
    sid = normalize(str(chunk.get("source_id", "")))
    authority = str(chunk.get("authority", ""))
    rb = 0.0
    for tok in route.get("source_id_contains", []):
        tok_norm = normalize(tok)
        if tok_norm and tok_norm in sid:
            rb += 0.45
    if authority in route.get("authority_filters", []):
        rb += 0.10
    return min(rb, 0.75)


def retrieve(query: str, top_k: int = 5) -> list[dict]:
    route = detect_route(query)
    expanded = query + "\n" + "\n".join(route.get("query_expansion", []))
    query_norm = normalize(expanded)

    print(f"RAG route: {route}", flush=True)
    print(f"RAG expanded query: {expanded[:500]}", flush=True)

    q_vec = EMBED_MODEL.encode(
        [expanded],
        normalize_embeddings=True,
        convert_to_numpy=True,
    ).astype(np.float32)

    # Search the entire small index. It has only 103 chunks, so this is cheap,
    # and it avoids losing good legal chunks due to a narrow semantic candidate set.
    k_search = len(CHUNKS)
    scores, indices = RAG_INDEX.search(q_vec, k_search)

    ranked: list[dict] = []

    for raw_score, idx in zip(scores[0], indices[0]):
        if idx < 0:
            continue
        chunk = CHUNKS[int(idx)].copy()
        chunk_norm = _chunk_text_for_matching(chunk)
        sem = _semantic_score(float(raw_score))
        lex = _lexical_boost(query_norm, chunk_norm)
        rb = _route_boost(chunk, route)
        chunk["_score"] = sem
        chunk["_final"] = sem + lex + rb
        chunk["_match_reason"] = "semantic+lexical+route"
        ranked.append(chunk)

    # Force route-matched legal sources into the candidate pool. This is the
    # safety net that prevents "No relevant regulatory chunks found" for short
    # but legally obvious prompts such as asset segregation / BCB authorization.
    route_tokens = [normalize(t) for t in route.get("source_id_contains", []) if normalize(t)]
    if route_tokens:
        for chunk0 in CHUNKS:
            sid = normalize(str(chunk0.get("source_id", "")))
            if not any(tok in sid for tok in route_tokens):
                continue
            chunk = chunk0.copy()
            chunk_norm = _chunk_text_for_matching(chunk)
            lex = _lexical_boost(query_norm, chunk_norm)
            rb = _route_boost(chunk, route)
            chunk["_score"] = 0.0
            chunk["_final"] = 0.85 + lex + rb
            chunk["_match_reason"] = "forced_route_source"
            ranked.append(chunk)

    ranked.sort(key=lambda r: float(r.get("_final", 0.0)), reverse=True)

    seen, unique = set(), []
    for r in ranked:
        raw_cid = r.get("chunk_id", r.get("source_id", ""))
        try:
            cid = json.dumps(raw_cid, sort_keys=True, ensure_ascii=False, default=str)
        except Exception:
            cid = str(raw_cid)
        if cid not in seen:
            seen.add(cid)
            unique.append(r)
        if len(unique) >= top_k:
            break

    print(f"RAG retrieved chunks after robust routing: {len(unique)}", flush=True)
    for i, r in enumerate(unique, 1):
        print(
            f"[RAG {i}] source_id={r.get('source_id')} | authority={r.get('authority')} | "
            f"article_hint={r.get('article_hint')} | final={float(r.get('_final', 0.0)):.3f} | "
            f"reason={r.get('_match_reason')}",
            flush=True,
        )

    return unique


def format_context(results: list[dict]) -> str:
    lines = []
    for i, r in enumerate(results, 1):
        art = f" — {r['article_hint']}" if r.get("article_hint") else ""
        norm = f" [{r['normative_reference_hint']}]" if r.get("normative_reference_hint") else ""
        lines.append(
            f"[SOURCE {i}] {r.get('source_label', '')}{art}{norm}\n"
            f"Source ID: {r.get('source_id', '?')} | Authority: {r.get('authority', '?')} | "
            f"Score: {float(r.get('_final', 0.0)):.3f} | Match: {r.get('_match_reason', '?')}\n"
            f"{str(r.get('text', ''))[:600]}..."
        )
    return "\n\n---\n\n".join(lines)


# ── Claude API (Model A) ──────────────────────────────────────────────

COMPLIANCE_SYSTEM = """You are RegTech BR, a specialist AI in Brazilian crypto asset regulation.
Analyze the compliance query and produce a structured JSON assessment.
Respond ONLY with valid JSON — no markdown fences.
Use EXACTLY these snake_case keys:
{
  "risk_level": "LOW | MEDIUM | HIGH | UNCLEAR",
  "compliance_status": "COMPLIANT | NON-COMPLIANT | REQUIRES_REVIEW | INSUFFICIENT_INFO",
  "applicable_regulations": ["list of regulation names"],
  "relevant_articles": ["list of article references"],
  "finding": "2-5 sentence assessment",
  "corrective_action": "specific steps or 'No action required'",
  "confidence": "HIGH | MEDIUM | LOW",
  "authority": "BCB | CVM | COAF | mixed | federal"
}
Rules:
- Both applicable_regulations and relevant_articles must be non-empty arrays.
- No segregation of client assets → HIGH risk.
- No BCB authorization → HIGH risk.
- Weak KYC / anonymous transactions → HIGH risk.
- Token with dividends/voting/public fundraising → CVM securities risk.
- Base answer strictly on the retrieved regulatory context.
"""


def call_claude(query: str, context: str) -> dict | None:
    api_key = os.environ.get("ANTHROPIC_API_KEY", "")
    if not api_key:
        return None
    try:
        r = requests.post(
            "https://api.anthropic.com/v1/messages",
            headers={"Content-Type": "application/json",
                     "x-api-key": api_key,
                     "anthropic-version": "2023-06-01"},
            json={"model": "claude-sonnet-4-20250514", "max_tokens": 1200,
                  "system": COMPLIANCE_SYSTEM,
                  "messages": [{"role": "user",
                                "content": f"COMPLIANCE QUERY:\n{query}\n\nREGULATORY CONTEXT:\n{context}\n\nProduce a structured compliance assessment."}]},
            timeout=90,
        )
        r.raise_for_status()
        raw = "".join(b.get("text", "") for b in r.json().get("content", []) if b.get("type") == "text")
        raw = re.sub(r"^```(?:json)?", "", raw.strip(), flags=re.IGNORECASE).strip().rstrip("```").strip()
        return json.loads(raw)
    except Exception as exc:
        print(f"Claude error: {exc}", flush=True)
        return None


# ── Fine-tuned Mixtral inference (Model B) ────────────────────────────

FT_SYSTEM = """You are RegTech BR, a specialist AI in Brazilian crypto asset regulation trained by Adaption AutoScientist.
Analyze the compliance query and produce a structured JSON assessment.
Respond ONLY with valid JSON — no markdown fences.
Use EXACTLY these snake_case keys:
{
  "risk_level": "LOW | MEDIUM | HIGH | UNCLEAR",
  "compliance_status": "COMPLIANT | NON-COMPLIANT | REQUIRES_REVIEW | INSUFFICIENT_INFO",
  "applicable_regulations": ["list"],
  "relevant_articles": ["list"],
  "finding": "2-5 sentence assessment",
  "corrective_action": "specific steps or 'No action required'",
  "confidence": "HIGH | MEDIUM | LOW",
  "authority": "BCB | CVM | COAF | mixed | federal"
}
Output contract:
- Return exactly one JSON object.
- Use double quotes for all JSON keys and string values.
- Do not include markdown, explanations, headings, bullets, or text before/after the JSON.
- Keep the JSON compact: finding <= 2 sentences, corrective_action <= 1 sentence.
- Use at most 4 applicable_regulations and at most 6 relevant_articles.
- If unsure, use "UNCLEAR", "REQUIRES_REVIEW", "LOW", and cite the closest retrieved context.
"""


@spaces.GPU(duration=GPU_DURATION_SECONDS)
def call_finetuned(query: str, context: str) -> dict | None:
    """Run inference on the fine-tuned Mixtral adapter (requires GPU)."""
    if not _load_finetuned():
        return {"error": _ft_load_error or "Failed to load fine-tuned model"}

    prompt = (
        f"<s>[INST] {FT_SYSTEM}\n\n"
        f"COMPLIANCE QUERY:\n{query}\n\n"
        f"REGULATORY CONTEXT:\n{context}\n\n"
        f"Produce a structured compliance assessment. [/INST]"
    )

    try:
        model_device = next(_ft_model.parameters()).device
        inputs = _ft_tokenizer(prompt, return_tensors="pt").to(model_device)
        with torch.no_grad():
            outputs = _ft_model.generate(
                **inputs,
                max_new_tokens=FT_MAX_NEW_TOKENS,
                do_sample=FT_TEMPERATURE > 0.0,
                temperature=FT_TEMPERATURE if FT_TEMPERATURE > 0.0 else None,
                pad_token_id=_ft_tokenizer.eos_token_id,
                eos_token_id=_ft_tokenizer.eos_token_id,
            )
        raw = _ft_tokenizer.decode(
            outputs[0][inputs["input_ids"].shape[1]:],
            skip_special_tokens=True,
        ).strip()

        print("\n" + "=" * 72, flush=True)
        print("FINE-TUNED RAW OUTPUT START", flush=True)
        print(raw[:3000] if raw else "<EMPTY>", flush=True)
        print("FINE-TUNED RAW OUTPUT END", flush=True)
        print("=" * 72 + "\n", flush=True)

        try:
            parsed = load_model_json(raw)
            parsed["_raw_adapter_output"] = raw
            return parsed
        except json.JSONDecodeError as json_exc:
            print(f"Fine-tuned JSON parse warning: {json_exc}", flush=True)
            return fallback_unstructured_adapter_report(raw)

    except Exception as exc:
        return {"error": f"{type(exc).__name__}: {exc}"}


# ── HTML rendering ────────────────────────────────────────────────────

RISK_COLOR = {"HIGH": "#dc2626", "MEDIUM": "#d97706", "LOW": "#16a34a", "UNCLEAR": "#6b7280"}
STATUS_ICON = {"NON-COMPLIANT": "⛔", "COMPLIANT": "✅", "REQUIRES_REVIEW": "⚠️", "INSUFFICIENT_INFO": "❓"}


def render_panel(report: dict | None, label: str, accent: str, error_msg: str = "") -> str:
    if report is None or "error" in (report or {}):
        err = (report or {}).get("error", error_msg or "No response")
        raw = (report or {}).get("raw") or (report or {}).get("raw_error") or ""
        raw_html = ""
        if raw:
            raw_html = f"""
  <details style="margin-top:0.9rem;background:#111827;border:1px solid #26364f;
                  border-radius:8px;padding:0.75rem;color:#cbd5e1">
    <summary style="cursor:pointer;color:#93c5fd;font-size:0.75rem">Show raw output / error</summary>
    <pre style="white-space:pre-wrap;word-break:break-word;font-size:0.72rem;
                line-height:1.45;color:#cbd5e1;margin-top:0.65rem">{esc(str(raw)[:3000])}</pre>
  </details>"""
        return f"""
<div style="font-family:'IBM Plex Mono',monospace;background:#0f172a;color:#fca5a5;
            padding:1.2rem;border-radius:12px;border:1px solid #7f1d1d;height:100%">
  <div style="color:{accent};font-size:0.72rem;text-transform:uppercase;
              letter-spacing:0.12em;margin-bottom:0.8rem;font-weight:700">{esc(label)}</div>
  ⚠️ {esc(err)}
  {raw_html}
</div>"""

    risk = str(report.get("risk_level", "UNCLEAR")).upper()
    status = str(report.get("compliance_status", "INSUFFICIENT_INFO")).upper().replace("_", "-")
    confidence = str(report.get("confidence", "LOW")).upper()
    authority = str(report.get("authority", "?"))
    color = RISK_COLOR.get(risk, "#6b7280")
    icon = STATUS_ICON.get(status, "❓")
    finding = esc(report.get("finding", "—"))
    corrective = esc(report.get("corrective_action", "—"))
    raw_adapter_output = str(report.get("_raw_adapter_output", "") or "")
    parse_warning = str(report.get("_parse_warning", "") or "")

    raw_adapter_html = ""
    if raw_adapter_output:
        warning_html = (
            f'<div style="color:#fbbf24;font-size:0.74rem;margin-bottom:0.45rem">{esc(parse_warning)}</div>'
            if parse_warning else ""
        )
        raw_adapter_html = f"""
  <details style="margin:0.8rem 0;background:#111827;border:1px solid #26364f;
                  border-radius:8px;padding:0.75rem;color:#cbd5e1">
    <summary style="cursor:pointer;color:{accent};font-size:0.75rem">Adapter raw output</summary>
    {warning_html}
    <pre style="white-space:pre-wrap;word-break:break-word;font-size:0.72rem;
                line-height:1.45;color:#cbd5e1;margin-top:0.65rem">{esc(raw_adapter_output[:3000])}</pre>
  </details>"""

    regs = as_list(report.get("applicable_regulations"))
    arts = as_list(report.get("relevant_articles"))

    regs_html = "".join(f'<div style="color:#93c5fd;font-size:0.78rem;padding:3px 0;border-bottom:1px solid #1e3a5f">• {esc(r)}</div>' for r in regs) or '<div style="color:#475569;font-size:0.78rem">Not specified</div>'
    arts_html = "".join(f'<div style="color:#86efac;font-size:0.78rem;padding:3px 0;border-bottom:1px solid #1e3a5f">• {esc(a)}</div>' for a in arts) or '<div style="color:#475569;font-size:0.78rem">Not specified</div>'

    return f"""
<div style="font-family:'IBM Plex Mono',monospace;background:#0f172a;color:#e2e8f0;
            padding:1.2rem;border-radius:12px;border:2px solid {accent};height:100%">
  <div style="color:{accent};font-size:0.72rem;text-transform:uppercase;
              letter-spacing:0.12em;margin-bottom:0.8rem;font-weight:700">{esc(label)}</div>
  <div style="display:flex;gap:0.6rem;align-items:center;margin-bottom:1rem;flex-wrap:wrap">
    <span style="background:{color};color:#fff;padding:3px 12px;border-radius:5px;
                 font-weight:700;font-size:0.85rem">{esc(risk)} RISK</span>
    <span style="background:#1e293b;color:#e2e8f0;padding:3px 12px;border-radius:5px;
                 font-size:0.85rem;border:1px solid #334155">{icon} {esc(status)}</span>
    <span style="color:#64748b;font-size:0.75rem">confidence: {esc(confidence)} · authority: {esc(authority)}</span>
  </div>
  <div style="background:#1e293b;border-radius:8px;padding:0.85rem;margin-bottom:0.8rem;
              border-left:3px solid {color}">
    <div style="color:#94a3b8;font-size:0.7rem;text-transform:uppercase;
                letter-spacing:0.1em;margin-bottom:0.4rem">Finding</div>
    <div style="color:#f8fafc;font-size:0.83rem;line-height:1.6">{finding}</div>
  </div>
  <div style="background:#1e293b;border-radius:8px;padding:0.85rem;margin-bottom:0.8rem">
    <div style="color:#94a3b8;font-size:0.7rem;text-transform:uppercase;
                letter-spacing:0.1em;margin-bottom:0.4rem">Corrective Action</div>
    <div style="color:#f8fafc;font-size:0.83rem;line-height:1.6">{corrective}</div>
  </div>
  {raw_adapter_html}
  <div style="display:grid;grid-template-columns:1fr 1fr;gap:0.6rem">
    <div style="background:#1e293b;border-radius:8px;padding:0.75rem">
      <div style="color:#94a3b8;font-size:0.7rem;text-transform:uppercase;
                  letter-spacing:0.1em;margin-bottom:0.4rem">Applicable Regulations</div>
      {regs_html}
    </div>
    <div style="background:#1e293b;border-radius:8px;padding:0.75rem">
      <div style="color:#94a3b8;font-size:0.7rem;text-transform:uppercase;
                  letter-spacing:0.1em;margin-bottom:0.4rem">Relevant Articles</div>
      {arts_html}
    </div>
  </div>
</div>"""


def render_diff(report_a: dict | None, report_b: dict | None) -> str:
    """Highlight agreements and divergences between the two models."""
    if not report_a or not report_b or "error" in report_a or "error" in report_b:
        return ""

    risk_a = str(report_a.get("risk_level", "")).upper()
    risk_b = str(report_b.get("risk_level", "")).upper()
    status_a = str(report_a.get("compliance_status", "")).upper().replace("_", "-")
    status_b = str(report_b.get("compliance_status", "")).upper().replace("_", "-")

    agree_risk = risk_a == risk_b
    agree_status = status_a == status_b

    def badge(match: bool) -> str:
        if match:
            return '<span style="background:#14532d;color:#86efac;padding:2px 8px;border-radius:4px;font-size:0.72rem">AGREE ✓</span>'
        return '<span style="background:#7f1d1d;color:#fca5a5;padding:2px 8px;border-radius:4px;font-size:0.72rem">DIVERGE ✗</span>'

    rows = [
        ("Risk Level", risk_a, risk_b, agree_risk),
        ("Compliance Status", status_a, status_b, agree_status),
        ("Confidence", str(report_a.get("confidence", "")).upper(), str(report_b.get("confidence", "")).upper(),
         report_a.get("confidence", "").upper() == report_b.get("confidence", "").upper()),
        ("Authority", str(report_a.get("authority", "?")), str(report_b.get("authority", "?")),
         report_a.get("authority", "").lower() == report_b.get("authority", "").lower()),
    ]

    table_rows = "".join(
        f"""<tr style="border-bottom:1px solid #1e3a5f">
          <td style="padding:6px 10px;color:#94a3b8;font-size:0.78rem">{esc(field)}</td>
          <td style="padding:6px 10px;color:#38bdf8;font-size:0.78rem">{esc(val_a)}</td>
          <td style="padding:6px 10px;color:#a78bfa;font-size:0.78rem">{esc(val_b)}</td>
          <td style="padding:6px 10px">{badge(match)}</td>
        </tr>"""
        for field, val_a, val_b, match in rows
    )

    agreed = sum(1 for *_, m in rows if m)
    total = len(rows)

    return f"""
<div style="font-family:'IBM Plex Mono',monospace;background:#0f172a;color:#e2e8f0;
            padding:1.1rem;border-radius:12px;border:1px solid #334155;margin:1rem 0">
  <div style="color:#94a3b8;font-size:0.72rem;text-transform:uppercase;
              letter-spacing:0.12em;margin-bottom:0.7rem">
    Agreement analysis — {agreed}/{total} fields match
  </div>
  <table style="width:100%;border-collapse:collapse">
    <thead>
      <tr style="border-bottom:1px solid #334155">
        <th style="padding:5px 10px;color:#475569;font-size:0.7rem;text-align:left">Field</th>
        <th style="padding:5px 10px;color:#38bdf8;font-size:0.7rem;text-align:left">Model A — RAG + Claude</th>
        <th style="padding:5px 10px;color:#a78bfa;font-size:0.7rem;text-align:left">Model B — Fine-tuned</th>
        <th style="padding:5px 10px;color:#475569;font-size:0.7rem;text-align:left">Verdict</th>
      </tr>
    </thead>
    <tbody>{table_rows}</tbody>
  </table>
</div>"""


def render_evidence(results: list[dict]) -> str:
    if not results:
        return ""
    rows = []
    for i, r in enumerate(results, 1):
        sid = esc(r.get("source_id", "?"))
        auth = esc(r.get("authority", "?"))
        art = esc(r.get("article_hint", "—") or "—")
        score = f"{float(r.get('_final', 0.0)):.3f}"
        text = esc(str(r.get("text", ""))[:220])
        rows.append(f"""
<div style="background:#111827;border:1px solid #26364f;border-radius:8px;padding:0.75rem;margin:0.4rem 0">
  <div style="display:flex;gap:0.5rem;align-items:center;flex-wrap:wrap;margin-bottom:0.3rem">
    <span style="color:#e2e8f0;font-weight:700;font-size:0.8rem">#{i}</span>
    <span style="background:#1e3a5f;color:#93c5fd;padding:2px 7px;border-radius:4px;font-size:0.7rem">{sid}</span>
    <span style="color:#94a3b8;font-size:0.7rem">{auth} · {art} · score {score}</span>
  </div>
  <div style="color:#cbd5e1;font-size:0.76rem;line-height:1.5">{text}...</div>
</div>""")
    return f"""
<div style="font-family:'IBM Plex Mono',monospace;background:#0f172a;padding:1rem;
            border-radius:12px;border:1px solid #1e3a5f;margin-top:0.8rem">
  <div style="color:#94a3b8;font-size:0.72rem;text-transform:uppercase;
              letter-spacing:0.1em;margin-bottom:0.5rem">Shared RAG context ({len(results)} chunks)</div>
  {''.join(rows)}
</div>"""


def render_error_html(msg: str) -> str:
    return f'<div style="background:#1c1917;color:#fca5a5;padding:1rem;border-radius:8px;font-family:monospace">⚠️ {esc(msg)}</div>'


def render_info_panel(label: str, accent: str, message: str) -> str:
    """Render a neutral panel used when Model B has not been run yet."""
    return f"""
<div style="font-family:'IBM Plex Mono',monospace;background:#0f172a;color:#cbd5e1;
            padding:1.2rem;border-radius:12px;border:2px dashed {accent};height:100%">
  <div style="color:{accent};font-size:0.72rem;text-transform:uppercase;
              letter-spacing:0.12em;margin-bottom:0.8rem;font-weight:700">{esc(label)}</div>
  <div style="background:#111827;border:1px solid #26364f;border-radius:8px;
              padding:0.9rem;font-size:0.82rem;line-height:1.65">
    {esc(message)}
  </div>
</div>"""


def make_model_b_runtime_error(exc: Exception | str) -> dict:
    """Convert ZeroGPU/runtime exceptions into a clear user-facing error."""
    raw = str(exc)
    if "exceeded your ZeroGPU quota" in raw or "Try again in" in raw:
        return {
            "error": (
                "ZeroGPU quota is not enough for this Model B run. "
                "Model A is still available. Wait for the quota reset shown by Hugging Face, "
                "or keep using RAG + Claude until more ZeroGPU time is available. "
                f"This app is currently requesting only {GPU_DURATION_SECONDS}s for Model B; "
                "if the HF error still says 180s requested, the Space is running an older app.py "
                "or a Space variable is overriding the duration."
            ),
            "raw_error": raw,
        }
    if "Some modules are dispatched on the CPU or the disk" in raw:
        return {
            "error": (
                "The Mixtral base model did not fit fully on the current GPU during loading. "
                "This build now defaults to a pre-quantized 4-bit Mixtral base to avoid CPU/disk dispatch. "
                "If this message persists, set BASE_MODEL_ID=ybelkada/Mixtral-8x7B-Instruct-v0.1-bnb-4bit "
                "and USE_PREQUANTIZED_BASE=1, then Factory rebuild. If it still fails, use a 48GB GPU."
            ),
            "raw_error": raw,
        }

    if "GPU task aborted" in raw:
        return {
            "error": (
                "ZeroGPU aborted the Model B task. This usually happens when the Mixtral base model "
                "and LoRA adapter cannot load/generate within the allocated ZeroGPU window. "
                "RAG + Claude is working; rerun Model B after quota reset or on A10G/Persistent GPU."
            ),
            "raw_error": raw,
        }
    return {"error": f"ZeroGPU/Model B runtime error: {type(exc).__name__}: {raw}"}



def model_b_preflight_error() -> dict | None:
    """Return a skip/error report before calling ZeroGPU when the setup is not viable.

    This is intentionally checked outside the @spaces.GPU function, so clicking
    the Model B button does not spend quota or trigger a GPU task abort when the
    configured duration is too small for Mixtral-8x7B + LoRA.
    """
    if not MODEL_B_ENABLED:
        return {
            "error": (
                "Model B GPU execution is disabled by configuration. "
                "Set ENABLE_MODEL_B_GPU=1 in the Space variables to enable it."
            )
        }

    if GPU_DURATION_SECONDS < MODEL_B_MIN_DURATION_SECONDS:
        return {
            "error": (
                f"Model B was not started because the current ZeroGPU request is only "
                f"{GPU_DURATION_SECONDS}s, below the safe minimum of "
                f"{MODEL_B_MIN_DURATION_SECONDS}s for loading Mixtral-8x7B + LoRA. "
                "This avoids another 'GPU task aborted'. "
                "Continue using Model A now. After quota reset, set "
                "ZEROGPU_DURATION_SECONDS=120 and ZEROGPU_MAX_DURATION_SECONDS=120 "
                "or run the Space on A10G/Persistent GPU."
            )
        }

    return None


# ── Main comparison functions ─────────────────────────────────────────

def analyze_with_claude(query: str):
    """Run only RAG + Claude.

    This avoids spending ZeroGPU quota while testing retrieval, Claude output,
    layout, and evidence rendering. Model B is run by a separate button.
    """
    if not query or not query.strip():
        err = render_error_html("Please enter a compliance query.")
        return err, render_info_panel("Model B — Fine-tuned Mixtral-8x7B", "#a78bfa", "Model B has not been run."), "", "", {}

    query = query.strip()
    print(f"\nCLAUDE-FIRST QUERY: {query}", flush=True)

    results = retrieve(query)
    if not results:
        err = render_error_html("No relevant regulatory chunks found.")
        return err, render_info_panel("Model B — Fine-tuned Mixtral-8x7B", "#a78bfa", "Model B was not run because no RAG context was retrieved."), "", "", {}

    context = format_context(results)

    print("Calling Claude only...", flush=True)
    report_a = call_claude(query, context)

    panel_a = render_panel(report_a, "Model A — RAG + Claude Sonnet", "#38bdf8")
    panel_b = render_info_panel(
        "Model B — Fine-tuned Mixtral-8x7B (AutoScientist)",
        "#a78bfa",
        f"Not run yet. RAG + Claude has been analyzed without spending ZeroGPU quota. "
        f"Click 'Run Model B — Fine-tuned Mixtral' only after quota reset or on persistent GPU. "
        f"Current Model B request: {GPU_DURATION_SECONDS}s; safe minimum: "
        f"{MODEL_B_MIN_DURATION_SECONDS}s.",
    )
    evidence = render_evidence(results)

    state = {
        "query": query,
        "context": context,
        "report_a": report_a,
        "results": results,
    }

    return panel_a, panel_b, "", evidence, state


def run_model_b(query: str, state: dict | None):
    """Run only the fine-tuned Mixtral adapter using the current or cached RAG context."""
    if not query or not query.strip():
        err = render_panel({"error": "Please enter a compliance query first."}, "Model B — Fine-tuned Mixtral-8x7B", "#a78bfa")
        return err, "", "", state or {}

    query = query.strip()
    state = state or {}

    same_query = state.get("query") == query and state.get("context") and state.get("results")
    if same_query:
        print("\nUsing cached RAG context for Model B.", flush=True)
        context = state["context"]
        results = state["results"]
        report_a = state.get("report_a")
    else:
        print("\nNo cached context for this query. Retrieving context before Model B.", flush=True)
        results = retrieve(query)
        if not results:
            err = render_panel({"error": "No relevant regulatory chunks found."}, "Model B — Fine-tuned Mixtral-8x7B", "#a78bfa")
            return err, "", "", state
        context = format_context(results)
        report_a = None

    preflight = model_b_preflight_error()
    if preflight is not None:
        print(f"Skipping Model B before ZeroGPU call: {preflight['error']}", flush=True)
        panel_b = render_panel(preflight, "Model B — Fine-tuned Mixtral-8x7B (AutoScientist)", "#a78bfa")
        evidence = render_evidence(results)
        new_state = {
            "query": query,
            "context": context,
            "report_a": report_a,
            "report_b": preflight,
            "results": results,
        }
        return panel_b, "", evidence, new_state

    print(
        f"Calling fine-tuned Mixtral with ZeroGPU duration={GPU_DURATION_SECONDS}s "
        f"and max_new_tokens={FT_MAX_NEW_TOKENS}...",
        flush=True,
    )

    try:
        report_b = call_finetuned(query, context)
    except Exception as exc:
        report_b = make_model_b_runtime_error(exc)
        print(f"Model B exception caught: {report_b.get('error')}", flush=True)
        if report_b.get("raw_error"):
            print(f"Model B raw error: {report_b['raw_error']}", flush=True)

    panel_b = render_panel(report_b, "Model B — Fine-tuned Mixtral-8x7B (AutoScientist)", "#a78bfa")
    diff = render_diff(report_a, report_b) if report_a else ""
    evidence = render_evidence(results)

    new_state = {
        "query": query,
        "context": context,
        "report_a": report_a,
        "report_b": report_b,
        "results": results,
    }

    return panel_b, diff, evidence, new_state


# ── Gradio UI ─────────────────────────────────────────────────────────

CSS = """
@import url('https://fonts.googleapis.com/css2?family=IBM+Plex+Mono:wght@400;500;700&family=IBM+Plex+Sans:wght@300;400;600&display=swap');
body { background: #0a0f1e !important; }
.gradio-container {
    background: #0a0f1e !important;
    font-family: 'IBM Plex Sans', sans-serif !important;
    max-width: 1200px !important;
    margin: 0 auto !important;
}
.query-box textarea {
    background: #0b1220 !important;
    border: 1px solid #1e3a5f !important;
    color: #e2e8f0 !important;
    font-family: 'IBM Plex Mono', monospace !important;
    font-size: 0.9rem !important;
    border-radius: 8px !important;
}
.query-box textarea:focus {
    border-color: #38bdf8 !important;
    box-shadow: 0 0 0 2px rgba(56,189,248,0.15) !important;
}
.compare-btn {
    background: #0369a1 !important;
    color: #fff !important;
    font-family: 'IBM Plex Mono', monospace !important;
    font-weight: 700 !important;
    border: 1px solid rgba(56,189,248,0.35) !important;
    border-radius: 8px !important;
    height: 48px !important;
    font-size: 0.84rem !important;
    box-shadow: none !important;
}
.compare-btn:hover { background: #0284c7 !important; }
.modelb-btn {
    background: #6d28d9 !important;
    color: #fff !important;
    font-family: 'IBM Plex Mono', monospace !important;
    font-weight: 700 !important;
    border: 1px solid rgba(167,139,250,0.45) !important;
    border-radius: 8px !important;
    height: 48px !important;
    font-size: 0.78rem !important;
    box-shadow: none !important;
}
.modelb-btn:hover { background: #7c3aed !important; }
label { color: #94a3b8 !important; font-size: 0.8rem !important; }
.quota-note {
    font-family: 'IBM Plex Mono', monospace;
    color: #94a3b8;
    font-size: 0.74rem;
    background: #0f172a;
    border: 1px solid #172554;
    border-radius: 10px;
    padding: 0.75rem 0.9rem;
    margin: 0.7rem 0 0.9rem;
}
.shortcut-panel {
    background: transparent;
    border: 1px solid #172554;
    border-radius: 12px;
    padding: 0.75rem 0.9rem;
    margin: 0.9rem 0 0.75rem;
}
.shortcut-title {
    color: #94a3b8;
    font-family: 'IBM Plex Mono', monospace;
    font-size: 0.72rem;
    text-transform: uppercase;
    letter-spacing: 0.14em;
    margin: 0;
}
.chip-btn {
    background: #0f172a !important;
    color: #bfdbfe !important;
    border: 1px solid #1e3a5f !important;
    border-radius: 999px !important;
    min-height: 34px !important;
    height: 34px !important;
    padding: 0 0.85rem !important;
    font-family: 'IBM Plex Mono', monospace !important;
    font-size: 0.76rem !important;
    font-weight: 600 !important;
    box-shadow: none !important;
}
.chip-btn:hover {
    background: #111827 !important;
    color: #e0f2fe !important;
    border-color: #38bdf8 !important;
}
"""

# Use regular buttons instead of gr.Examples.
# Gradio 4.44.0 can raise "TypeError: unhashable type: dict" or API schema errors
# when Examples is rendered on some Hugging Face Spaces builds.
SHORTCUT_QUERIES = {
    "Asset Segregation": "Our exchange does not segregate client virtual assets from corporate funds. Which Brazilian custody and asset segregation obligations apply?",
    "AML/KYC Controls": "Our platform performs full KYC only for transactions above BRL 100,000 and allows smaller transfers with simplified identification. Which AML/CFT controls should be assessed?",
    "Securities Token": "Our REV token pays revenue share, grants voting rights, and will be publicly offered to Brazilian investors. Does this create securities risk?",
    "BCB Authorization": "Our platform operates virtual asset exchange and custody services in Brazil without formal authorization from the Central Bank of Brazil. What is the compliance exposure?",
    "Cross-border Custody": "A foreign crypto custodian supports Brazilian clients and settles cross-border transactions using stablecoins. Which Brazilian regulatory obligations should be assessed?",
}


def load_shortcut(label: str) -> str:
    return SHORTCUT_QUERIES.get(label, "")

with gr.Blocks(css=CSS, title="RegTech BR — Model Comparison") as demo:
    gr.HTML("""
    <div style="text-align:center;padding:2rem 0 1.4rem;border-bottom:1px solid #172554;margin-bottom:1.5rem">
      <h1 style="font-family:'IBM Plex Mono',monospace;font-size:1.8rem;color:#38bdf8;margin:0 0 0.4rem">
        ⚖ RegTech BR — Model Comparison
      </h1>
      <p style="color:#cbd5e1;font-size:0.88rem;margin:0">
        A compliance reasoning benchmark: retrieval-augmented generation vs domain-adapted fine-tuning
      </p>
      <p style="color:#64748b;font-size:0.74rem;margin:0.45rem 0 0;font-family:'IBM Plex Mono',monospace">
        Same query · Same legal context · Side-by-side model evaluation
      </p>
    </div>
    """)

    session_state = gr.State({})

    with gr.Row():
        with gr.Column(scale=5):
            query_box = gr.Textbox(
                label="Compliance query",
                placeholder="Describe a crypto product, policy, control gap, or compliance question...",
                lines=4,
                elem_classes=["query-box"],
            )
        with gr.Column(scale=1, min_width=170):
            compare_btn = gr.Button("Run Model A — RAG + Claude →", elem_classes=["compare-btn"])
            run_b_btn = gr.Button("Run Model B — Fine-tuned Mixtral →", elem_classes=["modelb-btn"])

    gr.HTML("""
    <div class="quota-note">
      Model B uses ZeroGPU. Run Model A first; trigger Model B when quota is available.
    </div>
    """)

    gr.HTML("""
    <div class="shortcut-panel">
      <div class="shortcut-title">Reference scenarios</div>
    </div>
    """)
    with gr.Row():
        seg_btn = gr.Button("Asset Segregation", elem_classes=["chip-btn"])
        kyc_btn = gr.Button("AML/KYC Controls", elem_classes=["chip-btn"])
        cvm_btn = gr.Button("Securities Token", elem_classes=["chip-btn"])
        bcb_btn = gr.Button("BCB Authorization", elem_classes=["chip-btn"])
        en_btn = gr.Button("Cross-border Custody", elem_classes=["chip-btn"])

    diff_html = gr.HTML(label="Agreement analysis")

    with gr.Row(equal_height=True):
        panel_a_html = gr.HTML(label="Model A — RAG + Claude")
        panel_b_html = gr.HTML(label="Model B — Fine-tuned Mixtral")

    evidence_html = gr.HTML(label="Shared RAG context")

    gr.HTML("""
    <div style="text-align:center;color:#64748b;font-size:0.72rem;
                padding:1.5rem 0 0.5rem;font-family:'IBM Plex Mono',monospace;line-height:1.6">
      ⚠ Experimental compliance screening. Not legal advice. Results require review by qualified professionals. · RegTech BR · Fernando Rodrigues · AutoScientist Challenge 2026 · Adaption
    </div>
    """)

    seg_btn.click(fn=lambda: load_shortcut("Asset Segregation"), inputs=None, outputs=[query_box], api_name=False)
    kyc_btn.click(fn=lambda: load_shortcut("AML/KYC Controls"), inputs=None, outputs=[query_box], api_name=False)
    cvm_btn.click(fn=lambda: load_shortcut("Securities Token"), inputs=None, outputs=[query_box], api_name=False)
    bcb_btn.click(fn=lambda: load_shortcut("BCB Authorization"), inputs=None, outputs=[query_box], api_name=False)
    en_btn.click(fn=lambda: load_shortcut("Cross-border Custody"), inputs=None, outputs=[query_box], api_name=False)

    compare_btn.click(
        fn=analyze_with_claude,
        inputs=[query_box],
        outputs=[panel_a_html, panel_b_html, diff_html, evidence_html, session_state],
        api_name=False,
    )
    query_box.submit(
        fn=analyze_with_claude,
        inputs=[query_box],
        outputs=[panel_a_html, panel_b_html, diff_html, evidence_html, session_state],
        api_name=False,
    )
    run_b_btn.click(
        fn=run_model_b,
        inputs=[query_box, session_state],
        outputs=[panel_b_html, diff_html, evidence_html, session_state],
        api_name=False,
    )

if __name__ == "__main__":
    port = int(os.environ.get("PORT", 7860))
    demo.queue().launch(server_name="0.0.0.0", server_port=port,
                         share=True, show_api=False)