#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
執行 API Usage Stats 資料表修復 Migration

Usage:
    python migrations/fix_api_usage_stats.py
"""

import sys
from pathlib import Path

# 添加專案根目錄到 Python 路徑
project_root = Path(__file__).parent.parent
sys.path.insert(0, str(project_root))

from app import create_app, db
from sqlalchemy import text


def run_migration():
    """執行資料表結構修復"""
    app = create_app()

    with app.app_context():
        print("=" * 60)
        print("API Usage Stats 資料表結構修復")
        print("=" * 60)

        try:
            # 1. 檢查當前結構
            print("\n[1/8] 檢查當前資料表結構...")
            result = db.session.execute(text('SHOW COLUMNS FROM dt_api_usage_stats'))
            current_columns = {row[0]: row for row in result}
            print(f"   ✓ 當前欄位數量: {len(current_columns)}")

            # 2. 備份現有資料
            print("\n[2/8] 建立資料備份...")
            db.session.execute(text('''
                CREATE TABLE IF NOT EXISTS dt_api_usage_stats_backup_20251001
                AS SELECT * FROM dt_api_usage_stats
            '''))
            db.session.commit()

            backup_count = db.session.execute(
                text('SELECT COUNT(*) FROM dt_api_usage_stats_backup_20251001')
            ).scalar()
            print(f"   ✓ 已備份 {backup_count} 筆記錄")

            # 3. 修改欄位名稱：api_name → api_endpoint
            if 'api_name' in current_columns:
                print("\n[3/8] 重新命名 api_name → api_endpoint...")
                db.session.execute(text('''
                    ALTER TABLE dt_api_usage_stats
                    CHANGE COLUMN api_name api_endpoint VARCHAR(200) NOT NULL COMMENT 'API端點'
                '''))
                db.session.commit()
                print("   ✓ 已重新命名 api_name → api_endpoint")
            else:
                print("\n[3/8] 跳過（api_name 已不存在或已是 api_endpoint）")

            # 4. 新增 prompt_tokens 和 completion_tokens
            print("\n[4/8] 新增 prompt_tokens 和 completion_tokens...")
            if 'prompt_tokens' not in current_columns:
                db.session.execute(text('''
                    ALTER TABLE dt_api_usage_stats
                    ADD COLUMN prompt_tokens INT DEFAULT 0 COMMENT 'Prompt token數' AFTER api_endpoint
                '''))
            if 'completion_tokens' not in current_columns:
                db.session.execute(text('''
                    ALTER TABLE dt_api_usage_stats
                    ADD COLUMN completion_tokens INT DEFAULT 0 COMMENT 'Completion token數' AFTER prompt_tokens
                '''))
            db.session.commit()
            print("   ✓ 已新增 token 細分欄位")

            # 5. 重新命名 token_used → total_tokens
            if 'token_used' in current_columns:
                print("\n[5/8] 重新命名 token_used → total_tokens...")
                db.session.execute(text('''
                    ALTER TABLE dt_api_usage_stats
                    CHANGE COLUMN token_used total_tokens INT DEFAULT 0 COMMENT '總token數'
                '''))
                db.session.commit()
                print("   ✓ 已重新命名 token_used → total_tokens")
            else:
                print("\n[5/8] 跳過（token_used 已不存在或已是 total_tokens）")

            # 6. 新增計費相關欄位
            print("\n[6/8] 新增計費相關欄位...")
            if 'prompt_unit_price' not in current_columns:
                db.session.execute(text('''
                    ALTER TABLE dt_api_usage_stats
                    ADD COLUMN prompt_unit_price DECIMAL(10, 8) DEFAULT 0.00000000 COMMENT '單價' AFTER total_tokens
                '''))
            if 'prompt_price_unit' not in current_columns:
                db.session.execute(text('''
                    ALTER TABLE dt_api_usage_stats
                    ADD COLUMN prompt_price_unit VARCHAR(20) DEFAULT 'USD' COMMENT '價格單位' AFTER prompt_unit_price
                '''))
            db.session.commit()
            print("   ✓ 已新增計費欄位")

            # 7. 替換 status 欄位為 success (BOOLEAN)
            print("\n[7/8] 更新 status 欄位...")
            if 'status' in current_columns and 'success' not in current_columns:
                # 先新增 success 欄位
                db.session.execute(text('''
                    ALTER TABLE dt_api_usage_stats
                    ADD COLUMN success BOOLEAN DEFAULT TRUE COMMENT '是否成功' AFTER response_time_ms
                '''))

                # 將 status 資料轉換到 success
                db.session.execute(text('''
                    UPDATE dt_api_usage_stats
                    SET success = (status = 'SUCCESS')
                '''))

                # 刪除舊的 status 欄位
                db.session.execute(text('''
                    ALTER TABLE dt_api_usage_stats
                    DROP COLUMN status
                '''))
                db.session.commit()
                print("   ✓ 已將 status 轉換為 success (BOOLEAN)")
            else:
                print("   ℹ 跳過（已完成或不需要轉換）")

            # 8. 更新索引
            print("\n[8/8] 建立索引...")
            try:
                db.session.execute(text('''
                    ALTER TABLE dt_api_usage_stats
                    ADD INDEX IF NOT EXISTS idx_api_endpoint (api_endpoint)
                '''))
            except Exception as e:
                if 'Duplicate' not in str(e):
                    raise

            try:
                db.session.execute(text('''
                    ALTER TABLE dt_api_usage_stats
                    ADD INDEX IF NOT EXISTS idx_success (success)
                '''))
            except Exception as e:
                if 'Duplicate' not in str(e):
                    raise

            db.session.commit()
            print("   ✓ 已建立索引")

            # 9. 驗證最終結構
            print("\n" + "=" * 60)
            print("驗證最終資料表結構")
            print("=" * 60)

            result = db.session.execute(text('SHOW COLUMNS FROM dt_api_usage_stats'))
            final_columns = list(result)

            print(f"\n最終欄位列表 (共 {len(final_columns)} 個):")
            for col in final_columns:
                print(f"  - {col[0]:25} {col[1]:20} NULL={col[2]} Default={col[4]}")

            # 10. 統計資料
            print("\n" + "=" * 60)
            print("資料統計")
            print("=" * 60)

            total_records = db.session.execute(
                text('SELECT COUNT(*) FROM dt_api_usage_stats')
            ).scalar()
            print(f"總記錄數: {total_records}")

            if total_records > 0:
                stats = db.session.execute(text('''
                    SELECT
                        api_endpoint,
                        COUNT(*) as count,
                        SUM(total_tokens) as total_tokens,
                        SUM(cost) as total_cost
                    FROM dt_api_usage_stats
                    GROUP BY api_endpoint
                ''')).fetchall()

                print("\nAPI 使用統計:")
                for stat in stats:
                    print(f"  {stat[0]:40} | {stat[1]:5} 次 | {stat[2]:10} tokens | ${stat[3]:.4f}")

            print("\n" + "=" * 60)
            print("✅ Migration 執行完成!")
            print("=" * 60)

        except Exception as e:
            db.session.rollback()
            print(f"\n❌ Migration 失敗: {str(e)}")
            print("\n可以使用備份表還原資料:")
            print("   DROP TABLE dt_api_usage_stats;")
            print("   CREATE TABLE dt_api_usage_stats AS SELECT * FROM dt_api_usage_stats_backup_20251001;")
            raise


if __name__ == '__main__':
    run_migration()