"""数据库迁移脚本"""
import os
import psycopg2
from dotenv import load_dotenv

load_dotenv()

# 数据库连接配置
DB_HOST = os.getenv("DB_HOST", "8.156.90.138")
DB_PORT = os.getenv("DB_PORT", "5432")
DB_USER = os.getenv("DB_USER", "AigcSpace_test")
DB_PASSWORD = os.getenv("DB_PASSWORD", "pnk6FrPzYKmwHHjM")
DB_NAME = os.getenv("DB_NAME", "aigcspace_test")

def run_migration():
    """执行数据库迁移"""
    conn = psycopg2.connect(
        host=DB_HOST,
        port=DB_PORT,
        user=DB_USER,
        password=DB_PASSWORD,
        dbname=DB_NAME
    )
    conn.autocommit = False
    cur = conn.cursor()

    try:
        print("开始执行数据库迁移...")

        # ============================================
        # 1. models_new 表添加缺失字段
        # ============================================
        migrations = [
            # 添加缺失字段
            ("ALTER TABLE aigcspace.models_new ADD COLUMN IF NOT EXISTS group_name VARCHAR(100)", "添加 group_name 字段"),
            ("ALTER TABLE aigcspace.models_new ADD COLUMN IF NOT EXISTS encrypted_api_key TEXT", "添加 encrypted_api_key 字段"),
            ("ALTER TABLE aigcspace.models_new ADD COLUMN IF NOT EXISTS is_thinking BOOLEAN NOT NULL DEFAULT FALSE", "添加 is_thinking 字段"),
            ("ALTER TABLE aigcspace.models_new ADD COLUMN IF NOT EXISTS source_keys TEXT[]", "添加 source_keys 字段"),
            ("ALTER TABLE aigcspace.models_new ADD COLUMN IF NOT EXISTS normalized_keys TEXT[]", "添加 normalized_keys 字段"),
            ("ALTER TABLE aigcspace.models_new ADD COLUMN IF NOT EXISTS is_local BOOLEAN NOT NULL DEFAULT FALSE", "添加 is_local 字段"),
            ("ALTER TABLE aigcspace.models_new ADD COLUMN IF NOT EXISTS user_id VARCHAR(50)", "添加 user_id 字段"),
            ("ALTER TABLE aigcspace.models_new ADD COLUMN IF NOT EXISTS base_url VARCHAR(500)", "添加 base_url 字段"),
            ("ALTER TABLE aigcspace.models_new ADD COLUMN IF NOT EXISTS local_api_key VARCHAR(500)", "添加 local_api_key 字段"),
            ("ALTER TABLE aigcspace.models_new ADD COLUMN IF NOT EXISTS visibility VARCHAR(20) DEFAULT 'user'", "添加 visibility 字段"),
            # 添加索引
            ("CREATE INDEX IF NOT EXISTS idx_models_new_is_local ON aigcspace.models_new(is_local)", "创建 is_local 索引"),
            ("CREATE INDEX IF NOT EXISTS idx_models_new_group_name ON aigcspace.models_new(group_name)", "创建 group_name 索引"),
        ]

        for sql, desc in migrations:
            try:
                cur.execute(sql)
                print(f"✓ {desc}")
            except Exception as e:
                if "duplicate" in str(e).lower() or "already exists" in str(e).lower():
                    print(f"○ {desc} (已存在，跳过)")
                else:
                    print(f"✗ {desc}: {e}")

        # ============================================
        # 2. 创建可能缺失的表
        # ============================================
        tables = [
            # asr_recognition 表
            """CREATE TABLE IF NOT EXISTS aigcspace.asr_recognition (
                id SERIAL PRIMARY KEY,
                user_id VARCHAR(50) NOT NULL REFERENCES aigcspace.users(id) ON DELETE CASCADE,
                model VARCHAR(100) NOT NULL,
                audio_url VARCHAR(500),
                audio_base64 TEXT,
                language VARCHAR(20),
                enable_itn BOOLEAN DEFAULT FALSE,
                context TEXT,
                result_text TEXT NOT NULL,
                detected_language VARCHAR(20),
                emotion VARCHAR(20),
                duration INTEGER,
                input_tokens INTEGER DEFAULT 0,
                output_tokens INTEGER DEFAULT 0,
                bill NUMERIC(10,4) DEFAULT 0,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
            )""",

            # asr_recognition_v2 表
            """CREATE TABLE IF NOT EXISTS aigcspace.asr_recognition_v2 (
                id SERIAL PRIMARY KEY,
                user_id VARCHAR(50) NOT NULL REFERENCES aigcspace.users(id) ON DELETE CASCADE,
                task_id VARCHAR(100) UNIQUE NOT NULL,
                model VARCHAR(100) NOT NULL,
                file_url VARCHAR(500) NOT NULL,
                status VARCHAR(20) DEFAULT 'PENDING',
                result_text TEXT,
                result_url VARCHAR(500),
                duration INTEGER,
                bill NUMERIC(10,4) DEFAULT 0,
                error_message TEXT,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                completed_at TIMESTAMP
            )""",

            # audio_synthesis_v2 表
            """CREATE TABLE IF NOT EXISTS aigcspace.audio_synthesis_v2 (
                id SERIAL PRIMARY KEY,
                user_id VARCHAR(50) NOT NULL REFERENCES aigcspace.users(id) ON DELETE CASCADE,
                task_id VARCHAR(100) UNIQUE NOT NULL,
                model VARCHAR(100) NOT NULL,
                voice VARCHAR(100) NOT NULL,
                text TEXT NOT NULL,
                audio_url VARCHAR(500),
                status VARCHAR(20) DEFAULT 'PENDING',
                duration NUMERIC(10,2),
                format VARCHAR(20) DEFAULT 'mp3',
                characters INTEGER,
                bill NUMERIC(10,4) DEFAULT 0,
                custom_name VARCHAR(200),
                error_message TEXT,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                completed_at TIMESTAMP
            )""",

            # voice_clone_v2 表
            """CREATE TABLE IF NOT EXISTS aigcspace.voice_clone_v2 (
                id SERIAL PRIMARY KEY,
                user_id VARCHAR(50) NOT NULL REFERENCES aigcspace.users(id) ON DELETE CASCADE,
                task_id VARCHAR(100) UNIQUE NOT NULL,
                voice_id VARCHAR(200),
                target_model VARCHAR(100) NOT NULL,
                prefix VARCHAR(20) NOT NULL,
                voice_name VARCHAR(50),
                audio_url VARCHAR(500),
                status VARCHAR(20) DEFAULT 'PENDING',
                bill NUMERIC(10,4) DEFAULT 0,
                error_message TEXT,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                completed_at TIMESTAMP
            )""",

            # long_text_audio 表
            """CREATE TABLE IF NOT EXISTS aigcspace.long_text_audio (
                id SERIAL PRIMARY KEY,
                user_id VARCHAR(50) NOT NULL REFERENCES aigcspace.users(id) ON DELETE CASCADE,
                task_id VARCHAR(100) UNIQUE NOT NULL,
                model VARCHAR(100) NOT NULL,
                voice VARCHAR(100) NOT NULL,
                text TEXT NOT NULL,
                text_length INTEGER NOT NULL,
                segment_count INTEGER DEFAULT 0,
                segments JSONB DEFAULT '[]',
                audio_url VARCHAR(500),
                status VARCHAR(20) DEFAULT 'PENDING',
                progress INTEGER DEFAULT 0,
                duration NUMERIC(10,2),
                format VARCHAR(20) DEFAULT 'mp3',
                bill NUMERIC(10,4) DEFAULT 0,
                custom_name VARCHAR(200),
                error_message TEXT,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                completed_at TIMESTAMP
            )""",

            # system_voice 表
            """CREATE TABLE IF NOT EXISTS aigcspace.system_voice (
                id SERIAL PRIMARY KEY,
                voice_id VARCHAR(100) UNIQUE NOT NULL,
                name VARCHAR(50) NOT NULL,
                trait VARCHAR(100),
                age VARCHAR(20),
                category VARCHAR(50),
                languages JSONB DEFAULT '[]',
                models JSONB DEFAULT '[]',
                ssml_support BOOLEAN DEFAULT FALSE,
                instruct_support BOOLEAN DEFAULT FALSE,
                timestamp_support BOOLEAN DEFAULT FALSE,
                is_active BOOLEAN DEFAULT TRUE,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
            )""",

            # model_price_new 表
            """CREATE TABLE IF NOT EXISTS aigcspace.model_price_new (
                id SERIAL PRIMARY KEY,
                model_code VARCHAR(200) NOT NULL,
                label VARCHAR(200) NOT NULL,
                tier_min NUMERIC(20,2),
                tier_max NUMERIC(20,2),
                tier_unit VARCHAR(50),
                input_price_original NUMERIC(20,8) NOT NULL DEFAULT 0,
                output_price_original NUMERIC(20,8) NOT NULL DEFAULT 0,
                discount_rate NUMERIC(5,4) NOT NULL DEFAULT 1,
                discount_label VARCHAR(20),
                input_price_discounted NUMERIC(20,8) NOT NULL DEFAULT 0,
                output_price_discounted NUMERIC(20,8) NOT NULL DEFAULT 0,
                currency VARCHAR(10) NOT NULL DEFAULT 'CNY',
                unit VARCHAR(100) NOT NULL,
                display_multiplier INTEGER NOT NULL DEFAULT 1,
                source_url TEXT,
                crawled_at TIMESTAMP NOT NULL,
                is_active BOOLEAN NOT NULL DEFAULT TRUE,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
            )""",

            # crawler_sync_log 表
            """CREATE TABLE IF NOT EXISTS aigcspace.crawler_sync_log (
                id SERIAL PRIMARY KEY,
                crawler_version INTEGER NOT NULL,
                synced_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                model_count INTEGER NOT NULL DEFAULT 0,
                price_count INTEGER NOT NULL DEFAULT 0,
                status VARCHAR(20) NOT NULL DEFAULT 'success',
                error_message TEXT
            )""",

            # user_local_model_permission 表
            """CREATE TABLE IF NOT EXISTS aigcspace.user_local_model_permission (
                id SERIAL PRIMARY KEY,
                user_id VARCHAR(50) NOT NULL REFERENCES aigcspace.users(id) ON DELETE CASCADE,
                model_id INTEGER NOT NULL,
                granted_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                granted_by VARCHAR(50),
                UNIQUE(user_id, model_id)
            )""",
        ]

        print("\n创建缺失的表...")
        for sql in tables:
            table_name = sql.split("aigcspace.")[1].split(" ")[0] if "aigcspace." in sql else "unknown"
            try:
                cur.execute(sql)
                print(f"✓ 表 {table_name} 就绪")
            except Exception as e:
                if "already exists" in str(e).lower():
                    print(f"○ 表 {table_name} 已存在，跳过")
                else:
                    print(f"✗ 表 {table_name}: {e}")

        # 提交事务
        conn.commit()
        print("\n✓ 迁移完成！")

        # 验证
        print("\n验证表结构...")
        cur.execute("""
            SELECT table_name FROM information_schema.tables
            WHERE table_schema = 'aigcspace' ORDER BY table_name
        """)
        tables = cur.fetchall()
        print(f"数据库中共有 {len(tables)} 张表：")
        for t in tables:
            print(f"  - {t[0]}")

    except Exception as e:
        conn.rollback()
        print(f"\n✗ 迁移失败: {e}")
        raise
    finally:
        cur.close()
        conn.close()

if __name__ == "__main__":
    run_migration()