"""
SQLite-based data loader for protein structure data.
Reads from MuSProt.db: edge table (pairwise comparisons from CSV)
and node table (chain annotations from TSV).
"""
import json
import pandas as pd
from pathlib import Path
from typing import Any, Dict, List, Optional, Set
import logging

from app.protein.config import get_database_path, get_summary_path
from app.protein.database import connect_readonly

logger = logging.getLogger(__name__)


def _assign_cluster(tm: float) -> str:
    if tm >= 0.99:
        return "cluster_ultra_high"
    if tm >= 0.95:
        return "cluster_very_high"
    if tm >= 0.90:
        return "cluster_high"
    if tm >= 0.80:
        return "cluster_medium"
    return "cluster_low"


class DataManager:
    """Manages SQL queries against MuSProt.db edge and node tables."""

    def __init__(self, db_path: Optional[Path] = None, **_kwargs):
        self.db_path = db_path or get_database_path()
        self._loaded = False
        self._total_records = 0
        self._summary: Dict[str, Any] = {}

    def load_data(self) -> None:
        """Verify DB is accessible and load optional precomputed metadata."""
        if self._loaded:
            return
        summary_path = get_summary_path()
        if summary_path:
            self._summary = json.loads(summary_path.read_text(encoding="utf-8"))
            self._total_records = int(self._summary.get("total_records", 0))

        conn = connect_readonly(self.db_path)
        try:
            cur = conn.cursor()
            cur.execute("SELECT 1 FROM edge LIMIT 1")
            if not self._total_records:
                logger.warning(
                    "No musprot_summary.json found; counting edge rows during startup. "
                    "Publish the summary sidecar to avoid this scan."
                )
                cur.execute("SELECT COUNT(*) FROM edge")
                self._total_records = cur.fetchone()[0]
        finally:
            conn.close()
        self._loaded = True
        logger.info(f"DataManager ready: {self._total_records:,} edge records in {self.db_path}")

    def _connect(self):
        if not self._loaded:
            raise RuntimeError("Call load_data() first.")
        return connect_readonly(self.db_path)

    def get_chain_sample(self, limit: int = 500) -> Set[str]:
        """Return a small chain sample without loading the full node index."""
        conn = self._connect()
        try:
            cur = conn.cursor()
            cur.execute(
                "SELECT pdb_id, auth_asym_id FROM node "
                "WHERE pdb_id IS NOT NULL AND auth_asym_id IS NOT NULL LIMIT ?",
                (limit,),
            )
            return {
                f"{str(pdb).upper()}_{str(chain).upper()}"
                for pdb, chain in cur.fetchall()
            }
        finally:
            conn.close()

    def get_filters_data(self) -> Dict[str, Any]:
        """Return aggregated ranges for filter UI."""
        if self._summary:
            return {
                "tm_min": self._summary["tm_score_range"]["min"],
                "tm_max": self._summary["tm_score_range"]["max"],
                "rmsd_min": self._summary["rmsd_range"]["min"],
                "rmsd_max": self._summary["rmsd_range"]["max"],
                "length_min": self._summary["length_range"]["min"],
                "length_max": self._summary["length_range"]["max"],
                "total_records": self._total_records,
            }

        conn = self._connect()
        try:
            cur = conn.cursor()
            cur.execute(
                "SELECT MIN(CAST(TM1 AS REAL)), MAX(CAST(TM1 AS REAL)),"
                "       MIN(CAST(RMSD AS REAL)), MAX(CAST(RMSD AS REAL))"
                " FROM edge"
            )
            tm_min, tm_max, rmsd_min, rmsd_max = cur.fetchone()
        finally:
            conn.close()

        from app.protein import tsv_loader
        lengths = []
        for row in tsv_loader._get_index().values():
            try:
                lengths.append(int(row["sequence_length"]))
            except (ValueError, TypeError):
                pass

        return {
            "tm_min": tm_min or 0.0,
            "tm_max": tm_max or 1.0,
            "rmsd_min": rmsd_min or 0.0,
            "rmsd_max": rmsd_max or 10.0,
            "length_min": min(lengths) if lengths else 0,
            "length_max": max(lengths) if lengths else 0,
            "total_records": self._total_records,
        }

    def query_edge(
        self,
        pdb_id_a: Optional[str] = None,
        auth_asym_id_a: Optional[str] = None,
        chain_ids: Optional[List[str]] = None,
        rmsd_min: Optional[float] = None,
        rmsd_max: Optional[float] = None,
        tm_min: Optional[float] = None,
        tm_max: Optional[float] = None,
        limit: int = 1000,
    ) -> pd.DataFrame:
        """Query edge table with filters. Returns DataFrame with cluster_id and lengths."""
        where_clauses: List[str] = []
        params: List[Any] = []

        if pdb_id_a and auth_asym_id_a:
            where_clauses.append("LOWER(pdb_id_A) = LOWER(?) AND LOWER(auth_asym_id_A) = LOWER(?)")
            params.extend([pdb_id_a, auth_asym_id_a])
        elif chain_ids:
            sub: List[str] = []
            for cid in chain_ids:
                parts = cid.split("_", 1)
                if len(parts) == 2:
                    sub.append("(LOWER(pdb_id_A) = LOWER(?) AND LOWER(auth_asym_id_A) = LOWER(?))")
                    params.extend(parts)
            if sub:
                where_clauses.append("(" + " OR ".join(sub) + ")")

        if rmsd_min is not None:
            where_clauses.append("CAST(RMSD AS REAL) >= ?")
            params.append(rmsd_min)
        if rmsd_max is not None:
            where_clauses.append("CAST(RMSD AS REAL) <= ?")
            params.append(rmsd_max)
        if tm_min is not None:
            where_clauses.append("CAST(TM1 AS REAL) >= ?")
            params.append(tm_min)
        if tm_max is not None:
            where_clauses.append("CAST(TM1 AS REAL) <= ?")
            params.append(tm_max)

        where_sql = ("WHERE " + " AND ".join(where_clauses)) if where_clauses else ""
        sql = f"""
            SELECT e.pdb_id_A, e.auth_asym_id_A, e.pdb_id_B, e.auth_asym_id_B,
                   CAST(e.TM1 AS REAL) AS TM1,
                   CAST(e.RMSD AS REAL) AS RMSD,
                   CAST(e.structure_sim AS REAL) AS structure_sim,
                   e."delta_Rosetta", e."delta_FoldX", e."delta_EvoEF2",
                   e."delta_RW", e."delta_RW+",
                   e.state_id_B, e.state_fidelity, e.avg_sim, e.observation_fidelity
            FROM edge e
            {where_sql}
            LIMIT ?
        """
        params.append(limit)

        conn = self._connect()
        try:
            df = pd.read_sql_query(sql, conn, params=params)
        finally:
            conn.close()

        # Rename columns with special characters so itertuples works cleanly
        df = df.rename(columns={"delta_RW+": "delta_RW_plus"})

        df["cluster_id"] = df["TM1"].apply(
            lambda t: _assign_cluster(t) if pd.notna(t) else "cluster_low"
        )

        from app.protein import tsv_loader
        df["length_a"] = [
            tsv_loader.get_sequence(r.pdb_id_A, r.auth_asym_id_A)[1]
            for r in df.itertuples(index=False)
        ]
        df["length_b"] = [
            tsv_loader.get_sequence(r.pdb_id_B, r.auth_asym_id_B)[1]
            for r in df.itertuples(index=False)
        ]
        node_rows = [
            tsv_loader.get_node_row(r.pdb_id_B, r.auth_asym_id_B)
            for r in df.itertuples(index=False)
        ]
        df["experimental_method"] = [row.get("experimental_method") if row else None for row in node_rows]
        df["pH"] = [row.get("pH") if row else None for row in node_rows]
        df["temp_K"] = [row.get("temp_K") if row else None for row in node_rows]
        df["base_label"] = [row.get("base_label") if row else None for row in node_rows]
        df["chain_composition"] = [row.get("chain_composition") if row else None for row in node_rows]

        return df

    def get_summary_stats(
        self,
        chain_ids: Optional[List[str]] = None,
        rmsd_min: Optional[float] = None,
        rmsd_max: Optional[float] = None,
        tm_min: Optional[float] = None,
        tm_max: Optional[float] = None,
    ) -> Dict[str, Any]:
        """Compute aggregate statistics via SQL, with optional filters."""
        if not any([chain_ids, rmsd_min, rmsd_max, tm_min, tm_max]) and self._summary:
            return self._summary

        where_clauses: List[str] = []
        params: List[Any] = []

        if chain_ids:
            sub: List[str] = []
            for cid in chain_ids:
                parts = cid.split("_", 1)
                if len(parts) == 2:
                    sub.append("(LOWER(pdb_id_A) = LOWER(?) AND LOWER(auth_asym_id_A) = LOWER(?))")
                    params.extend(parts)
            if sub:
                where_clauses.append("(" + " OR ".join(sub) + ")")

        if rmsd_min is not None:
            where_clauses.append("CAST(RMSD AS REAL) >= ?")
            params.append(rmsd_min)
        if rmsd_max is not None:
            where_clauses.append("CAST(RMSD AS REAL) <= ?")
            params.append(rmsd_max)
        if tm_min is not None:
            where_clauses.append("CAST(TM1 AS REAL) >= ?")
            params.append(tm_min)
        if tm_max is not None:
            where_clauses.append("CAST(TM1 AS REAL) <= ?")
            params.append(tm_max)

        where_sql = ("WHERE " + " AND ".join(where_clauses)) if where_clauses else ""

        conn = self._connect()
        try:
            cur = conn.cursor()
            cur.execute(
                f"SELECT COUNT(*), AVG(CAST(RMSD AS REAL)), AVG(CAST(TM1 AS REAL))"
                f" FROM edge {where_sql}",
                params,
            )
            total, avg_rmsd, avg_tm = cur.fetchone()

            cur.execute(
                f"SELECT COUNT(DISTINCT LOWER(pdb_id_A) || '_' || LOWER(auth_asym_id_A))"
                f" FROM edge {where_sql}",
                params,
            )
            unique_chains = cur.fetchone()[0]

            def _dist(col, bins, labels):
                cases = " ".join(
                    f"WHEN CAST({col} AS REAL) >= {lo} AND CAST({col} AS REAL) < {hi} THEN '{lbl}'"
                    for (lo, hi), lbl in zip(zip(bins, bins[1:]), labels)
                )
                sql = (
                    f"SELECT CASE {cases} ELSE '{labels[-1]}' END AS bucket, COUNT(*)"
                    f" FROM edge {where_sql} GROUP BY bucket"
                )
                cur.execute(sql, params)
                return {r[0]: r[1] for r in cur.fetchall()}

            rmsd_dist = _dist(
                "RMSD",
                [0, 0.5, 1.0, 1.5, 2.0, 5.0, 1e9],
                ["0-0.5", "0.5-1.0", "1.0-1.5", "1.5-2.0", "2.0-5.0", ">5.0"],
            )
            tm_dist = _dist(
                "TM1",
                [0, 0.5, 0.7, 0.85, 0.95, 1.01],
                ["0-0.5", "0.5-0.7", "0.7-0.85", "0.85-0.95", "0.95-1.0"],
            )
        finally:
            conn.close()

        from app.protein import tsv_loader
        lengths = []
        for row in tsv_loader._get_index().values():
            try:
                lengths.append(int(row["sequence_length"]))
            except (ValueError, TypeError):
                pass
        avg_len = sum(lengths) / len(lengths) if lengths else 0.0
        len_dist = {}
        for length in lengths:
            if length <= 100:
                bucket = "0-100"
            elif length <= 200:
                bucket = "100-200"
            elif length <= 300:
                bucket = "200-300"
            elif length <= 500:
                bucket = "300-500"
            elif length <= 1000:
                bucket = "500-1000"
            else:
                bucket = ">1000"
            len_dist[bucket] = len_dist.get(bucket, 0) + 1

        return {
            "total_records": total or 0,
            "unique_chains": unique_chains or 0,
            "avg_rmsd": avg_rmsd or 0.0,
            "avg_tm_score": avg_tm or 0.0,
            "avg_sequence_length": avg_len,
            "rmsd_distribution": rmsd_dist,
            "tm_score_distribution": tm_dist,
            "length_distribution": len_dist,
        }

    def get_total(self) -> int:
        return self._total_records