"""xQTL Browser REST API — FastAPI + DuckDB over Phase 0 output (zero-copy).

Reads tensorQTL parquet output directly, with DuckDB SQL views performing all
transforms (column renames, variant_id parsing, gene annotation joins, FDR).

Usage:
    cd api && uvicorn main:app --reload --port 8000
"""
import json
from pathlib import Path
from typing import Optional

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

app = FastAPI(
    title="xQTL Browser API",
    version="0.3.0",
    description="Multiomics xQTL browser for neurodegenerative disease cohorts",
)

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

import os
PROJECT_ROOT = Path(os.environ.get("XQTL_DATA_DIR", Path(__file__).parent.parent))
REGISTRY_PATH = PROJECT_ROOT / "dataset_registry.json"
RSID_LOOKUP_PATH = PROJECT_ROOT / "references" / "rsid_lookup.parquet"

# Allowlist of valid layer names (prevents SQL injection via layer parameter)
VALID_LAYERS: set[str] = set()
LAYER_STATS: dict[str, dict] = {}


def load_registry() -> dict:
    with open(REGISTRY_PATH) as f:
        return json.load(f)


def init_db() -> duckdb.DuckDBPyConnection:
    """Create a persistent DuckDB connection with views over raw Phase 0 output."""
    db = duckdb.connect()
    registry = load_registry()

    # Load rsID lookup if available
    has_rsid = RSID_LOOKUP_PATH.exists()
    if has_rsid:
        db.execute(f"""
            CREATE TABLE rsid_lookup AS
            SELECT * FROM read_parquet('{RSID_LOOKUP_PATH}')
        """)

    # Collect SELECT statements per layer (multiple datasets may share a layer)
    layer_selects: dict[str, list[str]] = {}

    for ds_id, ds in registry["datasets"].items():
        data_dir = PROJECT_ROOT / ds["data_dir"]
        nom_dir = data_dir / "qtl" / "cis_nominal"
        gene_ann_path = data_dir / "gene_annotations.parquet"
        perm_fdr_path = data_dir / "permutation_fdr.parquet"

        if not nom_dir.exists() or not any(nom_dir.glob("*.parquet")):
            continue

        # Load small metadata tables
        if gene_ann_path.exists():
            db.execute(f"""
                CREATE TABLE IF NOT EXISTS gene_ann_{ds_id} AS
                SELECT * FROM read_parquet('{gene_ann_path}')
            """)
        if perm_fdr_path.exists():
            db.execute(f"""
                CREATE TABLE IF NOT EXISTS perm_fdr_{ds_id} AS
                SELECT * FROM read_parquet('{perm_fdr_path}')
            """)

        omics_layer = ds["omics_layer"]
        n_samples = ds["n_samples"]
        tissue = ds["tissue"]
        cohort = ds["cohort"]
        layer_name = ds["qtl_type"]  # e.g., "eqtl"

        rsid_join = ""
        rsid_col = "NULL AS rsid,"
        if has_rsid:
            rsid_join = "LEFT JOIN rsid_lookup r ON q.variant_id = r.variant_id"
            rsid_col = "r.rsid,"

        gene_join = ""
        gene_cols = "NULL AS feature_name, NULL AS feature_chr, CAST(NULL AS INTEGER) AS feature_start, CAST(NULL AS INTEGER) AS feature_end,"
        if gene_ann_path.exists():
            gene_join = f"LEFT JOIN gene_ann_{ds_id} g ON q.phenotype_id = g.feature_id"
            gene_cols = "g.feature_name, g.feature_chr, CAST(g.feature_start AS INTEGER) AS feature_start, CAST(g.feature_end AS INTEGER) AS feature_end,"

        fdr_join = ""
        fdr_cols = "NULL AS fdr, NULL AS is_egene, NULL AS pval_beta,"
        if perm_fdr_path.exists():
            fdr_join = f"LEFT JOIN perm_fdr_{ds_id} p ON q.phenotype_id = p.phenotype_id"
            fdr_cols = "p.qval AS fdr, (p.qval < 0.05) AS is_egene, p.pval_beta,"

        nom_glob = str(nom_dir / "*.parquet")
        select_sql = f"""
            SELECT
                q.variant_id,
                {rsid_col}
                split_part(q.variant_id, ':', 1) AS chr,
                CAST(split_part(q.variant_id, ':', 2) AS INTEGER) AS pos,
                split_part(q.variant_id, ':', 3) AS ref,
                split_part(q.variant_id, ':', 4) AS alt,
                LEAST(q.af, 1.0 - q.af) AS maf,
                q.phenotype_id AS feature_id,
                {gene_cols}
                '{omics_layer}' AS omics_layer,
                q.slope AS beta,
                q.slope_se AS se,
                q.pval_nominal AS pvalue,
                {fdr_cols}
                {n_samples} AS n_samples,
                '{tissue}' AS tissue,
                '{cohort}' AS cohort,
                '{layer_name}' AS layer_key
            FROM read_parquet('{nom_glob}', filename=true) q
            {gene_join}
            {fdr_join}
            {rsid_join}
        """
        layer_selects.setdefault(layer_name, []).append(select_sql)

    # Create one view per layer (UNION ALL if multiple datasets)
    for layer_name, selects in layer_selects.items():
        union_sql = " UNION ALL ".join(selects)
        db.execute(f"CREATE VIEW {layer_name} AS {union_sql}")
        VALID_LAYERS.add(layer_name)

    # Cross-layer view: UNION ALL of all per-layer views
    if VALID_LAYERS:
        all_union = " UNION ALL ".join(
            f"SELECT * FROM {ln}" for ln in sorted(VALID_LAYERS)
        )
        db.execute(f"CREATE VIEW all_layers AS {all_union}")

    # Build precomputed gene summary tables per layer from metadata parquets
    # (avoids scanning millions of nominal rows for /genes and /summary)
    gene_summary_selects: dict[str, list[str]] = {}
    for ds_id, ds in registry["datasets"].items():
        data_dir = PROJECT_ROOT / ds["data_dir"]
        gene_ann_path = data_dir / "gene_annotations.parquet"
        perm_fdr_path = data_dir / "permutation_fdr.parquet"
        layer_name = ds["qtl_type"]

        if layer_name not in VALID_LAYERS:
            continue
        if not gene_ann_path.exists() or not perm_fdr_path.exists():
            continue

        omics_layer = ds["omics_layer"]
        select_sql = f"""
            SELECT
                g.feature_id,
                g.feature_name,
                g.feature_chr,
                CAST(g.feature_start AS INTEGER) AS feature_start,
                CAST(g.feature_end AS INTEGER) AS feature_end,
                COALESCE(p.num_var, 0) AS n_variants,
                p.pval_beta,
                p.pval_nominal_threshold AS min_pvalue,
                p.qval AS fdr,
                (p.qval < 0.05) AS is_egene,
                '{omics_layer}' AS omics_layer,
                '{layer_name}' AS layer_key
            FROM gene_ann_{ds_id} g
            JOIN perm_fdr_{ds_id} p ON g.feature_id = p.phenotype_id
        """
        gene_summary_selects.setdefault(layer_name, []).append(select_sql)

    for layer_name, selects in gene_summary_selects.items():
        union_sql = " UNION ALL ".join(selects)
        db.execute(f"CREATE TABLE gene_summary_{layer_name} AS {union_sql}")

    # Cross-layer gene summary table
    if VALID_LAYERS:
        all_gs = " UNION ALL ".join(
            f"SELECT * FROM gene_summary_{ln}" for ln in sorted(VALID_LAYERS)
        )
        db.execute(f"CREATE TABLE gene_summary_all AS {all_gs}")

    # Precompute layer-level summary stats (avoids scanning nominal views on /summary)
    for layer_name in VALID_LAYERS:
        row = db.execute(f"""
            SELECT COUNT(*) AS n_pairs, COUNT(DISTINCT variant_id) AS n_variants
            FROM {layer_name}
        """).fetchone()
        gene_row = db.execute(f"""
            SELECT COUNT(*), COUNT(*) FILTER (WHERE is_egene), ANY_VALUE(omics_layer)
            FROM gene_summary_{layer_name}
        """).fetchone()
        LAYER_STATS[layer_name] = {
            "n_pairs": row[0], "n_variants": row[1],
            "n_genes": gene_row[0], "n_egenes": gene_row[1], "omics_layer": gene_row[2],
        }

    # Aggregate "all" stats
    if VALID_LAYERS:
        LAYER_STATS["all"] = {
            "n_pairs": sum(s["n_pairs"] for s in LAYER_STATS.values()),
            "n_variants": sum(s["n_variants"] for s in LAYER_STATS.values()),
            "n_genes": sum(s["n_genes"] for s in LAYER_STATS.values()),
            "n_egenes": sum(s["n_egenes"] for s in LAYER_STATS.values()),
            "omics_layer": "All Layers",
        }

    return db


