"""
Blood-Brain Omics Benchmark API

FastAPI + DuckDB backend serving benchmark results.
Queries Parquet files remotely via DuckDB httpfs — no data downloaded to disk.
DuckDB reads only the needed columns/row groups per query via HTTP range requests.

Data source modes (in priority order):
  1. Remote: HF Dataset repo via httpfs (default, zero RAM footprint)
  2. Local: files in BENCHMARK_DATA_DIR (for local dev)

Endpoints:
    GET /api/v1/registry              - Full registry metadata
    GET /api/v1/maxn/heatmap          - Blood x brain heatmap data
    GET /api/v1/maxn/detail           - Per-module results for one combination
    GET /api/v1/maxn/panel_summary    - One panel across all outcomes (mean/min/max)
    GET /api/v1/maxn/outcome_summary  - One outcome across all panels (mean/min/max)
    GET /api/v1/maxn/target_comparison - Per-panel performance for a specific target
    GET /api/v1/h2h/blood             - Blood H2H comparison
    GET /api/v1/h2h/blood/summary     - H2H win counts
    GET /api/v1/h2h/brain             - Brain H2H comparison
    GET /api/v1/h2h/models            - Model H2H comparison (future)
    GET /api/v1/temporal              - Temporal decay results
    GET /api/v1/features              - Feature importance for a target
    GET /api/v1/features/cross        - Cross-target feature importance
"""

import json
import os
import urllib.request
from typing import Optional

import duckdb
from fastapi import FastAPI, HTTPException, Query
from fastapi.middleware.cors import CORSMiddleware

# =============================================================================
# Configuration
# =============================================================================

# HF Dataset repo for remote Parquet queries
HF_DATASET = os.environ.get(
    "HF_DATASET", "stasaking/blood-brain-benchmark"
)
HF_BASE_URL = f"https://huggingface.co/datasets/{HF_DATASET}/resolve/main"

# Local fallback for dev
DATA_DIR = os.environ.get(
    "BENCHMARK_DATA_DIR",
    os.path.join(os.path.dirname(__file__), "..", "data")
)

app = FastAPI(
    title="Blood-Brain Omics Benchmark API",
    version="1.0.0",
    description="Interactive exploration of blood omics → brain phenotype predictions",
)

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

# =============================================================================
# Startup: configure DuckDB with httpfs views or local tables
# =============================================================================

db = None
registry = None


def _resolve_parquet(filename: str) -> str:
    """Return remote URL or local path for a Parquet file."""
    local_path = os.path.join(DATA_DIR, filename)
    if os.path.exists(local_path):
        return local_path
    return f"{HF_BASE_URL}/{filename}"


@app.on_event("startup")
def startup():
    global db, registry

    # Initialize DuckDB with httpfs
    db = duckdb.connect(":memory:")
    db.execute("INSTALL httpfs; LOAD httpfs;")

    # Resolve data sources (remote or local)
    results_src = _resolve_parquet("benchmark_results.parquet")
    features_src = _resolve_parquet("feature_importance.parquet")

    is_remote = results_src.startswith("http")
    mode = "remote httpfs" if is_remote else "local files"
    print(f"Data mode: {mode}")

    # Create views that query Parquet directly (no data copied into RAM)
    db.execute(f"""
        CREATE VIEW results AS
        SELECT * FROM read_parquet('{results_src}')
    """)
    db.execute(f"""
        CREATE VIEW features AS
        SELECT * FROM read_parquet('{features_src}')
    """)

    n_results = db.execute("SELECT COUNT(*) FROM results").fetchone()[0]
    n_features = db.execute("SELECT COUNT(*) FROM features").fetchone()[0]
    print(f"Views created: {n_results} results, {n_features} features")

    # Load registry JSON (small file, always downloaded)
    registry_local = os.path.join(DATA_DIR, "benchmark_registry.json")
    if os.path.exists(registry_local):
        with open(registry_local) as f:
            registry = json.load(f)
    else:
        # Download from HF
        registry_url = f"{HF_BASE_URL}/benchmark_registry.json"
        try:
            with urllib.request.urlopen(registry_url) as resp:
                registry = json.loads(resp.read().decode())
            print(f"Downloaded registry from {registry_url}")
        except Exception as e:
            print(f"WARNING: Could not load registry: {e}")
            registry = {"project": {}, "blood_platforms": {},
                         "brain_targets": {}, "phases": {}, "models": {}}


