server.py

"""
FastAPI Server Entrypoint
- Server starts INSTANTLY — frontend is available immediately
- All heavy data/ML/graph work runs in a background thread
- Frontend polls /api/status to show loading progress
"""
import json
import math
import os
import time
import threading
import uvicorn
from fastapi import FastAPI, HTTPException, Request
from fastapi.middleware.cors import CORSMiddleware
from fastapi.staticfiles import StaticFiles
from fastapi.responses import FileResponse, JSONResponse
import pickle
import pandas as pd
import networkx as nx


def _sanitize_nan(obj):
    """Recursively replace NaN/Inf floats with None so json.dumps never raises."""
    if isinstance(obj, float):
        return None if not math.isfinite(obj) else obj
    if isinstance(obj, dict):
        return {k: _sanitize_nan(v) for k, v in obj.items()}
    if isinstance(obj, list):
        return [_sanitize_nan(v) for v in obj]
    return obj


class _NaNSafeJSONResponse(JSONResponse):
    def render(self, content) -> bytes:
        return json.dumps(
            _sanitize_nan(content),
            ensure_ascii=False,
            separators=(',', ':'),
        ).encode('utf-8')

# Load .env before anything reads env vars
try:
    from dotenv import load_dotenv
    load_dotenv()
except ImportError:
    pass

from src.config_loader import get_config
from src.persistence import init_db
from src.state import AppState
from src.case_mgmt.db import init_cases_db
from src.case_mgmt.config import get_settings as get_case_settings
from src.case_mgmt.auth import get_current_user
from src.entity_graph_service import (
    load_accounts_reference,
    build_entity_bank_indices,
    build_bank_profiles,
    attach_risk_aggregates,
)
from src.entity_cluster_service import attach_cluster_aggregates, build_clusters_index
from src.risk_tiers import compute_percentile_tiers


# ── HF Hub model download (runs once if models/checkpoints are missing) ──
_HF_MODEL_REPO = "Aniket2006/fund-flow-models"
_MODEL_FILES = [
    ("models/gnn_model.pt",                  "models/gnn_model.pt"),
    ("models/xgb_fraud_model.ubj",           "models/xgb_fraud_model.ubj"),
    ("models/xgb_scaler.pkl",                "models/xgb_scaler.pkl"),
    ("models/hybrid_gnn_encoder.pth",        "models/hybrid_gnn_encoder.pth"),
    ("models/hybrid_xgb_model.json",         "models/hybrid_xgb_model.json"),
    ("models/hybrid_calibrator.pkl",         "models/hybrid_calibrator.pkl"),
    ("data/checkpoints/graph_edges.parquet", "data/checkpoints/graph_edges.parquet"),
    ("data/checkpoints/features.pkl",        "data/checkpoints/features.pkl"),
    ("data/checkpoints/alerts.pkl",          "data/checkpoints/alerts.pkl"),
    ("data/checkpoints/pagerank.pkl",        "data/checkpoints/pagerank.pkl"),
    ("data/checkpoints/betweenness.pkl",     "data/checkpoints/betweenness.pkl"),
    ("data/checkpoints/louvain.pkl",         "data/checkpoints/louvain.pkl"),
    ("data/checkpoints/gnn_metrics.pkl",     "data/checkpoints/gnn_metrics.pkl"),
]

def _is_lfs_pointer(path: str) -> bool:
    """Return True if the file is a Git LFS pointer (starts with 'version https://git-lfs')."""
    try:
        with open(path, 'rb') as f:
            header = f.read(27)
        return header.startswith(b'version https://git-lfs')
    except OSError:
        return False

def _ensure_models():
    """Download model and checkpoint files from HF Hub if missing or still an LFS pointer."""
    needed = [
        (repo_file, local_path)
        for repo_file, local_path in _MODEL_FILES
        if not os.path.exists(local_path) or _is_lfs_pointer(local_path)
    ]
    if not needed:
        return
    try:
        from huggingface_hub import hf_hub_download
        token = os.environ.get("HF_TOKEN")
        for repo_file, local_path in needed:
            os.makedirs(os.path.dirname(local_path), exist_ok=True)
            print(f"  ↓ Downloading {repo_file} from HF Hub...")
            hf_hub_download(
                repo_id=_HF_MODEL_REPO,
                filename=repo_file,
                local_dir=".",
                force_download=True,
                token=token,
            )
            print(f"  ✓ {local_path} ready.")
    except Exception as e:
        print(f"  ✗ Model/checkpoint download failed: {e} — assets may be unavailable.")


