from __future__ import annotations

import json
import logging
import sqlite3
from datetime import datetime, timedelta, timezone
from typing import List, Optional

logger = logging.getLogger(__name__)


class MeshDB:
    def __init__(self, db_path: str):
        self.db_path = db_path
        self.conn = sqlite3.connect(db_path, check_same_thread=False)
        self.conn.row_factory = sqlite3.Row
        self._create_tables()

    def _create_tables(self):
        self.conn.executescript("""
            CREATE TABLE IF NOT EXISTS nodes (
                node_id TEXT PRIMARY KEY,
                hostname TEXT,
                agent_url TEXT,
                status TEXT DEFAULT 'unknown',
                hardware_json TEXT,
                tags_json TEXT,
                registered_at TEXT,
                last_heartbeat TEXT
            );

            CREATE TABLE IF NOT EXISTS health_history (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                node_id TEXT NOT NULL,
                timestamp TEXT NOT NULL,
                cpu_pct REAL,
                memory_used_pct REAL,
                load_1m REAL,
                load_5m REAL,
                load_15m REAL,
                gpu_utilizations_json TEXT,
                disk_used_pct REAL,
                uptime_s REAL,
                FOREIGN KEY (node_id) REFERENCES nodes(node_id)
            );

            CREATE TABLE IF NOT EXISTS jobs (
                job_id TEXT PRIMARY KEY,
                job_type TEXT,
                command TEXT,
                args_json TEXT,
                status TEXT DEFAULT 'pending',
                assigned_node TEXT,
                submitted_at TEXT,
                started_at TEXT,
                completed_at TEXT,
                result_json TEXT,
                error TEXT,
                priority INTEGER DEFAULT 0
            );

            CREATE TABLE IF NOT EXISTS atlas_sync_log (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                node_id TEXT NOT NULL,
                atlas_db_name TEXT,
                data_points_count INTEGER,
                invariants_count INTEGER,
                synced_at TEXT
            );

            CREATE TABLE IF NOT EXISTS jepa_training_log (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                training_samples INTEGER,
                epochs INTEGER,
                final_loss REAL,
                mean_error REAL,
                duration_s REAL,
                trained_at TEXT
            );

            CREATE INDEX IF NOT EXISTS idx_health_node_ts
                ON health_history(node_id, timestamp);
            CREATE INDEX IF NOT EXISTS idx_jobs_status
                ON jobs(status);
        """)
        self.conn.commit()

    def _now_iso(self) -> str:
        return datetime.now(timezone.utc).isoformat()

    def _row_to_dict(self, row: sqlite3.Row | None) -> Optional[dict]:
        if row is None:
            return None
        return dict(row)

    # -- Nodes --

    def upsert_node(self, node_id: str, hostname: str, agent_url: str,
                    hardware_dict: dict | None = None, tags: list | None = None):
        now = self._now_iso()
        self.conn.execute(
            """INSERT INTO nodes (node_id, hostname, agent_url, status, hardware_json,
                                  tags_json, registered_at, last_heartbeat)
               VALUES (?, ?, ?, 'online', ?, ?, ?, ?)
               ON CONFLICT(node_id) DO UPDATE SET
                   hostname=excluded.hostname,
                   agent_url=excluded.agent_url,
                   hardware_json=excluded.hardware_json,
                   tags_json=excluded.tags_json,
                   last_heartbeat=excluded.last_heartbeat""",
            (node_id, hostname, agent_url,
             json.dumps(hardware_dict) if hardware_dict else None,
             json.dumps(tags) if tags else None,
             now, now)
        )
        self.conn.commit()

    def update_node_status(self, node_id: str, status: str):
        self.conn.execute(
            "UPDATE nodes SET status=? WHERE node_id=?",
            (status, node_id)
        )
        self.conn.commit()

    def update_node_heartbeat(self, node_id: str, timestamp: str):
        self.conn.execute(
            "UPDATE nodes SET last_heartbeat=? WHERE node_id=?",
            (timestamp, node_id)
        )
        self.conn.commit()

    def get_node(self, node_id: str) -> Optional[dict]:
        row = self.conn.execute(
            "SELECT * FROM nodes WHERE node_id=?", (node_id,)
        ).fetchone()
        d = self._row_to_dict(row)
        if d:
            d["hardware"] = json.loads(d.pop("hardware_json")) if d.get("hardware_json") else None
            d["tags"] = json.loads(d.pop("tags_json")) if d.get("tags_json") else None
        return d

    def get_all_nodes(self) -> List[dict]:
        rows = self.conn.execute(
            "SELECT * FROM nodes WHERE status != 'archived' ORDER BY node_id"
        ).fetchall()
        result = []
        for row in rows:
            d = dict(row)
            d["hardware"] = json.loads(d.pop("hardware_json")) if d.get("hardware_json") else None
            d["tags"] = json.loads(d.pop("tags_json")) if d.get("tags_json") else None
            result.append(d)
        return result

    def get_all_nodes_raw(self) -> List[dict]:
        """Return ALL nodes including archived, without transforming JSON fields."""
        rows = self.conn.execute(
            "SELECT * FROM nodes ORDER BY node_id"
        ).fetchall()
        return [dict(row) for row in rows]

    def delete_node(self, node_id: str):
        self.conn.execute("DELETE FROM nodes WHERE node_id=?", (node_id,))
        self.conn.commit()

    # -- Health History --

    def record_health(self, node_id: str, health_dict: dict):
        self.conn.execute(
            """INSERT INTO health_history
               (node_id, timestamp, cpu_pct, memory_used_pct, load_1m, load_5m, load_15m,
                gpu_utilizations_json, disk_used_pct, uptime_s)
               VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""",
            (node_id,
             health_dict.get("timestamp", self._now_iso()),
             health_dict.get("cpu_pct"),
             health_dict.get("memory_used_pct"),
             health_dict.get("load_1m"),
             health_dict.get("load_5m"),
             health_dict.get("load_15m"),
             json.dumps(health_dict.get("gpu_utilizations")) if health_dict.get("gpu_utilizations") else None,
             health_dict.get("disk_used_pct"),
             health_dict.get("uptime_s"))
        )
        self.conn.commit()

    def get_health_history(self, node_id: str, limit: int = 100) -> List[dict]:
        rows = self.conn.execute(
            """SELECT * FROM health_history WHERE node_id=?
               ORDER BY timestamp DESC LIMIT ?""",
            (node_id, limit)
        ).fetchall()
        result = []
        for row in rows:
            d = dict(row)
            d["gpu_utilizations"] = json.loads(d.pop("gpu_utilizations_json")) if d.get("gpu_utilizations_json") else None
            result.append(d)
        return result

    # -- Jobs --

    def create_job(self, job_id: str, job_type: str, command: str,
                   args: dict | None = None, priority: int = 0) -> dict:
        now = self._now_iso()
        self.conn.execute(
            """INSERT INTO jobs (job_id, job_type, command, args_json, status,
                                submitted_at, priority)
               VALUES (?, ?, ?, ?, 'pending', ?, ?)""",
            (job_id, job_type, command,
             json.dumps(args) if args else None,
             now, priority)
        )
        self.conn.commit()
        return self.get_job(job_id)

    def update_job(self, job_id: str, **kwargs):
        if not kwargs:
            return
        allowed = {"status", "assigned_node", "started_at", "completed_at",
                    "result_json", "error", "priority"}
        # Handle result dict -> json
        if "result" in kwargs:
            kwargs["result_json"] = json.dumps(kwargs.pop("result"))
        if "args" in kwargs:
            kwargs["args_json"] = json.dumps(kwargs.pop("args"))
        cols = {k: v for k, v in kwargs.items() if k in allowed}
        if not cols:
            return
        set_clause = ", ".join(f"{k}=?" for k in cols)
        values = list(cols.values()) + [job_id]
        self.conn.execute(
            f"UPDATE jobs SET {set_clause} WHERE job_id=?", values
        )
        self.conn.commit()

    def get_job(self, job_id: str) -> Optional[dict]:
        row = self.conn.execute(
            "SELECT * FROM jobs WHERE job_id=?", (job_id,)
        ).fetchone()
        d = self._row_to_dict(row)
        if d:
            d["args"] = json.loads(d.pop("args_json")) if d.get("args_json") else {}
            raw = d.pop("result_json", None)
            d["result"] = json.loads(raw) if raw else None
        return d

    def get_jobs(self, status: str | None = None, limit: int = 50) -> List[dict]:
        if status:
            rows = self.conn.execute(
                "SELECT * FROM jobs WHERE status=? ORDER BY priority DESC, submitted_at LIMIT ?",
                (status, limit)
            ).fetchall()
        else:
            rows = self.conn.execute(
                "SELECT * FROM jobs ORDER BY priority DESC, submitted_at LIMIT ?",
                (limit,)
            ).fetchall()
        result = []
        for row in rows:
            d = dict(row)
            d["args"] = json.loads(d.pop("args_json")) if d.get("args_json") else {}
            d["result"] = json.loads(d.pop("result_json")) if d.get("result_json") else None
            result.append(d)
        return result

    # -- Atlas Sync Log --

    def record_atlas_sync(self, node_id: str, atlas_db_name: str,
                          data_points_count: int, invariants_count: int):
        self.conn.execute(
            """INSERT INTO atlas_sync_log
               (node_id, atlas_db_name, data_points_count, invariants_count, synced_at)
               VALUES (?, ?, ?, ?, ?)""",
            (node_id, atlas_db_name, data_points_count, invariants_count, self._now_iso())
        )
        self.conn.commit()

    # -- JEPA Training Log --

    def record_jepa_training(self, training_samples: int, epochs: int,
                             final_loss: float, mean_error: float, duration_s: float):
        self.conn.execute(
            """INSERT INTO jepa_training_log
               (training_samples, epochs, final_loss, mean_error, duration_s, trained_at)
               VALUES (?, ?, ?, ?, ?, ?)""",
            (training_samples, epochs, final_loss, mean_error, duration_s, self._now_iso())
        )
        self.conn.commit()

    def get_jepa_training_log(self, limit: int = 20) -> List[dict]:
        rows = self.conn.execute(
            "SELECT * FROM jepa_training_log ORDER BY trained_at DESC LIMIT ?",
            (limit,)
        ).fetchall()
        return [dict(row) for row in rows]

    # -- Maintenance --

    def prune_health_history(self, max_age_hours: int = 168):
        cutoff = (datetime.now(timezone.utc) - timedelta(hours=max_age_hours)).isoformat()
        cursor = self.conn.execute(
            "DELETE FROM health_history WHERE timestamp < ?", (cutoff,)
        )
        self.conn.commit()
        deleted = cursor.rowcount
        if deleted > 0:
            logger.info("Pruned %d health records older than %d hours", deleted, max_age_hours)
        return deleted