# =============================================================================
# Helpers
# =============================================================================

VALID_METRICS = {"r2", "pearson", "mse", "mae"}


def validate_metric(metric: str) -> str:
    if metric not in VALID_METRICS:
        raise HTTPException(400, f"Invalid metric: {metric}. "
                            f"Must be one of {VALID_METRICS}")
    return metric


def covariate_phase(base_phase: str, covariates: str) -> str:
    """Map covariates param to actual phase name."""
    if covariates == "none":
        return base_phase
    elif covariates == "only":
        return f"{base_phase}_covonly"
    elif covariates == "included":
        return f"{base_phase}_withcov"
    return base_phase


# =============================================================================
# Endpoints
# =============================================================================

@app.get("/api/v1/registry")
def get_registry():
    """Full registry metadata."""
    return registry


@app.get("/api/v1/maxn/heatmap")
def maxn_heatmap(
    metric: str = Query("r2", description="Metric to aggregate"),
    covariates: str = Query("none", enum=["none", "only", "included"]),
    model: str = Query("TabPFN"),
):
    """
    Blood x brain heatmap data.
    Returns median metric across modules for each combination.
    """
    validate_metric(metric)
    phase = covariate_phase("maxn", covariates)

    rows = db.execute(f"""
        SELECT blood_platform, brain_target,
               MEDIAN({metric}) as value,
               MEDIAN(n_samples) as n_samples,
               COUNT(*) as n_targets
        FROM results
        WHERE phase = ? AND model = ?
        GROUP BY blood_platform, brain_target
        ORDER BY blood_platform, brain_target
    """, [phase, model]).fetchall()

    if not rows:
        return {"rows": [], "cols": [], "values": [], "n_samples": [],
                "metric": metric, "covariates": covariates}

    # Build matrix
    blood_set = sorted(set(r[0] for r in rows))
    brain_set = sorted(set(r[1] for r in rows))

    values = [[None] * len(brain_set) for _ in range(len(blood_set))]
    n_samples = [[None] * len(brain_set) for _ in range(len(blood_set))]

    blood_idx = {b: i for i, b in enumerate(blood_set)}
    brain_idx = {b: i for i, b in enumerate(brain_set)}

    for blood, brain, val, ns, nt in rows:
        i = blood_idx[blood]
        j = brain_idx[brain]
        values[i][j] = round(val, 4) if val is not None else None
        n_samples[i][j] = int(ns) if ns is not None else None

    return {
        "rows": blood_set,
        "cols": brain_set,
        "values": values,
        "n_samples": n_samples,
        "metric": metric,
        "covariates": covariates,
        "model": model,
    }


@app.get("/api/v1/maxn/detail")
def maxn_detail(
    blood: str = Query(..., description="Blood platform name"),
    brain: str = Query(..., description="Brain target name"),
    covariates: str = Query("none", enum=["none", "only", "included"]),
    model: str = Query("TabPFN"),
):
    """Per-module results for one blood x brain combination."""
    phase = covariate_phase("maxn", covariates)

    rows = db.execute("""
        SELECT target, n_samples, r2, pearson, mse, mae
        FROM results
        WHERE phase = ? AND blood_platform = ? AND brain_target = ?
              AND model = ?
        ORDER BY r2 DESC
    """, [phase, blood, brain, model]).fetchall()

    return {
        "blood": blood,
        "brain": brain,
        "covariates": covariates,
        "model": model,
        "targets": [
            {"target": r[0], "n_samples": r[1], "r2": r[2],
             "pearson": r[3], "mse": r[4], "mae": r[5]}
            for r in rows
        ],
    }