class Checkpoint:
    DIR = "data/checkpoints"

    @classmethod
    def load(cls, name):
        os.makedirs(cls.DIR, exist_ok=True)
        path = os.path.join(cls.DIR, f"{name}.pkl")
        if os.path.exists(path):
            print(f"  ✓ Loading {name} from checkpoint...")
            try:
                with open(path, 'rb') as f:
                    return pickle.load(f)
            except Exception as e:
                print(f"  ✗ Failed to load {name} checkpoint: {e}")
        return None

    @classmethod
    def save(cls, name, data):
        os.makedirs(cls.DIR, exist_ok=True)
        path = os.path.join(cls.DIR, f"{name}.pkl")
        try:
            with open(path, 'wb') as f:
                pickle.dump(data, f)
            print(f"  ✓ Saved {name} checkpoint to disk.")
        except Exception as e:
            print(f"  ✗ Failed to save {name} checkpoint: {e}")

    @classmethod
    def save_graph_parquet(cls, graph: nx.MultiDiGraph):
        """Save graph as a Parquet edge list — much smaller and faster to reload."""
        os.makedirs(cls.DIR, exist_ok=True)
        path = os.path.join(cls.DIR, "graph_edges.parquet")
        rows = []
        for u, v, d in graph.edges(data=True):
            rows.append({
                'source': u, 'target': v,
                'amount': d.get('amount', 0),
                'payment_type': d.get('payment_type', ''),
                'is_laundering': d.get('is_laundering', 0),
                'timestamp': d.get('timestamp', ''),
            })
        pd.DataFrame(rows).to_parquet(path, index=False)
        print(f"  ✓ Saved graph edge list ({len(rows):,} edges) to Parquet.")

    @classmethod
    def load_graph_parquet(cls) -> nx.MultiDiGraph:
        """Rebuild graph from Parquet edge list — 10-20x faster than pickle."""
        path = os.path.join(cls.DIR, "graph_edges.parquet")
        if not os.path.exists(path):
            return None
        t0 = time.time()
        df = pd.read_parquet(path)
        G = nx.MultiDiGraph()
        for row in df.itertuples(index=False):
            G.add_edge(
                row.source, row.target,
                amount=row.amount,
                payment_type=row.payment_type,
                is_laundering=int(row.is_laundering),
                timestamp=str(row.timestamp),
            )
        elapsed = time.time() - t0
        print(f"  ✓ Rebuilt graph from Parquet in {elapsed:.1f}s ({G.number_of_nodes():,} nodes, {G.number_of_edges():,} edges).")
        return G


# ── FastAPI App ──────────────────────────────────────────────
app = FastAPI(title="Fund Flow Tracker API", default_response_class=_NaNSafeJSONResponse)

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)


PUBLIC_API_PATHS = {"/api/status", "/api/config"}
# Copilot routes are served within the authenticated UI context;
# the server-level JWT guard is not required for these endpoints.
COPILOT_API_PREFIX = "/api/copilot"


@app.middleware("http")
async def require_supabase_auth_for_api(request: Request, call_next):
    """
    When Supabase auth is configured, protect API data routes with a user JWT.
    Public bootstrapping endpoints remain open so the frontend can load config.
    """
    path = request.url.path.rstrip("/") or "/"
    settings = get_case_settings()
    auth_enabled = bool(settings['supabase_url'] and settings['supabase_anon_key'])

    if (
        auth_enabled
        and request.method != "OPTIONS"
        and path.startswith("/api/")
        and path not in PUBLIC_API_PATHS
        and not path.startswith(COPILOT_API_PREFIX)
    ):
        authorization = request.headers.get("authorization", "")
        if not authorization.startswith("Bearer "):
            return JSONResponse(
                status_code=401,
                content={"detail": "Authentication required"},
            )
        try:
            await get_current_user(authorization)
        except HTTPException as exc:
            return JSONResponse(
                status_code=exc.status_code,
                content={"detail": exc.detail},
            )

    return await call_next(request)