# Global connection — created once at startup
_db: duckdb.DuckDBPyConnection | None = None


def get_db() -> duckdb.DuckDBPyConnection:
    global _db
    if _db is None:
        _db = init_db()
    return _db


@app.on_event("startup")
def startup():
    get_db()


# --- Models ---

QTL_COLS = [
    "variant_id", "rsid", "chr", "pos", "ref", "alt", "maf", "feature_id",
    "feature_name", "feature_chr", "feature_start", "feature_end",
    "omics_layer", "beta", "se", "pvalue", "fdr", "is_egene", "pval_beta",
    "n_samples", "tissue", "cohort", "layer_key",
]


class QTLResult(BaseModel):
    variant_id: str
    rsid: str | None = None
    chr: str
    pos: int
    ref: str
    alt: str
    maf: float
    feature_id: str
    feature_name: str | None = None
    feature_chr: str | None = None
    feature_start: int | None = None
    feature_end: int | None = None
    omics_layer: str
    beta: float
    se: float
    pvalue: float
    fdr: float | None = None
    is_egene: bool | None = None
    pval_beta: float | None = None
    n_samples: int
    tissue: str
    cohort: str
    layer_key: str | None = None


class GeneSummary(BaseModel):
    feature_id: str
    feature_name: str | None = None
    feature_chr: str | None = None
    feature_start: int | None = None
    feature_end: int | None = None
    n_variants: int
    min_pvalue: float
    pval_beta: float | None = None
    fdr: float | None = None
    is_egene: bool | None = None
    omics_layer: str
    layer_key: str | None = None


