app/data/data_store.py

"""
Persistent DataStore backed by SQLite.
- Uploaded CSVs are written to the DB and survive server restarts.
- Synthetic data is used only when the DB has no rows for that table.
- In-memory cache (_bms, _fleet, etc.) avoids repeated DB reads within
  the same server session.
"""
import json
import pandas as pd
from typing import Optional, Dict, List

from app.data.synthetic_bms import get_all_bms_data
from app.data.synthetic_fleet import get_fleet_data
from app.data.synthetic_supply_chain import get_suppliers, get_inspections
from app.database import get_conn


class DataStore:
    def __init__(self):
        # In-memory cache — populated lazily from DB or synthetic generators
        self._bms:         Optional[Dict[str, pd.DataFrame]] = None
        self._fleet:       Optional[List[Dict]]              = None
        self._suppliers:   Optional[List[Dict]]              = None
        self._inspections: Optional[List[Dict]]              = None

    # ── BMS ──────────────────────────────────────────────────────────

    def get_bms(self) -> Dict[str, pd.DataFrame]:
        if self._bms is not None:
            return self._bms

        # Try loading from DB first
        conn = get_conn()
        rows = conn.execute("SELECT * FROM bms_data ORDER BY vehicle_id, timestamp").fetchall()
        conn.close()

        if rows:
            df = pd.DataFrame([dict(r) for r in rows])
            df['timestamp'] = pd.to_datetime(df['timestamp'])
            self._bms = {vid: grp.sort_values('timestamp') for vid, grp in df.groupby('vehicle_id')}
            return self._bms

        # Fall back to synthetic
        return get_all_bms_data()

    def set_bms(self, df: pd.DataFrame):
        required = ['timestamp', 'vehicle_id', 'soc', 'voltage', 'current', 'temperature', 'cycle_count']
        if 'vehicle_id' not in df.columns:
            raise ValueError("BMS CSV must contain 'vehicle_id' column")
        missing = [c for c in required if c not in df.columns]
        if missing:
            raise ValueError(f"Missing columns in BMS CSV: {missing}")

        df['timestamp'] = pd.to_datetime(df['timestamp'])

        # Persist to DB
        conn = get_conn()
        conn.execute("DELETE FROM bms_data")  # replace with new upload
        for _, row in df.iterrows():
            conn.execute(
                """INSERT INTO bms_data
                   (vehicle_id, timestamp, soc, voltage, current, temperature, cycle_count, odometer_km)
                   VALUES (?, ?, ?, ?, ?, ?, ?, ?)""",
                (
                    row['vehicle_id'],
                    str(row['timestamp']),
                    float(row.get('soc', 0)),
                    float(row.get('voltage', 0)),
                    float(row.get('current', 0)),
                    float(row.get('temperature', 0)),
                    int(row.get('cycle_count', 0)),
                    float(row.get('odometer_km', 0)) if 'odometer_km' in row else 0.0,
                )
            )
        conn.commit()
        conn.close()

        # Update cache
        self._bms = {vid: grp.sort_values('timestamp') for vid, grp in df.groupby('vehicle_id')}

    # ── Fleet ─────────────────────────────────────────────────────────

    def get_fleet(self) -> List[Dict]:
        if self._fleet is not None:
            return self._fleet

        conn = get_conn()
        rows = conn.execute("SELECT * FROM fleet_data").fetchall()
        conn.close()

        if rows:
            self._fleet = [dict(r) for r in rows]
            return self._fleet

        return get_fleet_data()

    def set_fleet(self, records: List[Dict]):
        required = ['asset_id', 'category', 'route_distance_km', 'payload_tons',
                    'duty_cycle_hrs', 'dwell_time_hrs', 'terrain_score',
                    'diesel_l_per_100km', 'age_years']
        if records:
            missing = [c for c in required if c not in records[0]]
            if missing:
                raise ValueError(f"Missing columns in fleet CSV: {missing}")

        conn = get_conn()
        conn.execute("DELETE FROM fleet_data")
        for r in records:
            conn.execute(
                """INSERT OR REPLACE INTO fleet_data
                   (asset_id, category, route_distance_km, payload_tons,
                    duty_cycle_hrs, dwell_time_hrs, terrain_score,
                    diesel_l_per_100km, age_years)
                   VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)""",
                (
                    r['asset_id'], r['category'],
                    float(r['route_distance_km']), float(r['payload_tons']),
                    float(r['duty_cycle_hrs']),    float(r['dwell_time_hrs']),
                    float(r['terrain_score']),     float(r['diesel_l_per_100km']),
                    float(r['age_years']),
                )
            )
        conn.commit()
        conn.close()
        self._fleet = records

    # ── Suppliers ─────────────────────────────────────────────────────

    def get_suppliers(self) -> List[Dict]:
        if self._suppliers is not None:
            return self._suppliers

        conn = get_conn()
        rows = conn.execute("SELECT * FROM supplier_data").fetchall()
        conn.close()

        if rows:
            self._suppliers = [dict(r) for r in rows]
            return self._suppliers

        return get_suppliers()

    def set_suppliers(self, records: List[Dict]):
        required = ['supplier_id', 'name', 'tier', 'material', 'country',
                    'geopolitical_risk', 'esg_score', 'quality_score',
                    'lead_time_days', 'concentration_pct']
        if records:
            missing = [c for c in required if c not in records[0]]
            if missing:
                raise ValueError(f"Missing columns in supplier CSV: {missing}")

        conn = get_conn()
        conn.execute("DELETE FROM supplier_data")
        for r in records:
            conn.execute(
                """INSERT OR REPLACE INTO supplier_data
                   (supplier_id, name, tier, material, country,
                    geopolitical_risk, esg_score, quality_score,
                    lead_time_days, concentration_pct)
                   VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""",
                (
                    r['supplier_id'], r['name'], int(r['tier']), r['material'], r['country'],
                    float(r['geopolitical_risk']), float(r['esg_score']),
                    float(r['quality_score']),     int(r['lead_time_days']),
                    float(r['concentration_pct']),
                )
            )
        conn.commit()
        conn.close()
        self._suppliers = records

    # ── Inspections ───────────────────────────────────────────────────

    def get_inspections(self) -> List[Dict]:
        if self._inspections is not None:
            return self._inspections

        conn = get_conn()
        rows = conn.execute("SELECT * FROM inspection_data").fetchall()
        conn.close()

        if rows:
            self._inspections = [dict(r) for r in rows]
            return self._inspections

        return get_inspections()

    def set_inspections(self, records: List[Dict]):
        required = ['lot_id', 'supplier_id', 'material', 'defect_rate', 'inspection_date']
        if records:
            missing = [c for c in required if c not in records[0]]
            if missing:
                raise ValueError(f"Missing columns in inspection CSV: {missing}")

        conn = get_conn()
        conn.execute("DELETE FROM inspection_data")
        for r in records:
            conn.execute(
                """INSERT INTO inspection_data
                   (lot_id, supplier_id, material, defect_rate, inspection_date, root_cause_hint)
                   VALUES (?, ?, ?, ?, ?, ?)""",
                (
                    r['lot_id'], r['supplier_id'], r['material'],
                    float(r['defect_rate']), r['inspection_date'],
                    r.get('root_cause_hint', ''),
                )
            )
        conn.commit()
        conn.close()
        self._inspections = records

    # ── Reset ─────────────────────────────────────────────────────────

    def reset(self):
        """Clear DB tables + in-memory cache → synthetic data returns."""
        conn = get_conn()
        conn.execute("DELETE FROM bms_data")
        conn.execute("DELETE FROM fleet_data")
        conn.execute("DELETE FROM supplier_data")
        conn.execute("DELETE FROM inspection_data")
        conn.commit()
        conn.close()

        self._bms         = None
        self._fleet       = None
        self._suppliers   = None
        self._inspections = None


store = DataStore()