# ── Startup Event (instant — no blocking work) ──────────────
@app.on_event("startup")
async def startup_event():
    init_db()
    init_cases_db()
    _ensure_models()
    print("═══════════════════════════════════════════════════")
    print("  Server started. Frontend is available NOW.")
    print("  Background data loading initiated...")
    print("═══════════════════════════════════════════════════")
    # Spawn background thread for ALL heavy work
    t = threading.Thread(target=background_init, daemon=True)
    t.start()


def _update_status(step: str, done: int):
    """Helper to update the global loading status."""
    AppState.startup_status['current_step'] = step
    AppState.startup_status['steps_done'] = done
    print(f"  [{done}/{AppState.startup_status['total_steps']}] {step}")


def enrich_alerts_with_ml():
    """
    Globally enrich and calibrate alerts using a detector-dominant blend:

    The XGBoost model is miscalibrated (probabilities cluster near 0.97 for
    ~85% of the population — see src/risk_tiers.py). Blending it at 50% weight
    collapses every alert's risk_score to ~97. Instead we use the rule-engine
    detector score as the dominant signal and cap the XGBoost contribution.

    Blend weights (no GNN):
      70% detector (original rule-engine score) + 30% XGBoost percentile rank
    With GNN:
      60% detector + 20% XGBoost percentile rank + 20% GNN
    With Hybrid:
      50% detector + 15% XGBoost percentile rank + 10% GNN + 25% Hybrid

    XGBoost is converted to a PERCENTILE RANK score (0-99) over the full
    account population before blending — this de-compresses the 96-97 clump
    into a meaningful 0-99 distribution, exactly as risk_tier.py does for tiers.

    After blending, alert risk_tiers are re-ranked by percentile over the
    blended alert scores (not inherited from the per-account tier computed on
    raw XGBoost probabilities), so CRITICAL/HIGH/MEDIUM/LOW spread correctly
    across the alert population.
    """
    if not AppState.alerts or not AppState.features_by_account:
        return

    print("  [Enrichment] Calibrating alert risk scores with XGBoost & GNN...")

    # --- Build a percentile-rank lookup for xgb_score over the full population ---
    # This de-compresses the 96-97 probability clump into a 0-99 spread.
    from scipy.stats import rankdata as _rankdata
    _acct_xgb = {
        acct: (feats.get('risk_score') or 0)
        for acct, feats in AppState.features_by_account.items()
    }
    if _acct_xgb:
        _accounts_list = list(_acct_xgb.keys())
        _raw_scores = [_acct_xgb[a] for a in _accounts_list]
        _ranks = _rankdata(_raw_scores, method='average')
        _n = len(_accounts_list)
        _xgb_pct_rank = {
            acct: int(round((rank / _n) * 99))
            for acct, rank in zip(_accounts_list, _ranks)
        }
    else:
        _xgb_pct_rank = {}

    cfg = get_config()['ml']
    for a in AppState.alerts:
        acct_feats = AppState.features_by_account.get(a['account'], {})
        xgb_score_raw = acct_feats.get('risk_score')
        gnn_prob = acct_feats.get('gnn_fraud_score') or 0.0
        hybrid_score_prob = acct_feats.get('hybrid_score') or 0.0

        # Preserve the original detector score so multi-pass calls don't degrade it
        if 'detector_risk_score' not in a:
            a['detector_risk_score'] = int(a.get('risk_score') or 0)

        orig_score = a['detector_risk_score']

        if xgb_score_raw is not None:
            # Use percentile-rank version of XGBoost score (de-clumped)
            xgb_pct = _xgb_pct_rank.get(a['account'], int(xgb_score_raw))
            gnn_score = int(round(gnn_prob * 100.0))
            hybrid_score = int(round(hybrid_score_prob * 100.0))

            if cfg.get('hybrid_enabled', False) and hybrid_score > 0:
                blended = int(round(
                    orig_score * 0.50
                    + xgb_pct   * 0.15
                    + gnn_score * 0.10
                    + hybrid_score * 0.25
                ))
            elif gnn_score > 0:
                blended = int(round(
                    orig_score * 0.60
                    + xgb_pct   * 0.20
                    + gnn_score * 0.20
                ))
            else:
                # Detector-dominant: rule-engine score drives the number,
                # XGBoost percentile rank provides a modest population-relative lift
                blended = int(round(orig_score * 0.70 + xgb_pct * 0.30))

            blended = max(1, min(blended, 99))
            a['risk_score'] = blended
            a['fraud_probability'] = float(acct_feats.get('fraud_probability') or 0.0)

    # Re-rank alert tiers by percentile over the blended alert scores.
    # The per-account risk_tier (computed from raw XGBoost probabilities) is
    # unreliable when all accounts have the same miscalibrated probability —
    # every alert would inherit the same tier. Ranking by the blended
    # risk_score within the alert population ensures a meaningful
    # CRITICAL/HIGH/MEDIUM/LOW spread across what is already a pre-filtered
    # high-risk subset.
    alert_scores = {a['alert_id']: a.get('risk_score', 0) for a in AppState.alerts}
    alert_tiers = compute_percentile_tiers(alert_scores)
    for a in AppState.alerts:
        a['risk_tier'] = alert_tiers.get(a['alert_id'], 'LOW')

    # Re-compute global cached counts
    crit_count = 0
    typo_counts = {}
    for a in AppState.alerts:
        if a['risk_tier'] == 'CRITICAL':
            crit_count += 1
        for t in [t.strip() for t in a['typology'].split(',')]:
            typo_counts[t] = typo_counts.get(t, 0) + 1

    AppState.cached_crit_alerts = crit_count
    AppState.cached_typo_counts = typo_counts
    print(f"  [Enrichment] ✓ Calibrated {len(AppState.alerts):,} alerts. (Critical: {crit_count})",
          f" Score range: {min(alert_scores.values(), default=0)}-{max(alert_scores.values(), default=0)}")


