#!/usr/bin/env python3
"""
数据库迁移脚本：添加 revision_count, is_latest, original_entry_id 字段
"""
import os
import sys
from pathlib import Path

# 加载 .env 文件
env_file = Path(__file__).parent.parent / ".env"
if env_file.exists():
    with open(env_file) as f:
        for line in f:
            line = line.strip()
            if line and not line.startswith("#"):
                key, value = line.split("=", 1)
                os.environ[key.strip()] = value.strip()

from psycopg import connect

def run_migration():
    try:
        # 从环境变量获取数据库连接信息
        db_password = os.getenv("DB_PASSWORD")
        print(f"📛 连接数据库: {os.getenv('DB_HOST')}:{os.getenv('DB_PORT')}/{os.getenv('DB_NAME')}")
        print(f"👤 用户: {os.getenv('DB_USER')}")
        
        # 尝试用 njts_app 用户连接，不行的话尝试 postgres
        users_to_try = [
            {
                "user": os.getenv("DB_USER"),
                "password": db_password
            },
            {
                "user": "postgres",
                "password": ""  # 空密码或尝试无密码
            }
        ]
        
        conn = None
        last_error = None
        
        for cred in users_to_try:
            try:
                print(f"尝试连接用户: {cred['user']}")
                conn = connect(
                    host=os.getenv("DB_HOST"),
                    port=int(os.getenv("DB_PORT")),
                    dbname=os.getenv("DB_NAME"),
                    user=cred["user"],
                    password=cred["password"] if cred["password"] else None,
                )
                print(f"✓ 成功连接为用户: {cred['user']}")
                break
            except Exception as e:
                last_error = e
                print(f"✗ 连接失败: {e}")
                continue
        
        if not conn:
            raise Exception(f"无法连接到数据库。最后错误: {last_error}")
        
        with conn.cursor() as cur:
            # 首先检查字段是否存在
            print("🔄 检查字段是否存在...")
            cur.execute("""
                SELECT column_name FROM information_schema.columns 
                WHERE table_name = 'journal_entries' 
                AND column_name IN ('is_latest', 'revision_count', 'original_entry_id')
            """)
            existing_columns = {row[0] for row in cur.fetchall()}
            print(f"已存在的字段: {existing_columns}")
            
            # 添加 is_latest 字段
            if 'is_latest' not in existing_columns:
                print("🔄 添加 is_latest 字段...")
                try:
                    cur.execute("""
                        ALTER TABLE journal_entries 
                        ADD COLUMN is_latest BOOLEAN DEFAULT true;
                    """)
                    print("✓ is_latest 字段已添加")
                except Exception as e:
                    print(f"✗ 添加失败: {e}")
            else:
                print("✓ is_latest 字段已存在")
            print("✓ is_latest 字段处理完成")
            
            # 添加 revision_count 字段
            print("🔄 检查 revision_count 字段...")
            cur.execute("""
                DO $$
                BEGIN
                    IF NOT EXISTS (SELECT 1 FROM information_schema.columns 
                                   WHERE table_name = 'journal_entries' AND column_name = 'revision_count') THEN
                        ALTER TABLE journal_entries ADD COLUMN revision_count INTEGER DEFAULT 1;
                        RAISE NOTICE '字段 revision_count 已添加';
                    ELSE
                        RAISE NOTICE '字段 revision_count 已存在';
                    END IF;
                END $$;
            """)
            print("✓ revision_count 字段处理完成")
            
            # 添加 original_entry_id 字段
            print("🔄 检查 original_entry_id 字段...")
            cur.execute("""
                DO $$
                BEGIN
                    IF NOT EXISTS (SELECT 1 FROM information_schema.columns 
                                   WHERE table_name = 'journal_entries' AND column_name = 'original_entry_id') THEN
                        ALTER TABLE journal_entries ADD COLUMN original_entry_id INTEGER;
                        ALTER TABLE journal_entries ADD CONSTRAINT fk_original_entry_id 
                            FOREIGN KEY (original_entry_id) REFERENCES journal_entries(journal_entry_id) ON DELETE CASCADE;
                        RAISE NOTICE '字段 original_entry_id 已添加';
                    ELSE
                        RAISE NOTICE '字段 original_entry_id 已存在';
                    END IF;
                END $$;
            """)
            print("✓ original_entry_id 字段处理完成")
            
            # 创建索引
            print("🔄 创建索引...")
            cur.execute("""
                CREATE INDEX IF NOT EXISTS idx_journal_entries_is_latest
                ON journal_entries(is_latest, entry_date);
            """)
            print("✓ idx_journal_entries_is_latest 索引已创建")
            
            cur.execute("""
                CREATE INDEX IF NOT EXISTS idx_journal_entries_original_entry_id
                ON journal_entries(original_entry_id);
            """)
            print("✓ idx_journal_entries_original_entry_id 索引已创建")
            
            # 验证字段
            print("🔄 验证字段...")
            cur.execute("""
                SELECT column_name, data_type, is_nullable
                FROM information_schema.columns
                WHERE table_name = 'journal_entries'
                  AND column_name IN ('is_latest', 'revision_count', 'original_entry_id')
                ORDER BY ordinal_position;
            """)
            
            rows = cur.fetchall()
            print("\n📋 数据库字段信息:")
            for row in rows:
                print(f"  - {row[0]}: {row[1]} (nullable: {row[2]})")
            
        conn.commit()
        conn.close()
        
        print("\n✅ 迁移完成！")
        return 0
        
    except Exception as e:
        print(f"\n❌ 迁移失败: {e}")
        import traceback
        traceback.print_exc()
        return 1

if __name__ == "__main__":
    sys.exit(run_migration())