import sqlite3
from pathlib import Path

CACHE_DB_PATH = Path("agentm/data/agentM.db")
CACHE_DB_PATH.parent.mkdir(parents=True, exist_ok=True)

def get_db_conn():
    conn = sqlite3.connect(CACHE_DB_PATH)
    conn.row_factory = sqlite3.Row
    return conn

def initialize_database():
    with get_db_conn() as conn:
        conn.execute("PRAGMA foreign_keys = ON;")

        # Game metadata table
        conn.execute("""
            CREATE TABLE IF NOT EXISTS roms (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                title TEXT NOT NULL,
                rom_file TEXT NOT NULL UNIQUE,
                game_id TEXT NOT NULL UNIQUE,
                sha256 TEXT,
                difficulty_min INTEGER,
                difficulty_max INTEGER,
                characters TEXT,
                keywords TEXT,
                verified BOOLEAN NOT NULL DEFAULT 0,
                verified_at TEXT
            );
        """)

        # Agent definitions linked to game_id
        conn.execute("""
            CREATE TABLE IF NOT EXISTS agents (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                name TEXT NOT NULL,
                game_id TEXT NOT NULL,
                agent_type TEXT DEFAULT 'PPO',
                framework TEXT DEFAULT 'SB3',
                is_imitation BOOLEAN DEFAULT 0,
                dataset_path TEXT,
                config_json TEXT NOT NULL,
                notes TEXT,
                created_at TEXT DEFAULT CURRENT_TIMESTAMP,
                last_updated TEXT DEFAULT CURRENT_TIMESTAMP,
                FOREIGN KEY (game_id) REFERENCES roms(game_id)
                    ON DELETE CASCADE ON UPDATE CASCADE
            );
        """)

        # Run definitions linked to agents
        conn.execute("""
            CREATE TABLE IF NOT EXISTS runs (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                agent_id INTEGER NOT NULL,
                name TEXT NOT NULL,
                config_yaml TEXT NOT NULL,
                notes TEXT,
                pending BOOLEAN NOT NULL DEFAULT 1,
                created_at TEXT DEFAULT CURRENT_TIMESTAMP,
                updated_at TEXT DEFAULT CURRENT_TIMESTAMP,
                FOREIGN KEY (agent_id) REFERENCES agents(id)
                    ON DELETE CASCADE ON UPDATE CASCADE
            );
        """)

        # Model definitions (track training progress, lineage, resume state)
        conn.execute("""
            CREATE TABLE IF NOT EXISTS models (
                id INTEGER PRIMARY KEY AUTOINCREMENT,

                agent_id INTEGER NOT NULL,
                parent_model_id INTEGER,

                name TEXT NOT NULL,
                notes TEXT,

                total_steps_planned INTEGER NOT NULL,
                total_steps_completed INTEGER NOT NULL DEFAULT 0,
                average_reward REAL,

                current_learning_rate REAL,
                current_clip_range REAL,
                num_envs INTEGER DEFAULT 1,  -- New field added here

                config_patch_yaml TEXT,
                checkpoint_path TEXT,
                is_frozen BOOLEAN NOT NULL DEFAULT 0,

                status TEXT NOT NULL DEFAULT 'pending'
                    CHECK(status IN ('pending', 'training', 'paused', 'completed', 'failed')),

                created_at TEXT DEFAULT CURRENT_TIMESTAMP,
                updated_at TEXT DEFAULT CURRENT_TIMESTAMP,

                FOREIGN KEY (agent_id) REFERENCES agents(id)
                    ON DELETE CASCADE ON UPDATE CASCADE,
                FOREIGN KEY (parent_model_id) REFERENCES models(id)
                    ON DELETE CASCADE ON UPDATE CASCADE
            );
        """)


        conn.commit()