def background_init():
    """Two-phase startup:
    Phase 1 (~8s)  — load all checkpoints, score accounts → startup_ready = True
    Phase 2 (~80s) — build NetworkX graph + load df    → graph_ready   = True
    """
    try:
        from src.ml.predictor import load_model, score_accounts_batch

        # ── PHASE 1: checkpoint loading (dashboard-ready) ────────────

        # Step 1: ML features
        _update_status("Loading ML features...", 1)
        feats = Checkpoint.load('features')
        if feats is not None and len(feats) > 0:
            AppState.full_features = feats
            print(f"  ✓ Loaded {len(feats):,} ML features from checkpoint.")
        else:
            AppState.full_features = pd.DataFrame()
            print("  ✗ Features checkpoint empty — will re-engineer after graph loads.")

        # Step 2: Centrality checkpoints
        _update_status("Loading centrality checkpoints...", 2)
        AppState.pagerank_scores   = Checkpoint.load('pagerank')   or {}
        AppState.betweenness_scores= Checkpoint.load('betweenness')or {}
        AppState.louvain_partition = Checkpoint.load('louvain')    or {}
        AppState.gnn_metrics       = Checkpoint.load('gnn_metrics')or {}
        print(f"  ✓ PageRank({len(AppState.pagerank_scores):,})  "
              f"Betweenness({len(AppState.betweenness_scores):,})  "
              f"Louvain({len(AppState.louvain_partition):,})")

        # Step 3: Alerts
        _update_status("Loading alerts checkpoint...", 3)
        alerts = Checkpoint.load('alerts')
        AppState.alerts = alerts if alerts is not None else []
        print(f"  ✓ Loaded {len(AppState.alerts):,} alerts from checkpoint.")

        # Step 4: XGBoost model
        _update_status("Loading XGBoost model...", 4)
        bundle = load_model()
        if bundle:
            AppState.xgb_bundle = bundle
            feature_cols = bundle['feature_cols']
            AppState.model_metrics = {
                'n_features': len(feature_cols),
                'feature_cols': feature_cols,
            }
            print("  ✓ XGBoost model loaded.")
        else:
            AppState.xgb_bundle = None
            AppState.model_metrics = {}
            print("  ✗ XGBoost model not found — will train after graph loads.")

        # Step 5: Batch scoring
        _update_status("Scoring all accounts...", 5)
        try:
            acct_ids = AppState.full_features['account'].tolist() \
                if 'account' in AppState.full_features.columns else []
            if AppState.xgb_bundle and acct_ids:
                scores = score_accounts_batch(acct_ids, AppState.full_features, AppState.xgb_bundle)
                scores_df = pd.DataFrame([
                    {'account': a, 'risk_score': s['risk_score'] or 0,
                     'fraud_probability': s['fraud_probability'] or 0.0}
                    for a, s in scores.items()
                ])
                # Calibrate/de-compress the risk_score using percentile rank over the population
                from scipy.stats import rankdata as _rankdata
                _raw_probas = scores_df['fraud_probability'].values
                _ranks = _rankdata(_raw_probas, method='average')
                _n = len(_raw_probas)
                scores_df['risk_score'] = (_ranks / _n * 99).round().astype(int)

                AppState.full_features = AppState.full_features.drop(
                    columns=['risk_score', 'fraud_probability'], errors='ignore')
                AppState.full_features = AppState.full_features.merge(scores_df, on='account', how='left')
                AppState.full_features['risk_score'] = AppState.full_features['risk_score'].fillna(0).astype(int)
                AppState.full_features['fraud_probability'] = AppState.full_features['fraud_probability'].fillna(0.0)
                print(f"  ✓ Scored & calibrated {len(acct_ids):,} accounts.")
            else:
                if 'account' in AppState.full_features.columns:
                    AppState.full_features['risk_score'] = 0
                    AppState.full_features['fraud_probability'] = 0.0
        except Exception as e:
            print(f"  ✗ Scoring failed: {e}")
            AppState.startup_status['errors'].append(f"Scoring: {e}")
            if hasattr(AppState.full_features, 'columns') and 'risk_score' not in AppState.full_features.columns:
                AppState.full_features['risk_score'] = 0
                AppState.full_features['fraud_probability'] = 0.0

        # GNN scores deferred to Phase 2 (needs graph)
        if hasattr(AppState.full_features, 'columns') and 'gnn_fraud_score' not in AppState.full_features.columns:
            AppState.full_features['gnn_fraud_score'] = 0.0

        # Risk tier by percentile rank, not absolute score — the XGBoost
        # model's probabilities are badly miscalibrated (see src/risk_tiers.py),
        # so static score>=75 cutoffs flag the majority of accounts. Recomputed
        # again in Phase 2 once hybrid/GNN scores refine fraud_probability.
        if hasattr(AppState.full_features, 'columns') and 'account' in AppState.full_features.columns:
            tier_scores = dict(zip(AppState.full_features['account'], AppState.full_features['fraud_probability']))
            tiers = compute_percentile_tiers(tier_scores)
            AppState.full_features['risk_tier'] = AppState.full_features['account'].map(tiers).fillna('LOW')

        # model_auc was always 0.0 on the dashboard when a pre-trained model is
        # loaded from checkpoint, since 'auc_roc' is only set by trainer.py's
        # fresh-train path. Compute it here too — NOTE this scores against the
        # full population (no train/holdout split survives a checkpoint reload),
        # so it's an in-sample estimate, optimistic vs. the true holdout AUC.
        try:
            if (hasattr(AppState.full_features, 'columns')
                    and 'fraud_flag' in AppState.full_features.columns
                    and 'fraud_probability' in AppState.full_features.columns
                    and AppState.model_metrics is not None):
                from sklearn.metrics import roc_auc_score
                AppState.model_metrics['auc_roc'] = float(roc_auc_score(
                    AppState.full_features['fraud_flag'], AppState.full_features['fraud_probability']))
        except Exception as e:
            print(f"  ✗ model_auc computation failed: {e}")

        # Step 6: O(1) account lookup cache
        _update_status("Building account lookup cache...", 6)
        if hasattr(AppState.full_features, 'columns') and 'account' in AppState.full_features.columns:
            AppState.features_by_account = AppState.full_features.set_index('account').to_dict('index')
        else:
            AppState.features_by_account = {}

        # Load accounts metadata CSV for Account Intelligence + Network Intelligence features
        _accounts_csv_candidates = [
            "data/raw/accounts.csv",
            "data/raw/HI-Small_accounts.csv",
        ]
        for _csv_path in _accounts_csv_candidates:
            if os.path.exists(_csv_path):
                try:
                    AppState.accounts_by_number = load_accounts_reference(_csv_path)
                    print(f"  ✓ Loaded {len(AppState.accounts_by_number):,} account records from {_csv_path}")

                    AppState.entities_by_id, AppState.banks_by_id, _network_summary = \
                        build_entity_bank_indices(AppState.accounts_by_number)
                    AppState.bank_profiles_cache, _top_banks = \
                        build_bank_profiles(AppState.entities_by_id, AppState.banks_by_id)
                    _network_summary["top_banks_by_volume"] = _top_banks
                    AppState.network_summary_cache = _network_summary
                    print(f"  ✓ Indexed {len(AppState.entities_by_id):,} entities across {len(AppState.banks_by_id):,} banks")

                    _risk_summary = attach_risk_aggregates(
                        AppState.entities_by_id, AppState.banks_by_id,
                        AppState.bank_profiles_cache, AppState.features_by_account,
                    )
                    AppState.network_summary_cache["high_risk_entity_count"] = _risk_summary["high_risk_entity_count"]
                    print(f"  ✓ Entity risk scored: {_risk_summary['high_risk_entity_count']:,} high/critical-risk entities")

                    attach_cluster_aggregates(
                        AppState.entities_by_id, AppState.accounts_by_number,
                        AppState.louvain_partition, AppState.betweenness_scores,
                    )
                    AppState.clusters_by_id = build_clusters_index(AppState.entities_by_id)
                    print(f"  ✓ Mule-ring clusters indexed: {len(AppState.clusters_by_id):,} clusters")
                except Exception as _e:
                    print(f"  ✗ Accounts CSV load failed ({_csv_path}): {_e}")
                break

        # Globally enrich alerts with ML metadata using our unified calibration helper.
        # This blends the per-typology detector scores with XGBoost & GNN scores.
        enrich_alerts_with_ml()

        # O(1) account -> alerts index (avoids an O(n_accounts * n_alerts) scan
        # per entity when Network Intelligence aggregates typology breakdowns).
        _alerts_by_account: dict = {}
        for _a in AppState.alerts:
            _alerts_by_account.setdefault(_a['account'], []).append(_a)
        AppState.alerts_by_account = _alerts_by_account

        # Step 7: Pre-compute API caches from graph_edges.parquet (no graph object needed)
        _update_status("Pre-computing API caches...", 7)
        try:
            edge_path = os.path.join(Checkpoint.DIR, 'graph_edges.parquet')
            if os.path.exists(edge_path):
                edf = pd.read_parquet(edge_path)
                channel_counts       = edf['payment_type'].value_counts()
                fraud_channel_counts = edf[edf['is_laundering'] == 1]['payment_type'].value_counts()
                AppState.cached_channel_stats = [
                    {'channel': p, 'count': int(c), 'fraud_count': int(fraud_channel_counts.get(p, 0))}
                    for p, c in channel_counts.items()
                ]
                AppState.cached_overview = {
                    'total_transactions': len(edf),
                    'flagged_transactions': int(edf['is_laundering'].sum()),
                    'total_volume': float(edf['amount'].sum()),
                }
                del edf
        except Exception as e:
            print(f"  ✗ Cache pre-compute failed: {e}")

        # ── PHASE 1 COMPLETE ────────────────────────────────────────
        AppState.startup_ready = True
        AppState.startup_status['current_step'] = 'Ready'
        print("═══════════════════════════════════════════════════")
        print("  ✓ Dashboard ready!  (graph loading in background)")
        print("═══════════════════════════════════════════════════")

        # ── PHASE 2: graph + df (investigation / subgraph) ──────────
        _phase2_load_graph()

    except Exception as e:
        print(f"  ✗✗✗ CRITICAL background init error: {e}")
        AppState.startup_status['current_step'] = f'FAILED: {e}'
        AppState.startup_status['errors'].append(f"Critical: {e}")