@app.get("/api/v1/h2h/blood")
def h2h_blood(
    pair: str = Query(..., description="Platform pair, e.g. SomaScan_vs_TMT"),
    brain: str = Query(..., description="Brain target name"),
    metric: str = Query("r2"),
    covariates: str = Query("none", enum=["none", "included"]),
    model: str = Query("TabPFN"),
):
    """Per-module comparison for a blood platform pair on one brain target."""
    validate_metric(metric)
    phase = "h2h_blood_withcov" if covariates == "included" else "h2h_blood"

    rows = db.execute(f"""
        SELECT target, blood_platform, {metric}
        FROM results
        WHERE phase = ? AND h2h_pair = ? AND brain_target = ? AND model = ?
        ORDER BY target
    """, [phase, pair, brain, model]).fetchall()

    # Pivot: target -> {platform_a: val, platform_b: val}
    platforms = sorted(set(r[1] for r in rows))
    targets = {}
    for target, platform, val in rows:
        if target not in targets:
            targets[target] = {}
        targets[target][platform] = round(val, 4) if val is not None else None

    return {
        "pair": pair,
        "brain": brain,
        "platforms": platforms,
        "metric": metric,
        "targets": [
            {"target": t, **vals}
            for t, vals in sorted(targets.items())
        ],
    }


@app.get("/api/v1/h2h/blood/summary")
def h2h_blood_summary(
    metric: str = Query("r2"),
    covariates: str = Query("none", enum=["none", "included"]),
    model: str = Query("TabPFN"),
):
    """Win counts across all blood H2H pairs."""
    validate_metric(metric)
    phase = "h2h_blood_withcov" if covariates == "included" else "h2h_blood"

    rows = db.execute(f"""
        SELECT h2h_pair, brain_target, target, blood_platform, {metric}
        FROM results
        WHERE phase = ? AND model = ? AND h2h_pair IS NOT NULL
        ORDER BY h2h_pair, brain_target, target
    """, [phase, model]).fetchall()

    # Count wins per pair
    pair_wins = {}
    current = None
    buffer = {}

    for pair, brain, target, platform, val in rows:
        key = (pair, brain, target)
        if key != current:
            if current and len(buffer) == 2:
                pair_key = current[0]
                if pair_key not in pair_wins:
                    pair_wins[pair_key] = {}
                platforms = list(buffer.keys())
                v0, v1 = buffer[platforms[0]], buffer[platforms[1]]
                if v0 is not None and v1 is not None:
                    asc = metric in ("mse", "mae")
                    winner = platforms[0] if (v0 < v1 if asc else v0 > v1) else platforms[1]
                    pair_wins[pair_key][winner] = pair_wins[pair_key].get(winner, 0) + 1
            current = key
            buffer = {}
        buffer[platform] = val

    # Process last group
    if current and len(buffer) == 2:
        pair_key = current[0]
        if pair_key not in pair_wins:
            pair_wins[pair_key] = {}
        platforms = list(buffer.keys())
        v0, v1 = buffer[platforms[0]], buffer[platforms[1]]
        if v0 is not None and v1 is not None:
            asc = metric in ("mse", "mae")
            winner = platforms[0] if (v0 < v1 if asc else v0 > v1) else platforms[1]
            pair_wins[pair_key][winner] = pair_wins[pair_key].get(winner, 0) + 1

    return {
        "metric": metric,
        "covariates": covariates,
        "pairs": pair_wins,
    }


@app.get("/api/v1/h2h/brain")
def h2h_brain(
    blood: str = Query(..., description="Blood platform name"),
    metric: str = Query("r2"),
    covariates: str = Query("none", enum=["none", "included"]),
    model: str = Query("TabPFN"),
):
    """Compare brain targets for one blood platform."""
    validate_metric(metric)
    phase = "h2h_brain_withcov" if covariates == "included" else "h2h_brain"

    rows = db.execute(f"""
        SELECT h2h_pair, brain_target, target, {metric}
        FROM results
        WHERE phase = ? AND blood_platform = ? AND model = ?
        ORDER BY h2h_pair, target
    """, [phase, blood, model]).fetchall()

    # Group by pair
    pairs = {}
    for pair, brain, target, val in rows:
        if pair not in pairs:
            pairs[pair] = {}
        if target not in pairs[pair]:
            pairs[pair][target] = {}
        pairs[pair][target][brain] = round(val, 4) if val is not None else None

    return {
        "blood": blood,
        "metric": metric,
        "pairs": pairs,
    }


