db.py

"""
SQLite-backed database layer for the Hugging Face Space demo.

Mirrors the original Aurora PostgreSQL API surface (`get_metrics_distinct_values`,
`get_metrics_from_db_filtered`, `load_param_from_db`) so the agents and tools can be
re-used unchanged, but talks to a single-file `demo.sqlite` instead of a cloud DB.
"""
from __future__ import annotations

import os
import re
import sqlite3
from pathlib import Path
from typing import Iterable

import pandas as pd


HERE = Path(__file__).resolve().parent
DB_PATH = Path(os.environ.get("DEMO_DB_PATH", HERE / "demo.sqlite"))


# Column order matches the original sentinel_metrics CSV/Postgres schema.
METRICS_TABLE_COLUMNS = [
    "Period",
    "Region",
    "Level",
    "Product",
    "Country",
    "Cluster",
    "TA Market",
    "Class",
    "Current_Value",
    "Current_Value_MS",
    "Value_MS_Change",
    "Current_Value_Rank",
    "Value_Rank_Change",
    "Value_Growth",
    "Value_Trend_Reversal",
    "Current_Volume",
    "Current_Volume_MS",
    "Volume_MS_Change",
    "Current_Volume_Rank",
    "Volume_Rank_Change",
    "Volume_Growth",
    "Volume_Trend_Reversal",
    "ASP_Growth",
    "Calculation_Type",
]

# SQLite uses lowercase identifiers (no quoting headache) and we rename back when
# returning DataFrames to the agents.
_DF_TO_DB_COL: dict[str, str] = {
    "Period": "period",
    "Region": "region",
    "Level": "level",
    "Product": "product",
    "Country": "country",
    "Cluster": "cluster",
    "TA Market": "ta_market",
    "Class": "class",
    "Current_Value": "current_value",
    "Current_Value_MS": "current_value_ms",
    "Value_MS_Change": "value_ms_change",
    "Current_Value_Rank": "current_value_rank",
    "Value_Rank_Change": "value_rank_change",
    "Value_Growth": "value_growth",
    "Value_Trend_Reversal": "value_trend_reversal",
    "Current_Volume": "current_volume",
    "Current_Volume_MS": "current_volume_ms",
    "Volume_MS_Change": "volume_ms_change",
    "Current_Volume_Rank": "current_volume_rank",
    "Volume_Rank_Change": "volume_rank_change",
    "Volume_Growth": "volume_growth",
    "Volume_Trend_Reversal": "volume_trend_reversal",
    "ASP_Growth": "asp_growth",
    "Calculation_Type": "calculation_type",
}
_DB_TO_DF_COL: dict[str, str] = {v: k for k, v in _DF_TO_DB_COL.items()}

_METRICS_DDL = """
CREATE TABLE IF NOT EXISTS sentinel_metrics (
    period                TEXT,
    region                TEXT,
    level                 TEXT,
    product               TEXT,
    country               TEXT,
    cluster               TEXT,
    ta_market             TEXT,
    class                 TEXT,
    current_value         REAL,
    current_value_ms      REAL,
    value_ms_change       REAL,
    current_value_rank    INTEGER,
    value_rank_change     INTEGER,
    value_growth          REAL,
    value_trend_reversal  REAL,
    current_volume        REAL,
    current_volume_ms     REAL,
    volume_ms_change      REAL,
    current_volume_rank   INTEGER,
    volume_rank_change    INTEGER,
    volume_growth         REAL,
    volume_trend_reversal REAL,
    asp_growth            REAL,
    calculation_type      TEXT
);
"""

_METRICS_INDEXES = (
    "CREATE INDEX IF NOT EXISTS idx_metrics_region ON sentinel_metrics(region);",
    "CREATE INDEX IF NOT EXISTS idx_metrics_period ON sentinel_metrics(period);",
    "CREATE INDEX IF NOT EXISTS idx_metrics_cluster ON sentinel_metrics(cluster);",
    "CREATE INDEX IF NOT EXISTS idx_metrics_product ON sentinel_metrics(product);",
    "CREATE INDEX IF NOT EXISTS idx_metrics_calc ON sentinel_metrics(calculation_type);",
    "CREATE INDEX IF NOT EXISTS idx_metrics_class ON sentinel_metrics(class);",
)

# Parameter tables — same key names as the original code so callers don't need to change.
PARAM_TABLE_MAP: dict[str, str] = {
    "cluster_mapping": "sentinel_cluster_mapping",
    "market_summary_mapping": "sentinel_market_summary_mapping",
    "threshold_matrix": "sentinel_threshold_matrix",
    "country_region_mapping": "sentinel_country_region_mapping",
    "merck_products": "sentinel_company_products",  # legacy key kept for parameter_tools compatibility
    "company_products": "sentinel_company_products",
}

