from __future__ import annotations

import sqlite3

from backend.database.migrations.runner import Migration


def upgrade(connection: sqlite3.Connection) -> None:
    connection.execute(
        """
        CREATE TABLE llm_models (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            display_name TEXT NOT NULL,
            display_name_key TEXT NOT NULL UNIQUE,
            base_url TEXT NOT NULL,
            model_identifier TEXT NOT NULL,
            encrypted_api_key TEXT NOT NULL,
            created_at TEXT NOT NULL,
            updated_at TEXT NOT NULL,
            updated_by INTEGER NOT NULL REFERENCES users(id) ON DELETE RESTRICT
        )
        """
    )
    connection.execute(
        """
        CREATE TABLE llm_configuration (
            id INTEGER PRIMARY KEY CHECK (id = 1),
            primary_model_id INTEGER REFERENCES llm_models(id) ON DELETE RESTRICT,
            fallback_model_id INTEGER REFERENCES llm_models(id) ON DELETE RESTRICT,
            updated_at TEXT,
            updated_by INTEGER REFERENCES users(id) ON DELETE RESTRICT,
            CHECK (
                fallback_model_id IS NULL
                OR primary_model_id IS NULL
                OR fallback_model_id <> primary_model_id
            )
        )
        """
    )
    connection.execute("INSERT INTO llm_configuration (id) VALUES (1)")


def downgrade(connection: sqlite3.Connection) -> None:
    connection.execute("DROP TABLE llm_configuration")
    connection.execute("DROP TABLE llm_models")


MIGRATION = Migration(
    version=2,
    name="create_model_pool",
    signature="model-pool:v1:models-primary-fallback-encrypted-key",
    upgrade=upgrade,
    downgrade=downgrade,
)