@app.get("/api/v1/h2h/models")
def h2h_models(
    blood: str = Query(...),
    brain: str = Query(...),
    metric: str = Query("r2"),
    covariates: str = Query("none", enum=["none", "included"]),
):
    """Compare models for one blood x brain combination (future)."""
    validate_metric(metric)
    phase = covariate_phase("maxn", covariates)

    rows = db.execute(f"""
        SELECT target, model, {metric}
        FROM results
        WHERE phase = ? AND blood_platform = ? AND brain_target = ?
        ORDER BY target, model
    """, [phase, blood, brain]).fetchall()

    models = sorted(set(r[1] for r in rows))
    targets = {}
    for target, model, val in rows:
        if target not in targets:
            targets[target] = {}
        targets[target][model] = round(val, 4) if val is not None else None

    return {
        "blood": blood,
        "brain": brain,
        "models": models,
        "metric": metric,
        "targets": [
            {"target": t, **vals}
            for t, vals in sorted(targets.items())
        ],
    }


@app.get("/api/v1/temporal")
def temporal(
    target: Optional[str] = Query(None, description="Specific target (e.g., gpath)"),
    metric: str = Query("r2"),
    covariates: str = Query("none", enum=["none", "included"]),
    model: str = Query("TabPFN"),
):
    """Temporal decay results: metric across time bins."""
    validate_metric(metric)
    phase = "temporal_withcov" if covariates == "included" else "temporal"

    if target:
        rows = db.execute(f"""
            SELECT temporal_bin, brain_target, target, {metric}, n_samples
            FROM results
            WHERE phase = ? AND model = ? AND target = ?
            ORDER BY temporal_bin
        """, [phase, model, target]).fetchall()
    else:
        rows = db.execute(f"""
            SELECT temporal_bin, brain_target, target, {metric}, n_samples
            FROM results
            WHERE phase = ? AND model = ?
            ORDER BY temporal_bin, target
        """, [phase, model]).fetchall()

    results = [
        {"bin": r[0], "brain_target": r[1], "target": r[2],
         "value": round(r[3], 4) if r[3] is not None else None,
         "n_samples": r[4]}
        for r in rows
    ]

    return {
        "metric": metric,
        "covariates": covariates,
        "results": results,
    }


@app.get("/api/v1/features")
def feature_importance(
    blood: str = Query(...),
    brain: str = Query(...),
    target: str = Query(...),
    phase: str = Query("maxn"),
    model: str = Query("TabPFN"),
    limit: int = Query(30, ge=1, le=100),
):
    """Top features for a specific module/target."""
    rows = db.execute("""
        SELECT feature_name, importance, rank
        FROM features
        WHERE phase = ? AND blood_platform = ? AND brain_target = ?
              AND target = ? AND model = ?
        ORDER BY rank
        LIMIT ?
    """, [phase, blood, brain, target, model, limit]).fetchall()

    return {
        "blood": blood,
        "brain": brain,
        "target": target,
        "phase": phase,
        "features": [
            {"feature": r[0], "importance": round(r[1], 4), "rank": r[2]}
            for r in rows
        ],
    }


@app.get("/api/v1/features/cross")
def feature_cross_target(
    blood: str = Query(...),
    brain: str = Query(...),
    phase: str = Query("maxn"),
    model: str = Query("TabPFN"),
    limit: int = Query(30, ge=1, le=100),
):
    """Aggregated feature importance across all modules in a brain target."""
    rows = db.execute("""
        SELECT feature_name,
               AVG(importance) as mean_importance,
               COUNT(DISTINCT target) as n_targets,
               MIN(rank) as best_rank
        FROM features
        WHERE phase = ? AND blood_platform = ? AND brain_target = ? AND model = ?
        GROUP BY feature_name
        ORDER BY mean_importance DESC
        LIMIT ?
    """, [phase, blood, brain, model, limit]).fetchall()

    return {
        "blood": blood,
        "brain": brain,
        "phase": phase,
        "features": [
            {"feature": r[0], "mean_importance": round(r[1], 4),
             "n_targets": r[2], "best_rank": r[3]}
            for r in rows
        ],
    }