def _phase2_load_graph():
    """Load transactions.parquet and build the NetworkX graph.
    Runs after startup_ready=True so it doesn't block the dashboard."""
    try:
        from src.data_loader import get_processed_data
        from src.graph_builder import build_graph
        from src.ml.gnn_predictor import predict_gnn_score

        print("  [Phase 2] Loading transaction dataset...")
        try:
            AppState.df, AppState.node_features = get_processed_data()
            # Refresh cached_overview / channel_stats with exact values from df
            if AppState.df is not None and len(AppState.df) > 0:
                channel_counts       = AppState.df['payment_type'].value_counts()
                fraud_channel_counts = AppState.df[AppState.df['is_laundering'] == 1]['payment_type'].value_counts()
                AppState.cached_channel_stats = [
                    {'channel': p, 'count': int(c), 'fraud_count': int(fraud_channel_counts.get(p, 0))}
                    for p, c in channel_counts.items()
                ]
                AppState.cached_overview = {
                    'total_transactions': len(AppState.df),
                    'flagged_transactions': int((AppState.df['is_laundering'] == 1).sum()),
                    'total_volume': float(AppState.df['amount'].sum()),
                }
        except Exception as e:
            print(f"  [Phase 2] ✗ Dataset load failed: {e}")
            AppState.startup_status['errors'].append(f"Dataset: {e}")

        print("  [Phase 2] Building transaction graph...")
        try:
            parquet_graph = Checkpoint.load_graph_parquet()
            if parquet_graph is not None:
                AppState.graph = parquet_graph
            elif AppState.df is not None:
                AppState.graph = build_graph(AppState.df)
                Checkpoint.save_graph_parquet(AppState.graph)
            else:
                AppState.graph = nx.MultiDiGraph()
        except Exception as e:
            print(f"  [Phase 2] ✗ Graph build failed: {e}")
            AppState.startup_status['errors'].append(f"Graph: {e}")
            AppState.graph = nx.MultiDiGraph()

        # Graph topology is usable now — Investigation's subgraph endpoints and
        # Network Intelligence's radius-3 traversal only need AppState.graph
        # itself, not the GNN/hybrid scores computed below. Flip the flag here
        # (not after scoring finishes) so both pages unblock as soon as they
        # actually can, instead of waiting on potentially much-slower ML inference.
        AppState.graph_ready = True
        print("═══════════════════════════════════════════════════")
        print("  ✓ Graph ready!  Investigation & subgraph enabled.")
        print("═══════════════════════════════════════════════════")

        # GNN scoring now that graph exists
        try:
            if AppState.graph and hasattr(AppState.full_features, 'columns') \
                    and 'account' in AppState.full_features.columns:
                gnn_scores = predict_gnn_score(AppState.graph, AppState.full_features)
                AppState.full_features['gnn_fraud_score'] = \
                    AppState.full_features['account'].map(gnn_scores).fillna(0.0)
                
                cfg = get_config()['ml']
                if cfg.get('hybrid_enabled', False):
                    from src.ml.hybrid_predictor import hybrid_predictor
                    account_list = AppState.full_features['account'].tolist()
                    print("  [Phase 2] Executing Hybrid Stack inference...")
                    hybrid_scores = hybrid_predictor.predict(AppState.graph, account_list)
                    AppState.full_features['hybrid_score'] = AppState.full_features['account'].map(hybrid_scores).fillna(0.0)
                else:
                    AppState.full_features['hybrid_score'] = 0.0

                # Recompute risk_tier by percentile rank over the BEST available
                # per-account score (hybrid > gnn > base XGBoost probability),
                # same preference order entity_graph_service.compute_entity_risk
                # uses when picking an entity's risk-driver account.
                effective_score = AppState.full_features['fraud_probability'].copy()
                gnn_mask = AppState.full_features['gnn_fraud_score'] > 0
                effective_score[gnn_mask] = AppState.full_features.loc[gnn_mask, 'gnn_fraud_score']
                hybrid_mask = AppState.full_features['hybrid_score'] > 0
                effective_score[hybrid_mask] = AppState.full_features.loc[hybrid_mask, 'hybrid_score']
                
                # Update fraud_probability with GNN/hybrid-refined probability
                AppState.full_features['fraud_probability'] = effective_score

                # Recompute risk_score as percentile rank over effective_score to de-compress it
                from scipy.stats import rankdata as _rankdata
                _ranks = _rankdata(effective_score.values, method='average')
                _n = len(effective_score)
                AppState.full_features['risk_score'] = (_ranks / _n * 99).round().astype(int)

                tier_scores = dict(zip(AppState.full_features['account'], effective_score))
                tiers = compute_percentile_tiers(tier_scores)
                AppState.full_features['risk_tier'] = AppState.full_features['account'].map(tiers).fillna('LOW')

                # Refresh lookup cache with updated gnn scores
                AppState.features_by_account = AppState.full_features.set_index('account').to_dict('index')
                # Re-calibrate all alerts with the newly computed GNN structural risk scores
                enrich_alerts_with_ml()
                # Entity/bank risk improves once GNN/hybrid scores are in — recompute.
                if AppState.entities_by_id:
                    _risk_summary = attach_risk_aggregates(
                        AppState.entities_by_id, AppState.banks_by_id,
                        AppState.bank_profiles_cache, AppState.features_by_account,
                    )
                    AppState.network_summary_cache["high_risk_entity_count"] = _risk_summary["high_risk_entity_count"]
                    AppState.clusters_by_id = build_clusters_index(AppState.entities_by_id)
        except Exception as e:
            print(f"  [Phase 2] ✗ GNN scoring failed (non-fatal): {e}")

    except Exception as e:
        print(f"  [Phase 2] ✗✗✗ Graph init error: {e}")
        AppState.startup_status['errors'].append(f"Phase2: {e}")


