import sqlite3
import json


DB_FILE = 'predictions.db'

def init_db():
    conn = sqlite3.connect(DB_FILE)
    c = conn.cursor()
    c.execute('''
        CREATE TABLE IF NOT EXISTS predictions (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            input TEXT,
            predicted_class TEXT,
            probability REAL,
            timestamp DATETIME DEFAULT CURRENT_TIMESTAMP
        )
    ''')
    conn.commit()
    conn.close()

def insert_prediction(input_data, result):
    conn = sqlite3.connect(DB_FILE)
    c = conn.cursor()
    c.execute('''
        INSERT INTO predictions (input, predicted_class, probability)
        VALUES (?, ?, ?)
    ''', (json.dumps(input_data), result['class'], result['probability']))
    conn.commit()
    conn.close()

def get_latest_prediction():
    conn = sqlite3.connect(DB_FILE)
    c = conn.cursor()
    c.execute('''
        SELECT input, predicted_class, probability, timestamp
        FROM predictions
        ORDER BY id DESC
        LIMIT 2
    ''')
    rows = c.fetchall()
    conn.close()
    
    if rows:
        return [
            {
                "input": json.loads(row[0]),
                "class": row[1],
                "probability": row[2],
                "timestamp": row[3]
            }
            for row in rows
        ]
    return []


def get_all_predictions():
    conn = sqlite3.connect(DB_FILE)
    c = conn.cursor()
    c.execute('''
        SELECT input, predicted_class, probability, timestamp
        FROM predictions
    ''')
    rows = c.fetchall()
    conn.close()

    results = []
    for row in rows:
        results.append({
            "input": json.loads(row[0]),
            "class": row[1],
            "probability": row[2],
            "timestamp": row[3]
        })

    return results