# database.py
import sqlite3
import json
from datetime import datetime

class DB:
    def __init__(self, path="experiments.db"):
        self.conn = sqlite3.connect(path, check_same_thread=False)
        self._create()

    def _create(self):
        cur = self.conn.cursor()
        cur.execute("""
            CREATE TABLE IF NOT EXISTS experiments (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                prompt TEXT,
                generated TEXT,
                layer_scores TEXT,
                created_at TEXT
            )
        """)
        self.conn.commit()

    def save_experiment(self, prompt, generated, layer_scores):
        cur = self.conn.cursor()
        cur.execute("INSERT INTO experiments (prompt, generated, layer_scores, created_at) VALUES (?, ?, ?, ?)",
                    (prompt, generated, json.dumps(layer_scores), datetime.utcnow().isoformat()))
        self.conn.commit()
        return cur.lastrowid

    def fetch_experiment(self, id):
        cur = self.conn.cursor()
        cur.execute("SELECT id, prompt, generated, layer_scores, created_at FROM experiments WHERE id = ?", (id,))
        row = cur.fetchone()
        if not row:
            return None
        return {
            "id": row[0],
            "prompt": row[1],
            "generated": row[2],
            "layer_scores": json.loads(row[3]),
            "created_at": row[4]
        }

    def list_experiments(self, limit=50):
        cur = self.conn.cursor()
        cur.execute("SELECT id, prompt, created_at FROM experiments ORDER BY id DESC LIMIT ?", (limit,))
        return cur.fetchall()