# ── Status API (always available, even during loading) ───────
@app.get("/api/status")
async def get_status():
    return {
        "ready": AppState.startup_ready,
        "graph_ready": AppState.graph_ready,
        "current_step": AppState.startup_status['current_step'],
        "steps_done": AppState.startup_status['steps_done'],
        "total_steps": AppState.startup_status['total_steps'],
        "errors": AppState.startup_status['errors']
    }


@app.get("/api/config")
async def get_frontend_config():
    """Expose safe public config to the frontend (anon key, Supabase URL)."""
    s = get_case_settings()
    return {
        "supabase_url":      s['supabase_url'],
        "supabase_anon_key": s['supabase_anon_key'],
        "auth_enabled":      bool(s['supabase_url'] and s['supabase_anon_key']),
        "storage_enabled":   bool(s['supabase_url'] and s['supabase_service_role_key']),
        "db_backend":        s['database_backend'],
    }


# ── Include Routers ─────────────────────────────────────────
from api.routes import overview, alerts, investigation, graph_api, report, model, insights, upload, network  # noqa: E402
from api.routes import cases as cases_routes  # noqa: E402

app.include_router(overview.router, prefix="/api")
app.include_router(alerts.router, prefix="/api")
app.include_router(investigation.router, prefix="/api")
app.include_router(graph_api.router, prefix="/api")
app.include_router(report.router, prefix="/api")
app.include_router(model.router, prefix="/api")
app.include_router(insights.router, prefix="/api")
app.include_router(upload.router, prefix="/api")
app.include_router(network.router, prefix="/api")
app.include_router(cases_routes.router, prefix="/api")