_PARAM_DDL: dict[str, str] = {
    "sentinel_cluster_mapping": """
        CREATE TABLE IF NOT EXISTS sentinel_cluster_mapping (
            country  TEXT,
            market   TEXT,
            cluster  TEXT,
            region   TEXT
        );
    """,
    "sentinel_market_summary_mapping": """
        CREATE TABLE IF NOT EXISTS sentinel_market_summary_mapping (
            region             TEXT,
            market             TEXT,
            value_summary      TEXT,
            volume_summary     TEXT,
            timeframe_summary  TEXT
        );
    """,
    "sentinel_threshold_matrix": """
        CREATE TABLE IF NOT EXISTS sentinel_threshold_matrix (
            cluster           TEXT,
            market            TEXT,
            value_ms_change   REAL,
            value_growth      REAL,
            volume_ms_change  REAL,
            volume_growth     REAL,
            asp_growth        REAL,
            region            TEXT
        );
    """,
    "sentinel_country_region_mapping": """
        CREATE TABLE IF NOT EXISTS sentinel_country_region_mapping (
            country_code  TEXT,
            country_name  TEXT,
            region        TEXT
        );
    """,
    "sentinel_company_products": """
        CREATE TABLE IF NOT EXISTS sentinel_company_products (
            product  TEXT
        );
    """,
}


def _connect() -> sqlite3.Connection:
    """Open a SQLite connection with row factory and foreign keys enabled."""
    DB_PATH.parent.mkdir(parents=True, exist_ok=True)
    conn = sqlite3.connect(str(DB_PATH))
    conn.execute("PRAGMA journal_mode = WAL;")
    return conn


def create_all_tables() -> None:
    """Create the metrics table, indexes, and all parameter tables if missing."""
    with _connect() as conn:
        conn.execute(_METRICS_DDL)
        for stmt in _METRICS_INDEXES:
            conn.execute(stmt)
        for ddl in _PARAM_DDL.values():
            conn.execute(ddl)
        conn.commit()


def db_is_seeded() -> bool:
    """Return True when the metrics table has at least one row."""
    if not DB_PATH.exists():
        return False
    try:
        with _connect() as conn:
            row = conn.execute("SELECT COUNT(1) FROM sentinel_metrics;").fetchone()
            return bool(row and row[0] > 0)
    except sqlite3.OperationalError:
        return False


# ---------------------------------------------------------------------------
# Metrics queries
# ---------------------------------------------------------------------------

def _rename_db_to_df(df: pd.DataFrame) -> pd.DataFrame:
    """Rename SQLite (lowercase) columns back to the canonical DataFrame names."""
    return df.rename(columns=_DB_TO_DF_COL)


def _period_sort_key(period: str) -> tuple[int, int]:
    match = re.match(r"^\s*(\d{2})\s*Q([1-4])\s*$", str(period or ""), re.IGNORECASE)
    return (int(match.group(1)), int(match.group(2))) if match else (-1, -1)


def get_metrics_distinct_values() -> dict:
    """
    Return distinct values per filterable column in sentinel_metrics.

    Also returns ``_region_latest_period``: a dict mapping each region to its most
    recent period, used by the filter extractor to default smartly when the user
    doesn't specify a period.
    """
    if not DB_PATH.exists():
        return {}
    with _connect() as conn:
        cols_lower = ["region", "period", "cluster", "product", "calculation_type",
                      "ta_market", "class", "level"]
        result: dict = {}
        for db_col in cols_lower:
            df_col = _DB_TO_DF_COL[db_col]
            rows = conn.execute(
                f"SELECT DISTINCT {db_col} FROM sentinel_metrics "
                f"WHERE {db_col} IS NOT NULL ORDER BY {db_col};"
            ).fetchall()
            result[df_col] = [str(r[0]) for r in rows if r[0] is not None]

        region_period_rows = conn.execute(
            "SELECT region, period FROM sentinel_metrics "
            "WHERE region IS NOT NULL AND period IS NOT NULL "
            "GROUP BY region, period;"
        ).fetchall()

    region_periods: dict[str, list[str]] = {}
    for region, period in region_period_rows:
        region_periods.setdefault(str(region), []).append(str(period))
    region_latest: dict[str, str] = {
        r: sorted(ps, key=_period_sort_key)[-1] for r, ps in region_periods.items() if ps
    }
    result["_region_latest_period"] = region_latest
    return result


