import sqlite3
import time

DB_NAME = "chrono_ronin.db"

def get_conn():
    conn = sqlite3.connect(DB_NAME, check_same_thread=False)
    conn.row_factory = sqlite3.Row
    return conn

def init_db():
    with get_conn() as conn:
        conn.execute("""
            CREATE TABLE IF NOT EXISTS agents (
                id TEXT PRIMARY KEY, role TEXT, bank_account REAL,
                position REAL, entry_price REAL, age INTEGER
            )
        """)
        conn.execute("""
            CREATE TABLE IF NOT EXISTS trades (
                id INTEGER PRIMARY KEY AUTOINCREMENT, timestamp REAL,
                agent_id TEXT, action TEXT, quantity REAL,
                price REAL, fees REAL
            )
        """)
        conn.execute("""
            CREATE TABLE IF NOT EXISTS metrics (key TEXT PRIMARY KEY, value REAL)
        """)
        # CORRECTION ICI
        conn.execute("""
            CREATE TABLE IF NOT EXISTS memory_cells (
                regime TEXT, rsi_bucket INTEGER,
                wins REAL, total REAL, last_update REAL,
                PRIMARY KEY (regime, rsi_bucket)
            )
        """)
        conn.execute("""
            CREATE TABLE IF NOT EXISTS trade_contexts (
                id INTEGER PRIMARY KEY AUTOINCREMENT, timestamp REAL,
                regime TEXT, rsi_bucket INTEGER, sentiment_bucket INTEGER,
                entry_price REAL, exit_price REAL, pnl_pct REAL,
                duration_min INTEGER, exit_reason TEXT
            )
        """)
        conn.commit()

def save_agent(agent):
    with get_conn() as conn:
        conn.execute("""INSERT OR REPLACE INTO agents
            (id, role, bank_account, position, entry_price, age)
            VALUES (?, ?, ?, ?, ?, ?)""",
            (agent.id, agent.role, agent.bank_account, agent.position, agent.entry_price, agent.age))
        conn.commit()

def delete_agent(agent_id):
    with get_conn() as conn:
        conn.execute("DELETE FROM agents WHERE id = ?", (agent_id,))
        conn.commit()

def load_agents():
    with get_conn() as conn:
        return [dict(r) for r in conn.execute("SELECT * FROM agents").fetchall()]

def save_trade(agent_id, action, quantity, price, fees):
    with get_conn() as conn:
        conn.execute("""INSERT INTO trades (timestamp, agent_id, action, quantity, price, fees)
            VALUES (?, ?, ?, ?, ?, ?)""",
            (time.time(), agent_id, action, quantity, price, fees))
        conn.commit()

def save_metrics(metrics_dict):
    with get_conn() as conn:
        for k, v in metrics_dict.items():
            conn.execute("INSERT OR REPLACE INTO metrics (key, value) VALUES (?, ?)", (k, v))
        conn.commit()

def load_metrics():
    with get_conn() as conn:
        return {r["key"]: r["value"] for r in conn.execute("SELECT * FROM metrics").fetchall()}

def load_memory_cells():
    with get_conn() as conn:
        try:
            return [dict(r) for r in conn.execute("SELECT * FROM memory_cells").fetchall()]
        except Exception:
            return []

def save_memory_cell(regime, rsi_bucket, wins, total, last_update):
    with get_conn() as conn:
        conn.execute("""INSERT OR REPLACE INTO memory_cells
            (regime, rsi_bucket, wins, total, last_update)
            VALUES (?, ?, ?, ?, ?)""",
            (regime, rsi_bucket, wins, total, last_update))
        conn.commit()

def save_trade_context(regime, rsi_bucket, sentiment_bucket, entry_price, exit_price, pnl_pct, duration_min, exit_reason):
    with get_conn() as conn:
        conn.execute("""INSERT INTO trade_contexts
            (timestamp, regime, rsi_bucket, sentiment_bucket, entry_price, exit_price, pnl_pct, duration_min, exit_reason)
            VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)""",
            (time.time(), regime, rsi_bucket, sentiment_bucket, entry_price, exit_price, pnl_pct, duration_min, exit_reason))
        conn.commit()

def clear_db():
    with get_conn() as conn:
        conn.execute("DELETE FROM agents")
        conn.execute("DELETE FROM trades")
        conn.execute("DELETE FROM metrics")
        conn.execute("DELETE FROM trade_contexts")
        # On garde memory_cells pour que l'IA n'oublie pas
        conn.commit()