#!/usr/bin/env python3
"""
Migration script to create initial users in the database from environment variables
Run this once after creating the users table
"""

import os
import psycopg2
import bcrypt
from dotenv import load_dotenv
from datetime import datetime

load_dotenv()

def get_connection():
    return psycopg2.connect(
        host=os.getenv("POSTGRES_HOST", "postgres"),
        port=os.getenv("POSTGRES_PORT", "5432"),
        dbname=os.getenv("POSTGRES_DB"),
        user=os.getenv("POSTGRES_USER"),
        password=os.getenv("POSTGRES_PASSWORD"),
        options='-c client_encoding=UTF8'
    )

def hash_password(password):
    return bcrypt.hashpw(password.encode('utf-8'), bcrypt.gensalt()).decode('utf-8')

def migrate_users():
    conn = get_connection()
    cur = conn.cursor()
    
    try:
        # Check if users already exist
        cur.execute("SELECT COUNT(*) FROM users")
        user_count = cur.fetchone()[0]
        
        if user_count > 0:
            print(f"Users table already contains {user_count} users. Skipping migration.")
            return
        
        # Get admin users from environment
        admin1_user = os.getenv("ADMIN1_USERNAME")
        admin1_pass = os.getenv("ADMIN1_PASSWORD")
        admin2_user = os.getenv("ADMIN2_USERNAME")
        admin2_pass = os.getenv("ADMIN2_PASSWORD")
        
        if not all([admin1_user, admin1_pass, admin2_user, admin2_pass]):
            print("ERROR: Admin credentials not found in environment variables!")
            return
        
        # Insert admin users
        users = [
            (admin1_user, hash_password(admin1_pass), f"{admin1_user}@v2-admin.local"),
            (admin2_user, hash_password(admin2_pass), f"{admin2_user}@v2-admin.local")
        ]
        
        for username, password_hash, email in users:
            cur.execute("""
                INSERT INTO users (username, password_hash, email, totp_enabled, created_at)
                VALUES (%s, %s, %s, %s, %s)
            """, (username, password_hash, email, False, datetime.now()))
            print(f"Created user: {username}")
        
        conn.commit()
        print("\nMigration completed successfully!")
        print("Users can now log in with their existing credentials.")
        print("They can enable 2FA from their profile page.")
        
    except Exception as e:
        conn.rollback()
        print(f"ERROR during migration: {e}")
    finally:
        cur.close()
        conn.close()

if __name__ == "__main__":
    print("Starting user migration...")
    migrate_users()