results_db.py

"""Complete SQLite store for SWMM model inputs, outputs, summaries and AI retrieval."""
from __future__ import annotations

import json
import re
import sqlite3
import tempfile
from pathlib import Path
from typing import Any

import pandas as pd


class ResultDatabase:
    """File-backed SQLite database containing the complete model and simulation dataset.

    The database is intentionally file-backed rather than ``:memory:`` so the user can
    download it, inspect it in any SQLite client, and retain an auditable simulation
    artefact. The LLM still receives only bounded query results.
    """

    def __init__(self, db_path: str | Path | None = None) -> None:
        if db_path is None:
            fd, name = tempfile.mkstemp(prefix="swmm_model_", suffix=".sqlite")
            Path(name).unlink(missing_ok=True)
            try:
                import os
                os.close(fd)
            except OSError:
                pass
            self.path = Path(name)
        else:
            self.path = Path(db_path)
        self.connection = sqlite3.connect(str(self.path), check_same_thread=False)
        self.connection.execute("PRAGMA journal_mode=WAL")
        self.connection.execute("PRAGMA synchronous=NORMAL")
        self.connection.execute("PRAGMA foreign_keys=ON")

    @staticmethod
    def _safe_table_name(section: str) -> str:
        clean = re.sub(r"[^a-zA-Z0-9_]+", "_", section.strip().lower()).strip("_")
        return f"inp_{clean or 'unknown'}"

    def _load_complete_input(self, inp_path: str | Path) -> None:
        path = Path(inp_path)
        text = path.read_text(encoding="utf-8", errors="replace")
        line_rows: list[dict[str, Any]] = []
        section_rows: dict[str, list[dict[str, Any]]] = {}
        catalog: list[dict[str, Any]] = []
        current_section = "PREAMBLE"
        section_row_no = 0

        for line_no, raw in enumerate(text.splitlines(), start=1):
            stripped = raw.strip()
            is_section = stripped.startswith("[") and "]" in stripped
            if is_section:
                current_section = stripped[1:stripped.index("]")].strip().upper()
                section_row_no = 0
            elif stripped and not stripped.startswith(";"):
                section_row_no += 1

            line_rows.append({
                "line_no": line_no,
                "section_name": current_section,
                "section_row_no": section_row_no if not is_section else 0,
                "raw_text": raw,
                "stripped_text": stripped,
                "is_blank": int(not stripped),
                "is_comment": int(stripped.startswith(";")),
                "is_section_header": int(is_section),
            })

            if stripped and not stripped.startswith(";") and not is_section:
                # Keep both the exact raw line and all whitespace-delimited values.
                tokens = stripped.split()
                row = {
                    "row_no": section_row_no,
                    "source_line_no": line_no,
                    "raw_text": raw,
                }
                for i, token in enumerate(tokens, start=1):
                    row[f"value_{i}"] = token
                section_rows.setdefault(current_section, []).append(row)

        pd.DataFrame(line_rows).to_sql(
            "model_input_lines", self.connection, if_exists="replace", index=False
        )
        pd.DataFrame([{
            "file_name": path.name,
            "absolute_path_at_run": str(path),
            "byte_size": path.stat().st_size,
            "line_count": len(line_rows),
            "full_text": text,
        }]).to_sql("model_input_file", self.connection, if_exists="replace", index=False)

        for section, rows in section_rows.items():
            table = self._safe_table_name(section)
            pd.DataFrame(rows).to_sql(table, self.connection, if_exists="replace", index=False)
            catalog.append({
                "section_name": section,
                "table_name": table,
                "row_count": len(rows),
                "max_values_per_row": max(
                    (sum(1 for k in r if k.startswith("value_")) for r in rows), default=0
                ),
            })

        pd.DataFrame(catalog).to_sql(
            "model_input_section_catalog", self.connection, if_exists="replace", index=False
        )

    @staticmethod
    def _time_strings(times: list[Any]) -> list[str]:
        return [t.isoformat(sep=" ") if hasattr(t, "isoformat") else str(t) for t in times]

    def _load_complete_outputs(self, results: dict[str, Any]) -> None:
        times = results.get("times", [])
        time_strings = self._time_strings(times)

        node_rows: list[dict[str, Any]] = []
        for node_id, data in results.get("node_ts", {}).items():
            n = max((len(data.get(k, [])) for k in (
                "depth", "flooding", "inflow", "head", "outflow", "volume"
            )), default=0)
            for i in range(n):
                node_rows.append({
                    "time_index": i,
                    "timestamp": time_strings[i] if i < len(time_strings) else str(i),
                    "node_id": node_id,
                    "depth": data.get("depth", [None] * n)[i],
                    "flooding": data.get("flooding", [None] * n)[i],
                    "inflow": data.get("inflow", [None] * n)[i],
                    "head": data.get("head", [None] * n)[i],
                    "outflow": data.get("outflow", [None] * n)[i],
                    "volume": data.get("volume", [None] * n)[i],
                })
        pd.DataFrame(node_rows, columns=[
            "time_index", "timestamp", "node_id", "depth", "flooding",
            "inflow", "head", "outflow", "volume"
        ]).to_sql("node_timeseries", self.connection, if_exists="replace", index=False, chunksize=5000)

        node_static = [{
            "node_id": node_id,
            "invert_elevation": data.get("invert_elevation"),
            "full_depth": data.get("full_depth"),
        } for node_id, data in results.get("node_ts", {}).items()]
        pd.DataFrame(node_static, columns=["node_id", "invert_elevation", "full_depth"]).to_sql(
            "node_output_metadata", self.connection, if_exists="replace", index=False
        )

        link_rows: list[dict[str, Any]] = []
        for link_id, data in results.get("link_ts", {}).items():
            n = max((len(data.get(k, [])) for k in (
                "flow", "depth", "velocity", "volume", "capacity"
            )), default=0)
            for i in range(n):
                link_rows.append({
                    "time_index": i,
                    "timestamp": time_strings[i] if i < len(time_strings) else str(i),
                    "link_id": link_id,
                    "flow": data.get("flow", [None] * n)[i],
                    "depth": data.get("depth", [None] * n)[i],
                    "velocity": data.get("velocity", [None] * n)[i],
                    "volume": data.get("volume", [None] * n)[i],
                    "capacity": data.get("capacity", [None] * n)[i],
                })
        pd.DataFrame(link_rows, columns=[
            "time_index", "timestamp", "link_id", "flow", "depth",
            "velocity", "volume", "capacity"
        ]).to_sql("link_timeseries", self.connection, if_exists="replace", index=False, chunksize=5000)

        link_static = [{
            "link_id": link_id,
            "length": data.get("length"),
            "roughness": data.get("roughness"),
            "diameter": data.get("diameter"),
        } for link_id, data in results.get("link_ts", {}).items()]
        pd.DataFrame(link_static, columns=["link_id", "length", "roughness", "diameter"]).to_sql(
            "link_output_metadata", self.connection, if_exists="replace", index=False
        )

        sub_rows: list[dict[str, Any]] = []
        for sub_id, data in results.get("sub_ts", {}).items():
            n = max((len(data.get(k, [])) for k in ("runoff", "rainfall", "infil")), default=0)
            for i in range(n):
                sub_rows.append({
                    "time_index": i,
                    "timestamp": time_strings[i] if i < len(time_strings) else str(i),
                    "subcatchment_id": sub_id,
                    "runoff": data.get("runoff", [None] * n)[i],
                    "rainfall": data.get("rainfall", [None] * n)[i],
                    "infiltration": data.get("infil", [None] * n)[i],
                })
        pd.DataFrame(sub_rows, columns=[
            "time_index", "timestamp", "subcatchment_id", "runoff", "rainfall", "infiltration"
        ]).to_sql("subcatchment_timeseries", self.connection, if_exists="replace", index=False, chunksize=5000)

        metadata = results.get("metadata", {})
        meta_rows = []
        warnings = metadata.get("warnings", []) or []
        for key, value in metadata.items():
            if key == "warnings":
                continue
            if isinstance(value, (dict, list, tuple)):
                value = json.dumps(value, default=str)
            elif hasattr(value, "isoformat"):
                value = value.isoformat()
            meta_rows.append({"key": key, "value": value})
        pd.DataFrame(meta_rows, columns=["key", "value"]).to_sql(
            "simulation_metadata", self.connection, if_exists="replace", index=False
        )
        pd.DataFrame(warnings, columns=["code", "message"]).to_sql(
            "simulation_warnings", self.connection, if_exists="replace", index=False
        )

        # Indexes materially reduce retrieval cost for large models.
        self.connection.executescript("""
        CREATE INDEX IF NOT EXISTS idx_node_ts_id_time ON node_timeseries(node_id, time_index);
        CREATE INDEX IF NOT EXISTS idx_node_ts_flood ON node_timeseries(flooding DESC);
        CREATE INDEX IF NOT EXISTS idx_link_ts_id_time ON link_timeseries(link_id, time_index);
        CREATE INDEX IF NOT EXISTS idx_link_ts_capacity ON link_timeseries(capacity DESC);
        CREATE INDEX IF NOT EXISTS idx_sub_ts_id_time ON subcatchment_timeseries(subcatchment_id, time_index);
        CREATE INDEX IF NOT EXISTS idx_input_section ON model_input_lines(section_name, section_row_no);
        """)

    def load(
        self,
        node_summary: pd.DataFrame,
        link_summary: pd.DataFrame,
        sub_summary: pd.DataFrame,
        *,
        inp_path: str | Path,
        results: dict[str, Any],
    ) -> None:
        node_summary.to_sql("node_summary", self.connection, if_exists="replace", index=False)
        link_summary.to_sql("link_summary", self.connection, if_exists="replace", index=False)
        sub_summary.to_sql("subcatchment_summary", self.connection, if_exists="replace", index=False)
        self._load_complete_input(inp_path)
        self._load_complete_outputs(results)
        self.connection.commit()

    def table_catalog(self) -> pd.DataFrame:
        return pd.read_sql_query("""
            SELECT name AS table_name
            FROM sqlite_master
            WHERE type='table' AND name NOT LIKE 'sqlite_%'
            ORDER BY name
        """, self.connection)

    def export_bytes(self) -> bytes:
        self.connection.commit()
        # Checkpoint WAL so the downloaded main file is self-contained.
        try:
            self.connection.execute("PRAGMA wal_checkpoint(FULL)")
        except sqlite3.DatabaseError:
            pass
        return self.path.read_bytes()

    @staticmethod
    def _quoted(value: str) -> str:
        return value.replace("'", "''")

    def engineering_context(self, question: str, limit: int = 20) -> str:
        """Return compact SQL-derived context while retaining the full database locally."""
        q = question.lower()
        parts: list[str] = []

        # Asset IDs in the question are used to retrieve exact time series.
        ids = re.findall(r"\b[A-Za-z][A-Za-z0-9_.:-]*\b", question)
        known_noise = {"which", "what", "when", "where", "show", "compare", "node", "nodes",
                       "link", "links", "pipe", "pipes", "conduit", "flow", "depth", "runoff",
                       "the", "and", "for", "from", "with", "during", "model"}
        ids = [x for x in ids if x.lower() not in known_noise][:8]

        if any(k in q for k in ("input", "roughness", "diameter", "length", "invert", "elevation",
                                "option", "rain gage", "timeseries", "control", "weir", "orifice",
                                "pump", "storage", "infiltration", "subarea")):
            section_terms = {
                "roughness": "CONDUITS", "diameter": "XSECTIONS", "length": "CONDUITS",
                "invert": "JUNCTIONS", "elevation": "JUNCTIONS", "control": "CONTROLS",
                "pump": "PUMPS", "weir": "WEIRS", "orifice": "ORIFICES",
                "storage": "STORAGE", "infiltration": "INFILTRATION", "subarea": "SUBAREAS",
                "rain": "RAINGAGES", "option": "OPTIONS",
            }
            selected = {v for k, v in section_terms.items() if k in q}
            if not selected:
                selected = {"OPTIONS", "JUNCTIONS", "CONDUITS", "XSECTIONS", "SUBCATCHMENTS"}
            names = ",".join(f"'{self._quoted(s)}'" for s in selected)
            sql = f"""SELECT section_name, section_row_no, raw_text
                      FROM model_input_lines
                      WHERE section_name IN ({names}) AND is_comment=0 AND is_blank=0
                      LIMIT {int(limit * 2)}"""
            parts.append("MODEL INPUT\n" + pd.read_sql_query(sql, self.connection).to_csv(index=False))

        if any(k in q for k in ("flood", "node", "manhole", "head", "inflow", "outflow")):
            parts.append("NODE SUMMARY\n" + pd.read_sql_query(
                f'''SELECT * FROM node_summary ORDER BY "Peak Flooding (m³/s)" DESC, "Depth Ratio" DESC LIMIT {int(limit)}''',
                self.connection).to_csv(index=False))
            for asset_id in ids:
                df = pd.read_sql_query(
                    """SELECT * FROM node_timeseries WHERE lower(node_id)=lower(?)
                       ORDER BY time_index LIMIT ?""", self.connection, params=(asset_id, int(limit * 3)))
                if not df.empty:
                    parts.append(f"NODE TIMESERIES {asset_id}\n" + df.to_csv(index=False))

        if any(k in q for k in ("conduit", "pipe", "link", "surcharge", "velocity", "capacity", "flow")):
            parts.append("LINK SUMMARY\n" + pd.read_sql_query(
                f'''SELECT * FROM link_summary ORDER BY "Depth Ratio" DESC, "Peak Flow (m³/s)" DESC LIMIT {int(limit)}''',
                self.connection).to_csv(index=False))
            for asset_id in ids:
                df = pd.read_sql_query(
                    """SELECT * FROM link_timeseries WHERE lower(link_id)=lower(?)
                       ORDER BY time_index LIMIT ?""", self.connection, params=(asset_id, int(limit * 3)))
                if not df.empty:
                    parts.append(f"LINK TIMESERIES {asset_id}\n" + df.to_csv(index=False))

        if any(k in q for k in ("subcatch", "runoff", "rain", "catchment", "hydrology", "infiltration")):
            parts.append("SUBCATCHMENT SUMMARY\n" + pd.read_sql_query(
                f'''SELECT * FROM subcatchment_summary ORDER BY "Peak Runoff (m³/s)" DESC LIMIT {int(limit)}''',
                self.connection).to_csv(index=False))
            for asset_id in ids:
                df = pd.read_sql_query(
                    """SELECT * FROM subcatchment_timeseries WHERE lower(subcatchment_id)=lower(?)
                       ORDER BY time_index LIMIT ?""", self.connection, params=(asset_id, int(limit * 3)))
                if not df.empty:
                    parts.append(f"SUBCATCHMENT TIMESERIES {asset_id}\n" + df.to_csv(index=False))

        if not parts:
            for table in ("simulation_metadata", "node_summary", "link_summary", "subcatchment_summary"):
                parts.append(table.upper() + "\n" + pd.read_sql_query(
                    f"SELECT * FROM {table} LIMIT 10", self.connection).to_csv(index=False))

        return "\n".join(parts)