@app.get("/api/v1/maxn/panel_summary")
def maxn_panel_summary(
    blood: str = Query(..., description="Blood platform name"),
    metric: str = Query("r2"),
    covariates: str = Query("none", enum=["none", "only", "included"]),
    model: str = Query("TabPFN"),
):
    """
    Summary of one predictor panel across all outcome panels.
    Returns mean, min, max, n_targets, and median n_samples per outcome.
    Replaces N calls to /maxn/detail with a single query.
    """
    validate_metric(metric)
    phase = covariate_phase("maxn", covariates)

    rows = db.execute(f"""
        SELECT brain_target,
               AVG({metric}) as mean_val,
               MIN({metric}) as min_val,
               MAX({metric}) as max_val,
               COUNT(*) as n_targets,
               MEDIAN(n_samples) as n_samples
        FROM results
        WHERE phase = ? AND blood_platform = ? AND model = ?
        GROUP BY brain_target
        ORDER BY mean_val DESC
    """, [phase, blood, model]).fetchall()

    return {
        "blood": blood,
        "metric": metric,
        "covariates": covariates,
        "outcomes": [
            {
                "brain_target": r[0],
                "mean": round(r[1], 4) if r[1] is not None else None,
                "min": round(r[2], 4) if r[2] is not None else None,
                "max": round(r[3], 4) if r[3] is not None else None,
                "n_targets": r[4],
                "n_samples": int(r[5]) if r[5] is not None else None,
            }
            for r in rows
        ],
    }


@app.get("/api/v1/maxn/outcome_summary")
def maxn_outcome_summary(
    brain: str = Query(..., description="Brain target name"),
    metric: str = Query("r2"),
    covariates: str = Query("none", enum=["none", "only", "included"]),
    model: str = Query("TabPFN"),
):
    """
    Summary of one outcome panel across all predictor panels.
    Returns mean, min, max, n_targets, and median n_samples per predictor.
    Replaces N calls to /maxn/detail with a single query.
    """
    validate_metric(metric)
    phase = covariate_phase("maxn", covariates)

    rows = db.execute(f"""
        SELECT blood_platform,
               AVG({metric}) as mean_val,
               MIN({metric}) as min_val,
               MAX({metric}) as max_val,
               COUNT(*) as n_targets,
               MEDIAN(n_samples) as n_samples
        FROM results
        WHERE phase = ? AND brain_target = ? AND model = ?
        GROUP BY blood_platform
        ORDER BY mean_val DESC
    """, [phase, brain, model]).fetchall()

    return {
        "brain": brain,
        "metric": metric,
        "covariates": covariates,
        "predictors": [
            {
                "blood_platform": r[0],
                "mean": round(r[1], 4) if r[1] is not None else None,
                "min": round(r[2], 4) if r[2] is not None else None,
                "max": round(r[3], 4) if r[3] is not None else None,
                "n_targets": r[4],
                "n_samples": int(r[5]) if r[5] is not None else None,
            }
            for r in rows
        ],
    }


@app.get("/api/v1/maxn/target_comparison")
def maxn_target_comparison(
    brain: str = Query(..., description="Brain target name"),
    target: str = Query(..., description="Specific target/module name"),
    metric: str = Query("r2"),
    covariates: str = Query("none", enum=["none", "only", "included"]),
    model: str = Query("TabPFN"),
):
    """
    Per-predictor performance for a specific target within an outcome panel.
    E.g., how each blood platform predicts module m4 of bulkrnaseq_dlpfc.
    Replaces N calls to /maxn/detail with a single query.
    """
    validate_metric(metric)
    phase = covariate_phase("maxn", covariates)

    rows = db.execute(f"""
        SELECT blood_platform, {metric}, n_samples
        FROM results
        WHERE phase = ? AND brain_target = ? AND target = ? AND model = ?
        ORDER BY {metric} DESC
    """, [phase, brain, target, model]).fetchall()

    return {
        "brain": brain,
        "target": target,
        "metric": metric,
        "covariates": covariates,
        "predictors": [
            {
                "blood_platform": r[0],
                "value": round(r[1], 4) if r[1] is not None else None,
                "n_samples": r[2],
            }
            for r in rows
        ],
    }


@app.get("/api/v1/health")
def health():
    """Health check."""
    n_results = db.execute("SELECT COUNT(*) FROM results").fetchone()[0]
    n_features = db.execute("SELECT COUNT(*) FROM features").fetchone()[0]
    return {"status": "ok", "n_results": n_results, "n_features": n_features}