app/core/database.py

"""SQLite + ChromaDB initialization and management."""
import sqlite3, json, logging
from pathlib import Path
from datetime import datetime
from typing import List, Dict, Any
from contextlib import contextmanager
from app.core.config import settings

logger = logging.getLogger("deltamind.db")
DB_PATH = settings.data_dir / "deltamind.db"

def init_database():
    """Initialize SQLite database with operational tables."""
    Path(settings.data_dir).mkdir(parents=True, exist_ok=True)
    with get_db() as conn:
        c = conn.cursor()
        c.execute("""CREATE TABLE IF NOT EXISTS flare_detections (
            id INTEGER PRIMARY KEY AUTOINCREMENT, latitude REAL, longitude REAL,
            brightness_temp REAL, frp REAL, confidence REAL, satellite TEXT,
            detection_time TEXT, oml_id TEXT, severity TEXT DEFAULT 'medium',
            created_at TEXT DEFAULT CURRENT_TIMESTAMP)""")
        c.execute("""CREATE TABLE IF NOT EXISTS spill_reports (
            id INTEGER PRIMARY KEY AUTOINCREMENT, latitude REAL, longitude REAL,
            spill_source TEXT, spill_cause TEXT, spill_volume REAL, status TEXT,
            report_date TEXT, oml_id TEXT, severity TEXT DEFAULT 'medium',
            created_at TEXT DEFAULT CURRENT_TIMESTAMP)""")
        c.execute("""CREATE TABLE IF NOT EXISTS anomaly_detections (
            id INTEGER PRIMARY KEY AUTOINCREMENT, source TEXT, anomaly_type TEXT,
            latitude REAL, longitude REAL, severity_score REAL, description TEXT,
            raw_data TEXT, status TEXT DEFAULT 'active',
            created_at TEXT DEFAULT CURRENT_TIMESTAMP)""")
        c.execute("""CREATE TABLE IF NOT EXISTS maintenance_predictions (
            id INTEGER PRIMARY KEY AUTOINCREMENT, equipment_id TEXT, equipment_type TEXT,
            oml_id TEXT, failure_probability REAL, predicted_days_to_failure INTEGER,
            recommended_action TEXT, severity TEXT, status TEXT DEFAULT 'pending',
            created_at TEXT DEFAULT CURRENT_TIMESTAMP)""")
        c.execute("""CREATE TABLE IF NOT EXISTS chat_history (
            id INTEGER PRIMARY KEY AUTOINCREMENT, session_id TEXT, role TEXT,
            content TEXT, provider TEXT, metadata TEXT,
            created_at TEXT DEFAULT CURRENT_TIMESTAMP)""")
        c.execute("""CREATE TABLE IF NOT EXISTS production_data (
            id INTEGER PRIMARY KEY AUTOINCREMENT, oml_id TEXT, well_id TEXT,
            field_name TEXT, oil_rate_bpd REAL, gas_rate_mmscfd REAL,
            water_cut_pct REAL, wellhead_pressure_psi REAL, temperature_f REAL,
            status TEXT, record_date TEXT, created_at TEXT DEFAULT CURRENT_TIMESTAMP)""")
        c.execute("""CREATE TABLE IF NOT EXISTS alerts (
            id INTEGER PRIMARY KEY AUTOINCREMENT, alert_type TEXT, severity TEXT,
            title TEXT, message TEXT, source TEXT, latitude REAL, longitude REAL,
            oml_id TEXT, acknowledged INTEGER DEFAULT 0,
            created_at TEXT DEFAULT CURRENT_TIMESTAMP)""")
        conn.commit()
    logger.info("SQLite Database initialized successfully")

@contextmanager
def get_db():
    conn = sqlite3.connect(str(DB_PATH))
    conn.row_factory = sqlite3.Row
    try:
        yield conn
    finally:
        conn.close()

