"""
데이터베이스 마이그레이션 실행 스크립트
users 테이블에 pro_trial_count, upload_count 컬럼 추가
신규 회원은 파일 업로드 2회 + Pro 무료체험 2번 제공
"""
import sys
import os
from dotenv import load_dotenv
import psycopg2

# 환경변수 로딩
load_dotenv()

# 프로젝트 루트 경로 추가
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

def get_db_connection():
    """데이터베이스 연결 (pandas 없이 직접 연결)"""
    DB_CONFIG = {
        'dbname': os.getenv('DB_NAME', 'bridge_db'),
        'user': os.getenv('DB_USER', 'postgres'),
        'password': os.getenv('DB_PASSWORD', 'Maruta82'),
        'host': os.getenv('DB_HOST', 'infrasmart.co.kr'),
        'port': os.getenv('DB_PORT', '5433')
    }
    return psycopg2.connect(**DB_CONFIG)

def run_migration():
    """pro_trial_count, upload_count 컬럼 추가 마이그레이션 실행"""
    conn = None
    try:
        print("데이터베이스 연결 중...")
        conn = get_db_connection()
        cur = conn.cursor()
        
        # 1. pro_trial_count 컬럼 추가
        cur.execute("""
            SELECT column_name 
            FROM information_schema.columns 
            WHERE table_name = 'users' 
            AND column_name = 'pro_trial_count'
        """)
        
        if cur.fetchone():
            print("✓ pro_trial_count 컬럼이 이미 존재합니다.")
        else:
            print("pro_trial_count 컬럼 추가 중...")
            cur.execute("""
                ALTER TABLE users 
                ADD COLUMN pro_trial_count INTEGER DEFAULT 2
            """)
            conn.commit()
            print("✓ pro_trial_count 컬럼이 성공적으로 추가되었습니다.")
        
        # 2. upload_count 컬럼 추가
        cur.execute("""
            SELECT column_name 
            FROM information_schema.columns 
            WHERE table_name = 'users' 
            AND column_name = 'upload_count'
        """)
        
        if cur.fetchone():
            print("✓ upload_count 컬럼이 이미 존재합니다.")
        else:
            print("upload_count 컬럼 추가 중...")
            cur.execute("""
                ALTER TABLE users 
                ADD COLUMN upload_count INTEGER DEFAULT 2
            """)
            conn.commit()
            print("✓ upload_count 컬럼이 성공적으로 추가되었습니다.")
        
        # 3. plan 컬럼 추가
        cur.execute("""
            SELECT column_name 
            FROM information_schema.columns 
            WHERE table_name = 'users' 
            AND column_name = 'plan'
        """)
        
        if cur.fetchone():
            print("✓ plan 컬럼이 이미 존재합니다.")
        else:
            print("plan 컬럼 추가 중...")
            cur.execute("""
                ALTER TABLE users 
                ADD COLUMN plan VARCHAR(20) DEFAULT 'free'
            """)
            conn.commit()
            print("✓ plan 컬럼이 성공적으로 추가되었습니다.")
            
            # 기존 사용자들은 free 플랜으로 설정
            cur.execute("""
                UPDATE users 
                SET plan = 'free' 
                WHERE plan IS NULL
            """)
            conn.commit()
            print("✓ 기존 사용자들을 free 플랜으로 설정했습니다.")
        
        # 4. is_admin 컬럼 추가
        cur.execute("""
            SELECT column_name 
            FROM information_schema.columns 
            WHERE table_name = 'users' 
            AND column_name = 'is_admin'
        """)
        
        if cur.fetchone():
            print("✓ is_admin 컬럼이 이미 존재합니다.")
        else:
            print("is_admin 컬럼 추가 중...")
            cur.execute("""
                ALTER TABLE users 
                ADD COLUMN is_admin BOOLEAN DEFAULT FALSE
            """)
            conn.commit()
            print("✓ is_admin 컬럼이 성공적으로 추가되었습니다.")
            
            # 기존 사용자들은 관리자가 아님
            cur.execute("""
                UPDATE users 
                SET is_admin = FALSE 
                WHERE is_admin IS NULL
            """)
            conn.commit()
            print("✓ 기존 사용자들을 일반 사용자로 설정했습니다.")
        
        # 모든 기존 사용자에게 업로드 횟수 2회 부여
        cur.execute("""
            UPDATE users 
            SET upload_count = 2
            WHERE upload_count IS NULL OR upload_count != 2
        """)
        all_updated_count = cur.rowcount
        if all_updated_count > 0:
            conn.commit()
            print(f"✓ 모든 기존 사용자 {all_updated_count}명에게 파일 업로드 2회를 부여했습니다.")
        else:
            print("✓ 모든 기존 사용자가 이미 업로드 횟수 2회를 가지고 있습니다.")
        
        # 모든 기존 사용자에게 Pro 무료체험 2회 부여
        cur.execute("""
            UPDATE users 
            SET pro_trial_count = 2
            WHERE pro_trial_count IS NULL OR pro_trial_count != 2
        """)
        pro_updated_count = cur.rowcount
        if pro_updated_count > 0:
            conn.commit()
            print(f"✓ 모든 기존 사용자 {pro_updated_count}명에게 Pro 무료체험 2회를 부여했습니다.")
        else:
            print("✓ 모든 기존 사용자가 이미 Pro 무료체험 2회를 가지고 있습니다.")
        
        # 테스트: 컬럼 확인
        cur.execute("SELECT upload_count, pro_trial_count FROM users LIMIT 1")
        result = cur.fetchone()
        if result:
            print(f"✓ 테스트 조회 성공: upload_count = {result[0]}, pro_trial_count = {result[1]}")
        
        print("\n✅ 마이그레이션이 완료되었습니다!")
        
    except Exception as e:
        print(f"❌ 마이그레이션 실패: {e}")
        import traceback
        traceback.print_exc()
        if conn:
            conn.rollback()
    finally:
        if conn:
            cur.close()
            conn.close()

if __name__ == '__main__':
    run_migration()