import sqlite3
import os

DB_PATH = 'data/scud_orion_ai.db'


def migrate():
    if not os.path.exists(DB_PATH):
        print(f"База данных {DB_PATH} не найдена. Создание будет выполнено при первом запуске.")
        return

    conn = sqlite3.connect(DB_PATH)
    cursor = conn.cursor()

    cursor.execute("PRAGMA table_info(scud_logs);")
    columns = [row[1] for row in cursor.fetchall()]

    if 'first_activity' not in columns:
        cursor.execute("ALTER TABLE scud_logs ADD COLUMN first_activity TEXT DEFAULT '—';")
        print("[✓] Добавлена колонка 'first_activity'")

    if 'anomaly_flag' not in columns:
        cursor.execute("ALTER TABLE scud_logs ADD COLUMN anomaly_flag TEXT DEFAULT 'NONE';")
        print("[✓] Добавлена колонка 'anomaly_flag'")

    if 'snapshot_time' not in columns:
        cursor.execute("ALTER TABLE scud_logs ADD COLUMN snapshot_time TEXT DEFAULT NULL;")
        print("[✓] Добавлена колонка 'snapshot_time'")

    if 'snapshot_id' not in columns:
        cursor.execute("ALTER TABLE scud_logs ADD COLUMN snapshot_id TEXT DEFAULT NULL;")
        print("[✓] Добавлена колонка 'snapshot_id'")

        # Проставляем номера для ранее сохраненных срезов
        cursor.execute("SELECT DISTINCT snapshot_time FROM scud_logs WHERE snapshot_time IS NOT NULL ORDER BY snapshot_time ASC")
        snaps = cursor.fetchall()
        for idx, (s_time,) in enumerate(snaps, 1):
            s_id = f"{idx:07d}"
            cursor.execute("UPDATE scud_logs SET snapshot_id = ? WHERE snapshot_time = ?", (s_id, s_time))

    conn.commit()
    conn.close()
    print("[✓] Миграция схемы БД успешно завершена!")


if __name__ == "__main__":
    migrate()