from __future__ import annotations

import sqlite3

from backend.database.migrations.runner import Migration


def extend_llm_audit(connection: sqlite3.Connection) -> None:
    columns = {
        str(row["name"])
        for row in connection.execute("PRAGMA table_info(llm_usage)")
    }
    additions = (
        ("role", "TEXT NOT NULL DEFAULT ''"),
        ("prompt_version", "TEXT NOT NULL DEFAULT ''"),
        ("error_code", "TEXT NOT NULL DEFAULT ''"),
        ("input_tokens", "INTEGER NOT NULL DEFAULT 0"),
        ("output_tokens", "INTEGER NOT NULL DEFAULT 0"),
    )
    for name, declaration in additions:
        if name not in columns:
            connection.execute(
                f"ALTER TABLE llm_usage ADD COLUMN {name} {declaration}"
            )


MIGRATION = Migration(
    version="0003",
    name="extend_llm_audit",
    action=extend_llm_audit,
    signature="llm-audit:v1:role,prompt-version,error-code,input-tokens,output-tokens",
)