class SummaryStats(BaseModel):
    layer_key: str
    omics_layer: str
    n_genes: int
    n_egenes: int
    n_variants: int
    n_pairs: int


def _validate_layer(layer: str) -> str:
    if layer != "all" and layer not in VALID_LAYERS:
        raise HTTPException(404, f"Unknown layer '{layer}'. Available: {['all'] + sorted(VALID_LAYERS)}")
    return layer


def _layer_view(layer: str) -> str:
    """Return the DuckDB view name for nominal QTL data."""
    return "all_layers" if layer == "all" else layer


def _gene_summary_table(layer: str) -> str:
    """Return the DuckDB table name for gene summaries."""
    return "gene_summary_all" if layer == "all" else f"gene_summary_{layer}"


# --- Endpoints ---

@app.get("/api/v1/layers", response_model=list[str])
def list_layers():
    """List available omics layers."""
    return sorted(VALID_LAYERS)


@app.get("/api/v1/summary", response_model=list[SummaryStats])
def summary():
    """Summary statistics per omics layer (precomputed at startup)."""
    get_db()  # ensure init
    return [
        SummaryStats(layer_key=layer, **LAYER_STATS[layer])
        for layer in ["all"] + sorted(VALID_LAYERS)
    ]


@app.get("/api/v1/genes", response_model=list[GeneSummary])
def list_genes(
    layer: str = Query("all", description="Omics layer (or 'all' for cross-layer)"),
    search: Optional[str] = Query(None, description="Search gene name (prefix match)"),
    chr: Optional[str] = Query(None, description="Filter by chromosome"),
    egenes_only: bool = Query(False, description="Only return significant eGenes"),
    limit: int = Query(500, le=2000),
    offset: int = Query(0),
):
    """List genes with QTL summary statistics."""
    _validate_layer(layer)
    db = get_db()
    table = _gene_summary_table(layer)

    conditions = []
    params = []
    if search:
        conditions.append(f"(feature_name ILIKE ${len(params) + 1} OR feature_id ILIKE ${len(params) + 1})")
        params.append(f"{search}%")
    if chr:
        conditions.append(f"feature_chr = ${len(params) + 1}")
        params.append(chr)
    if egenes_only:
        conditions.append("is_egene = true")
    where = "WHERE " + " AND ".join(conditions) if conditions else ""

    # When layer=all and no filters, sample proportionally from each layer
    # so all layers are represented in the Manhattan plot
    if layer == "all" and not search and not chr and not egenes_only:
        per_layer_limit = max(limit // len(VALID_LAYERS), 1) if VALID_LAYERS else limit
        per_layer_selects = []
        for ln in sorted(VALID_LAYERS):
            per_layer_selects.append(f"""
                (SELECT feature_id, feature_name, feature_chr, feature_start, feature_end,
                        n_variants, min_pvalue, pval_beta, fdr, is_egene, omics_layer, layer_key
                 FROM gene_summary_{ln}
                 ORDER BY pval_beta NULLS LAST
                 LIMIT {per_layer_limit} OFFSET {offset})
            """)
        union_sql = " UNION ALL ".join(per_layer_selects)
        rows = db.execute(f"""
            SELECT * FROM ({union_sql}) sub
            ORDER BY pval_beta NULLS LAST
            LIMIT {limit}
        """).fetchall()
    else:
        rows = db.execute(f"""
            SELECT feature_id, feature_name, feature_chr, feature_start, feature_end,
                   n_variants, min_pvalue,
                   pval_beta, fdr, is_egene, omics_layer, layer_key
            FROM {table}
            {where}
            ORDER BY pval_beta NULLS LAST
            LIMIT {limit} OFFSET {offset}
        """, params).fetchall()

    return [GeneSummary(
        feature_id=r[0], feature_name=r[1], feature_chr=r[2],
        feature_start=r[3], feature_end=r[4], n_variants=r[5],
        min_pvalue=r[6], pval_beta=r[7], fdr=r[8], is_egene=r[9],
        omics_layer=r[10], layer_key=r[11],
    ) for r in rows]


@app.get("/api/v1/genes/{gene}/eqtls", response_model=list[QTLResult])
def get_gene_eqtls(
    gene: str,
    layer: str = Query("all", description="Omics layer (or 'all' for cross-layer)"),
    pvalue_threshold: float = Query(1.0, description="Max p-value"),
    limit: int = Query(100, le=5000),
):
    """Get QTL results for a specific gene (by name or Ensembl ID)."""
    _validate_layer(layer)
    db = get_db()
    view = _layer_view(layer)

    if gene.startswith("ENSG"):
        condition = "feature_id = $1"
    else:
        condition = "feature_name ILIKE $1"

    rows = db.execute(f"""
        SELECT {', '.join(QTL_COLS)}
        FROM {view}
        WHERE {condition} AND pvalue <= $2
        ORDER BY pvalue
        LIMIT {limit}
    """, [gene, pvalue_threshold]).fetchall()

    if not rows:
        raise HTTPException(404, f"No results for gene '{gene}' in {layer}")

    return [QTLResult(**dict(zip(QTL_COLS, r))) for r in rows]


@app.get("/api/v1/variants/{variant_id:path}", response_model=list[QTLResult])
def get_variant(
    variant_id: str,
    layer: str = Query("all", description="Omics layer (or 'all' for cross-layer)"),
    limit: int = Query(100, le=5000),
):
    """Get all QTL associations for a specific variant (by CHR:POS:REF:ALT or rsID)."""
    _validate_layer(layer)
    db = get_db()
    view = _layer_view(layer)

    # Resolve rsID to variant_id if needed
    if variant_id.startswith("rs"):
        resolved = db.execute(
            "SELECT variant_id FROM rsid_lookup WHERE rsid = $1 LIMIT 1",
            [variant_id],
        ).fetchone()
        if not resolved:
            raise HTTPException(404, f"rsID '{variant_id}' not found in lookup")
        variant_id = resolved[0]

    rows = db.execute(f"""
        SELECT {', '.join(QTL_COLS)}
        FROM {view}
        WHERE variant_id = $1
        ORDER BY pvalue
        LIMIT {limit}
    """, [variant_id]).fetchall()

    if not rows:
        raise HTTPException(404, f"No results for variant '{variant_id}'")

    return [QTLResult(**dict(zip(QTL_COLS, r))) for r in rows]


@app.get("/api/v1/region", response_model=list[QTLResult])
def query_region(
    chr: str = Query(..., description="Chromosome"),
    start: int = Query(..., description="Start position"),
    end: int = Query(..., description="End position"),
    layer: str = Query("all", description="Omics layer (or 'all' for cross-layer)"),
    pvalue_threshold: float = Query(0.05, description="Max p-value"),
    limit: int = Query(500, le=5000),
):
    """Get QTL results in a genomic region."""
    _validate_layer(layer)
    db = get_db()
    view = _layer_view(layer)

    rows = db.execute(f"""
        SELECT {', '.join(QTL_COLS)}
        FROM {view}
        WHERE chr = $1 AND pos BETWEEN $2 AND $3
              AND pvalue <= $4
        ORDER BY pvalue
        LIMIT {limit}
    """, [chr, start, end, pvalue_threshold]).fetchall()

    return [QTLResult(**dict(zip(QTL_COLS, r))) for r in rows]