#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
多租户迁移脚本
为现有数据表添加 tenant_id 字段，已有数据填充为 'default'
"""

import sys
import os
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from src.core.database import db_manager
from sqlalchemy import text, inspect

TABLES_TO_MIGRATE = [
    "work_orders",
    "chat_sessions",
    "conversations",
    "knowledge_entries",
    "analytics",
    "alerts",
    "users",
]


def migrate():
    print("=" * 50)
    print("多租户迁移: 添加 tenant_id 字段")
    print("=" * 50)

    with db_manager.get_session() as session:
        inspector = inspect(session.bind)

        for table in TABLES_TO_MIGRATE:
            # 检查表是否存在
            if table not in inspector.get_table_names():
                print(f"  [跳过] 表 {table} 不存在")
                continue

            # 检查字段是否已存在
            columns = [col["name"] for col in inspector.get_columns(table)]
            if "tenant_id" in columns:
                print(f"  [已有] 表 {table} 已包含 tenant_id")
                continue

            # 添加字段
            print(f"  [迁移] 表 {table} 添加 tenant_id ...")
            session.execute(text(
                f"ALTER TABLE {table} ADD COLUMN tenant_id VARCHAR(50) DEFAULT 'default' NOT NULL"
            ))
            session.commit()
            print(f"  [完成] 表 {table} 迁移成功")

    print("\n迁移完成!")


if __name__ == "__main__":
    migrate()