import sqlite3
import os

DB_PATH = "/home/maxi/Projekt/grimhold.waiswuaschd.de/data/grimhold.db"

def migrate():
    conn = sqlite3.connect(DB_PATH)
    cur = conn.cursor()
    
    # Check users table columns
    cur.execute("PRAGMA table_info(users)")
    cols = [r[1] for r in cur.fetchall()]
    print("Existing users cols:", cols)

    if "shield_until" not in cols:
        print("Adding shield_until column to users...")
        cur.execute("ALTER TABLE users ADD COLUMN shield_until REAL NOT NULL DEFAULT 0")

    # Create user_army table if not exists
    cur.execute("""
    CREATE TABLE IF NOT EXISTS user_army (
        user_id INTEGER NOT NULL,
        unit_type TEXT NOT NULL,
        pool_count INTEGER NOT NULL DEFAULT 0,
        defense_count INTEGER NOT NULL DEFAULT 0,
        PRIMARY KEY (user_id, unit_type),
        FOREIGN KEY (user_id) REFERENCES users(id) ON DELETE CASCADE
    )
    """)

    # Seed default starter army for existing users if empty
    cur.execute("SELECT id FROM users")
    users = [r[0] for r in cur.fetchall()]
    for uid in users:
        cur.execute("SELECT COUNT(*) FROM user_army WHERE user_id = ?", (uid,))
        if cur.fetchone()[0] == 0:
            print(f"Seeding starter army for user {uid}...")
            # Starter units: 8 scouts, 3 tanks, 1 defense scout
            cur.execute("INSERT INTO user_army (user_id, unit_type, pool_count, defense_count) VALUES (?, 'scout', 8, 1)", (uid,))
            cur.execute("INSERT INTO user_army (user_id, unit_type, pool_count, defense_count) VALUES (?, 'tank', 3, 1)", (uid,))
            cur.execute("INSERT INTO user_army (user_id, unit_type, pool_count, defense_count) VALUES (?, 'crusher', 0, 0)", (uid,))
            cur.execute("INSERT INTO user_army (user_id, unit_type, pool_count, defense_count) VALUES (?, 'infiltrator', 0, 0)", (uid,))
            cur.execute("INSERT INTO user_army (user_id, unit_type, pool_count, defense_count) VALUES (?, 'gargoyle', 0, 0)", (uid,))

    conn.commit()
    conn.close()
    print("DB migration completed successfully!")

if __name__ == "__main__":
    migrate()