#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
数据库初始化脚本
用于创建数据库表结构
"""

import os
import sys
from pathlib import Path

# 添加项目根目录到系统路径
project_root = Path(__file__).parent.parent.absolute()
sys.path.insert(0, str(project_root))

from sqlalchemy import create_engine, text
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker

# 导入模型
from data.models.base import BaseModel
from data.models.vwedtask import (
    VWEDTaskDef,
    VWEDTaskRecord,
    VWEDTaskLog,
    VWEDBlockRecord,
    VWEDTaskTemplate,
    VWEDDataCache,
    VWEDScript,
    VWEDScriptVersion,
    VWEDScriptLog,
    InterfaceDefHistory
)

# 数据库连接信息
DB_USER = os.environ.get("DB_USER", "root")
DB_PASSWORD = os.environ.get("DB_PASSWORD", "password")
DB_HOST = os.environ.get("DB_HOST", "localhost")
DB_PORT = os.environ.get("DB_PORT", "3306")
DB_NAME = os.environ.get("DB_NAME", "vwed_task")

# 创建数据库连接URL
DATABASE_URL = f"mysql+pymysql://{DB_USER}:{DB_PASSWORD}@{DB_HOST}:{DB_PORT}/{DB_NAME}"

# 创建数据库引擎
engine = create_engine(DATABASE_URL, echo=True)

# 创建session工厂
SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)


def create_database():
    """创建数据库（如果不存在）"""
    # 连接到MySQL服务器
    root_engine = create_engine(f"mysql+pymysql://{DB_USER}:{DB_PASSWORD}@{DB_HOST}:{DB_PORT}/mysql")
    
    with root_engine.connect() as conn:
        # 创建数据库（如果不存在）
        conn.execute(text(f"CREATE DATABASE IF NOT EXISTS {DB_NAME} CHARACTER SET utf8mb4 COLLATE utf8mb4_unicode_ci"))
        print(f"数据库 {DB_NAME} 已创建或已存在")


def create_tables():
    """创建所有表"""
    # 创建所有表
    BaseModel.metadata.create_all(bind=engine)
    print("所有表已创建")


def main():
    """主函数"""
    try:
        # 创建数据库（如果不存在）
        create_database()
        
        # 创建表
        create_tables()
        
        print("数据库初始化完成")
    except Exception as e:
        print(f"初始化数据库时出错: {e}")
        sys.exit(1)


if __name__ == "__main__":
    main()