import os
import pg8000.dbapi as pg

# Load .env manually since run_command environment might miss it
env_path = r"c:\Users\WIN10\Desktop\work\11th-week\法律风险提示-new\市监局-lawRisk-backend\.env"
if os.path.exists(env_path):
    print(f"Loading env from {env_path}")
    with open(env_path, "r", encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if not line or line.startswith("#"): continue
            if "=" in line:
                k, v = line.split("=", 1)
                # Remove quotes if present
                v = v.strip('\'"')
                os.environ[k.strip()] = v
                
from lawrisk.services.licensing_repo import _lic_pg_conn

def apply_migration():
    migration_file = r"c:\Users\WIN10\Desktop\work\11th-week\法律风险提示-new\市监局-lawRisk-backend\data\migration_v2.sql"
    
    with open(migration_file, "r", encoding="utf-8") as f:
        sql = f.read()
        
    print(f"Applying migration from {migration_file}...")
    
    conn = None
    try:
        # Manually create connection if _lic_pg_conn context manager is tricky to use directly with raw SQL script splitting
        # But _lic_pg_conn returns a connection.
        with _lic_pg_conn() as conn:
            cursor = conn.cursor()
            # Split by semilcolon to execute statements one by one if needed, or just execute block
            # pgcrypto usually requires separation, but simple DDL is fine.
            # However, pg8000 might prefer single statements per execute.
            statements = [s.strip() for s in sql.split(';') if s.strip()]
            for stmt in statements:
                print(f"Executing: {stmt[:50]}...")
                cursor.execute(stmt)
            conn.commit()
            print("Migration applied successfully!")
            
    except Exception as e:
        print(f"Migration failed: {e}")
        if conn:
            conn.rollback()

if __name__ == "__main__":
    apply_migration()