"""Run database migration for training_status fields."""
import psycopg2
import os

# Read password from .env file
password = ""
try:
    with open(".env") as f:
        for line in f:
            if line.startswith("DB_PASSWORD="):
                password = line.strip().split("=", 1)[1].strip('"').strip("'")
                break
except Exception as e:
    print(f"Error reading .env: {e}")

print(f"Password found: {bool(password)}")

conn = psycopg2.connect(
    host="192.168.68.31",
    port=5432,
    database="docmaster",
    user="docmaster",
    password=password
)
conn.autocommit = True
cur = conn.cursor()

# Add training_status column
try:
    cur.execute("ALTER TABLE training_datasets ADD COLUMN training_status VARCHAR(20) DEFAULT NULL")
    print("Added training_status column")
except Exception as e:
    print(f"training_status: {e}")

# Add active_training_task_id column
try:
    cur.execute("ALTER TABLE training_datasets ADD COLUMN active_training_task_id UUID DEFAULT NULL")
    print("Added active_training_task_id column")
except Exception as e:
    print(f"active_training_task_id: {e}")

# Create indexes
try:
    cur.execute("CREATE INDEX IF NOT EXISTS idx_training_datasets_training_status ON training_datasets(training_status)")
    print("Created training_status index")
except Exception as e:
    print(f"index training_status: {e}")

try:
    cur.execute("CREATE INDEX IF NOT EXISTS idx_training_datasets_active_training_task_id ON training_datasets(active_training_task_id)")
    print("Created active_training_task_id index")
except Exception as e:
    print(f"index active_training_task_id: {e}")

# Update existing datasets that have been used in completed training tasks to trained status
try:
    cur.execute("""
        UPDATE training_datasets d
        SET status = 'trained'
        WHERE d.status = 'ready'
        AND EXISTS (
            SELECT 1 FROM training_tasks t
            WHERE t.dataset_id = d.dataset_id
            AND t.status = 'completed'
        )
    """)
    print(f"Updated {cur.rowcount} datasets to trained status")
except Exception as e:
    print(f"update status: {e}")

cur.close()
conn.close()
print("Migration complete!")