# ── Copilot Router (optional, requires API keys) ──────────────
try:
    from api.routes import copilot as copilot_routes  # noqa: E402
    app.include_router(copilot_routes.router)
    print("  [OK] Copilot routes registered")
except Exception as e:
    print(f"  [WARN] Copilot routes unavailable: {e}")
    print("  [INFO] Run 'python setup_copilot.py' to enable copilot")

# Make AppState available to copilot
@app.middleware("http")
async def add_app_state_to_request(request: Request, call_next):
    """Make AppState available to copilot routes"""
    request.app.state.app_state = AppState
    return await call_next(request)


# ── Serve Static Frontend ───────────────────────────────────
SERVE_FRONTEND = os.environ.get("SERVE_FRONTEND", "true").lower() == "true"

if SERVE_FRONTEND:
    frontend_dir = os.path.join(os.path.dirname(__file__), 'frontend')
    os.makedirs(frontend_dir, exist_ok=True)
    app.mount("/static", StaticFiles(directory=frontend_dir), name="static")
    app.mount("/js", StaticFiles(directory=os.path.join(frontend_dir, "js")), name="frontend-js")
    app.mount("/pages", StaticFiles(directory=os.path.join(frontend_dir, "pages")), name="frontend-pages")

    @app.get("/")
    async def serve_landing():
        return FileResponse(os.path.join(frontend_dir, "pages", "landing.html"))

    @app.get("/app")
    @app.get("/overview")
    @app.get("/alerts")
    @app.get("/investigation")
    @app.get("/model")
    @app.get("/cases")
    @app.get("/login")
    async def serve_spa():
        return FileResponse(os.path.join(frontend_dir, "index.html"))
else:
    @app.get("/")
    async def root():
        return {
            "status": "active",
            "message": "Fund Flow Tracker Backend is Running",
            "api_docs": "/docs",
            "frontend": "disabled_on_this_host"
        }

if __name__ == "__main__":
    import os
    cfg = get_config()
    # Use PORT env var if available (standard for HF/Heroku/etc)
    port = int(os.environ.get("PORT", cfg['server']['port']))
    uvicorn.run("server:app", host=cfg['server']['host'], port=port, reload=False) # Disable reload in prod