def insert_flare(data: Dict) -> int:
    with get_db() as conn:
        c = conn.cursor()
        c.execute("INSERT INTO flare_detections (latitude,longitude,brightness_temp,frp,confidence,satellite,detection_time,oml_id,severity) VALUES (?,?,?,?,?,?,?,?,?)",
            (data.get("latitude"),data.get("longitude"),data.get("brightness_temp"),
             data.get("frp"),data.get("confidence"),data.get("satellite"),
             data.get("detection_time"),data.get("oml_id"),data.get("severity","medium")))
        conn.commit()
        return c.lastrowid

def insert_alert(data: Dict) -> int:
    with get_db() as conn:
        c = conn.cursor()
        c.execute("INSERT INTO alerts (alert_type,severity,title,message,source,latitude,longitude,oml_id) VALUES (?,?,?,?,?,?,?,?)",
            (data.get("alert_type"),data.get("severity"),data.get("title"),
             data.get("message"),data.get("source"),data.get("latitude"),
             data.get("longitude"),data.get("oml_id")))
        conn.commit()
        return c.lastrowid

def insert_production(data: Dict) -> int:
    with get_db() as conn:
        c = conn.cursor()
        c.execute("INSERT INTO production_data (oml_id,well_id,field_name,oil_rate_bpd,gas_rate_mmscfd,water_cut_pct,wellhead_pressure_psi,temperature_f,status,record_date) VALUES (?,?,?,?,?,?,?,?,?,?)",
            (data.get("oml_id"),data.get("well_id"),data.get("field_name"),
             data.get("oil_rate_bpd"),data.get("gas_rate_mmscfd"),data.get("water_cut_pct"),
             data.get("wellhead_pressure_psi"),data.get("temperature_f"),
             data.get("status","active"),data.get("record_date")))
        conn.commit()
        return c.lastrowid

def get_recent_alerts(limit: int = 50) -> List[Dict]:
    with get_db() as conn:
        rows = conn.cursor().execute("SELECT * FROM alerts ORDER BY created_at DESC LIMIT ?", (limit,)).fetchall()
        return [dict(r) for r in rows]

def get_production_data(oml_id: str = None, limit: int = 100) -> List[Dict]:
    with get_db() as conn:
        if oml_id:
            rows = conn.cursor().execute("SELECT * FROM production_data WHERE oml_id=? ORDER BY record_date DESC LIMIT ?", (oml_id, limit)).fetchall()
        else:
            rows = conn.cursor().execute("SELECT * FROM production_data ORDER BY record_date DESC LIMIT ?", (limit,)).fetchall()
        return [dict(r) for r in rows]

def save_chat(session_id: str, role: str, content: str, provider: str = "", metadata: dict = None):
    with get_db() as conn:
        conn.cursor().execute("INSERT INTO chat_history (session_id,role,content,provider,metadata) VALUES (?,?,?,?,?)",
            (session_id, role, content, provider, json.dumps(metadata or {})))
        conn.commit()

def get_chat_history(session_id: str, limit: int = 20) -> List[Dict]:
    with get_db() as conn:
        rows = conn.cursor().execute("SELECT * FROM chat_history WHERE session_id=? ORDER BY created_at DESC LIMIT ?", (session_id, limit)).fetchall()
        return [dict(r) for r in rows]

def get_dashboard_stats() -> Dict:
    with get_db() as conn:
        c = conn.cursor()
        return {
            "total_alerts": c.execute("SELECT COUNT(*) FROM alerts").fetchone()[0],
            "active_alerts": c.execute("SELECT COUNT(*) FROM alerts WHERE acknowledged=0").fetchone()[0],
            "critical_alerts": c.execute("SELECT COUNT(*) FROM alerts WHERE severity='critical'").fetchone()[0],
            "total_flares": c.execute("SELECT COUNT(*) FROM flare_detections").fetchone()[0],
            "total_anomalies": c.execute("SELECT COUNT(*) FROM anomaly_detections").fetchone()[0],
            "active_wells": c.execute("SELECT COUNT(DISTINCT well_id) FROM production_data WHERE status='active'").fetchone()[0],
        }

