#!/usr/bin/env python3
"""
修复脚本：确保 gsdh_data 表有 industry_company 列
"""

import psycopg2
import json
import os

def load_config():
    """加载数据库配置"""
    if os.path.exists("config.json"):
        with open("config.json", "r", encoding="utf-8") as f:
            return json.load(f)
    return {}

def safe_add_column(conn, table, col, dtype):
    """安全地添加列（如果不存在）"""
    cur = conn.cursor()
    try:
        # 首先检查列是否存在
        cur.execute("""
            SELECT column_name FROM information_schema.columns
            WHERE table_name = %s AND column_name = %s
        """, (table, col))
        
        if not cur.fetchone():
            cur.execute(f"ALTER TABLE {table} ADD COLUMN {col} {dtype}")
            print(f"✅ 添加列 {col} 到 {table} 表")
            return True
        else:
            print(f"ℹ️  列 {col} 已存在于 {table} 表")
            return False
    except Exception as e:
        print(f"❌ 添加列 {col} 时出错: {e}")
        return False
    finally:
        cur.close()

def main():
    config = load_config()
    
    if not config.get('db_host'):
        print("❌ 未找到数据库配置")
        return
    
    print(f"🔗 连接到数据库: {config.get('db_host')}:{config.get('db_port')}/{config.get('db_name')}")
    
    try:
        conn = psycopg2.connect(
            host=config.get("db_host", "localhost"),
            port=config.get("db_port", "5432"),
            user=config.get("db_user", "gsdh"),
            password=config.get("db_password", "123gsdh"),
            database=config.get("db_name", "gsdh"),
            connect_timeout=5
        )
        
        # 确保 gsdh_data 表有 industry_company 列
        print("\n📋 检查 gsdh_data 表结构...")
        added = safe_add_column(conn, 'gsdh_data', 'industry_company', 'VARCHAR(200)')
        
        if added:
            conn.commit()
            print("\n✅ 数据库修复成功！")
            print("   现在签到功能应该可以正常工作了。")
        else:
            print("\n✅ 数据库已经是最新状态，无需修复")
        
        # 显示当前表结构
        print("\n📊 当前 gsdh_data 表的列：")
        cur = conn.cursor()
        cur.execute("""
            SELECT column_name, data_type
            FROM information_schema.columns
            WHERE table_name = 'gsdh_data'
            ORDER BY ordinal_position
        """)
        for col in cur.fetchall():
            print(f"   - {col[0]}: {col[1]}")
        cur.close()
        
        conn.close()
        
    except Exception as e:
        print(f"\n❌ 连接数据库失败: {e}")
        import traceback
        traceback.print_exc()

if __name__ == "__main__":
    main()