_FILTER_COL_TO_DB: dict[str, str] = {
    "Region": "region",
    "Period": "period",
    "Cluster": "cluster",
    "Product": "product",
    "Calculation_Type": "calculation_type",
    "TA Market": "ta_market",
    "Class": "class",
    "Level": "level",
}


def get_metrics_from_db_filtered(filters: dict) -> pd.DataFrame:
    """
    Return the sentinel_metrics rows that match *filters*.

    *filters* keys must be DataFrame column names (e.g. ``Region``, ``Period``,
    ``Cluster``, ``Product``, ``Calculation_Type``, ``TA Market``, ``Class``).
    Unrecognised keys are silently ignored.  Empty *filters* returns all rows.
    """
    where_clauses: list[str] = []
    params: list = []
    for col, filter_values in (filters or {}).items():
        db_col = _FILTER_COL_TO_DB.get(col)
        if not db_col or not filter_values:
            continue
        clean = [str(v).strip() for v in filter_values if str(v).strip()]
        if col == "Period":
            clean = [v.replace(" ", "") for v in clean]
        if not clean:
            continue
        placeholders = ", ".join(["?"] * len(clean))
        where_clauses.append(f"{db_col} IN ({placeholders})")
        params.extend(clean)

    where_sql = ("WHERE " + " AND ".join(where_clauses)) if where_clauses else ""
    query = f"SELECT * FROM sentinel_metrics {where_sql};"
    with _connect() as conn:
        df = pd.read_sql_query(query, conn, params=params if params else None)
    df = _rename_db_to_df(df)
    # Synthesize Cluster from Country/Region when missing (kept for backward compat
    # with the original Postgres helper, even though our schema always populates it).
    if "Cluster" not in df.columns:
        if {"Level", "Country", "Region"}.issubset(df.columns):
            df["Cluster"] = df.apply(
                lambda r: r["Country"] if r["Level"] == "Country" else r["Region"], axis=1
            )
        elif "Country" in df.columns:
            df["Cluster"] = df["Country"].fillna("").astype(str)
    return df


# ---------------------------------------------------------------------------
# Parameter tables
# ---------------------------------------------------------------------------

def load_param_from_db(key: str) -> pd.DataFrame:
    """
    Load a parameter table by its logical key.

    Keys: ``cluster_mapping``, ``market_summary_mapping``, ``threshold_matrix``,
    ``country_region_mapping``, ``merck_products`` (alias of ``company_products``).
    Returns a DataFrame matching the original CSV column casing.
    """
    if key not in PARAM_TABLE_MAP:
        raise ValueError(
            f"Unknown parameter key: {key!r}. Valid keys: {sorted(set(PARAM_TABLE_MAP))}"
        )
    table = PARAM_TABLE_MAP[key]
    with _connect() as conn:
        df = pd.read_sql_query(f"SELECT * FROM {table};", conn)
    if key == "threshold_matrix":
        df = df.rename(columns={
            "cluster": "Cluster",
            "value_ms_change": "Value_MS_Change",
            "value_growth": "Value_Growth",
            "volume_ms_change": "Volume_MS_Change",
            "volume_growth": "Volume_Growth",
            "asp_growth": "ASP_Growth",
        })
    return df


# ---------------------------------------------------------------------------
# Bulk insert helpers (used by seed_data.py)
# ---------------------------------------------------------------------------

def replace_param_table(key: str, df: pd.DataFrame) -> int:
    """Replace the contents of a parameter table with *df*. Returns rows inserted."""
    if key not in PARAM_TABLE_MAP:
        raise ValueError(f"Unknown parameter key: {key!r}")
    table = PARAM_TABLE_MAP[key]
    df = df.copy()
    df.columns = [c.lower() for c in df.columns]
    with _connect() as conn:
        conn.execute(f"DELETE FROM {table};")
        df.to_sql(table, conn, if_exists="append", index=False)
        conn.commit()
    return len(df)


def insert_metrics(df: pd.DataFrame) -> int:
    """Insert metrics rows (DataFrame uses mixed-case column names)."""
    if df is None or df.empty:
        return 0
    db_df = df.rename(columns=_DF_TO_DB_COL)
    keep_cols = [c for c in _DF_TO_DB_COL.values() if c in db_df.columns]
    db_df = db_df[keep_cols]
    with _connect() as conn:
        db_df.to_sql("sentinel_metrics", conn, if_exists="append", index=False)
        conn.commit()
    return len(db_df)


def replace_metrics(df: pd.DataFrame) -> int:
    """Truncate the metrics table and insert *df*. Returns rows inserted."""
    with _connect() as conn:
        conn.execute("DELETE FROM sentinel_metrics;")
        conn.commit()
    return insert_metrics(df)