# ── ChromaDB RAG Setup ──
_chroma_client = None
_chroma_collection = None

def init_chromadb():
    global _chroma_client, _chroma_collection
    try:
        import chromadb
        from chromadb.config import Settings as CS
        Path(settings.chroma_persist_dir).mkdir(parents=True, exist_ok=True)
        _chroma_client = chromadb.PersistentClient(path=settings.chroma_persist_dir, settings=CS(anonymized_telemetry=False))
        _chroma_collection = _chroma_client.get_or_create_collection(name=settings.chroma_collection)
        logger.info(f"ChromaDB initialized: {_chroma_collection.count()} documents")
    except Exception as e:
        logger.error(f"ChromaDB init failed: {e}")

def add_documents(documents: List[str], metadatas: List[Dict], ids: List[str]):
    global _chroma_collection
    if _chroma_collection is None: init_chromadb()
    if _chroma_collection: 
        _chroma_collection.upsert(documents=documents, metadatas=metadatas, ids=ids)

def query_documents(query: str, n_results: int = 5) -> Dict:
    global _chroma_collection
    if _chroma_collection is None: init_chromadb()
    if not _chroma_collection: return {"documents":[],"metadatas":[],"distances":[]}
    try:
        return _chroma_collection.query(query_texts=[query], n_results=n_results)
    except Exception as e:
        logger.error(f"ChromaDB Query failed: {e}")
        return {"documents":[],"metadatas":[],"distances":[]}

def seed_knowledge_base():
    """Seed initial SOPs and regulations into Vector DB."""
    docs = [
        "Renaissance Africa Energy operates 18 Oil Mining Leases across onshore, swamp, and shallow water terrains in the Niger Delta. Manages Bonny and Forcados terminals and Sea Eagle FPSO.",
        "SOP for ESP failure: 1) Confirm alarm via SCADA. 2) Check vibration and temperature. 3) If vibration >0.5 in/s, shutdown. 4) Create work order. 5) Notify supervisor. 6) Schedule workover within 5 days.",
        "Pipeline integrity: Monthly aerial surveillance. SAR satellite for ground disturbance. Report unauthorized excavation within 50m of ROW immediately.",
        "Gas flaring: NUPRC mandates max 5% flaring. Renaissance targets zero routine flaring by 2028. All flaring reported within 24 hours.",
        "Daily production reports by 0700. Include oil rate bpd, gas rate MMSCF/D, water cut %, wellhead pressure, deferred volumes.",
        "Security alert: Pipeline breach detected: 1) Alert security. 2) Dispatch patrol. 3) Notify ops manager. 4) GPS + photos. 5) Incident report within 4hrs. 6) Notify partners if loss >100 bpd.",
        "Niger Delta: swamps, creeks, mangroves. Rainy season April-October limits access. Many locations boat-only. Community engagement critical.",
        "Maintenance: ESP every 12 months. Gas compressors every 6 months. Pipeline assessment annually. Corrosion quarterly. All tracked in SAP PM.",
        "Partner reporting: NNPC monthly allocation. TotalEnergies quarterly JV. AENR annual review. Include volumes, costs, HSE stats.",
        "HSE: Level 1 Minor, Level 2 Serious, Level 3 Major (>100bbl spill), Level 4 Critical. Level 3+ reported to NUPRC within 24hrs.",
        "Emergency response: In case of major spill, activate Tier 2 response. Deploy boom within 2 hours. Notify NOSDRA. Begin wildlife assessment. Coordinate with community leaders.",
    ]
    categories = ["general","maintenance","integrity","compliance","operations","security","general","maintenance","commercial","hse","hse"]
    metas = [{"source":f"doc_{i}","category":categories[i % len(categories)]} for i in range(len(docs))]
    ids = [f"kb_{i}" for i in range(len(docs))]
    add_documents(docs, metas, ids)
    logger.info("Knowledge base seeded.")