from sqlalchemy.ext.asyncio import create_async_engine, async_sessionmaker, AsyncSession
from sqlalchemy.orm import DeclarativeBase
from sqlalchemy import text
from config import settings


class Base(DeclarativeBase):
    """SQLAlchemy ORM 基类，所有数据库模型均继承此类。"""
    pass


# 异步数据库引擎，连接池大小 20，最大溢出 40，启用连接健康检查
async_engine = create_async_engine(
    settings.DATABASE_URL,
    pool_size=20,
    max_overflow=40,
    pool_pre_ping=True,
    pool_recycle=3600,
    echo=False,
)

# 异步数据库会话工厂，用于创建数据库会话实例
AsyncSessionLocal = async_sessionmaker(
    async_engine,
    class_=AsyncSession,
    expire_on_commit=False,
)


async def init_db():
    """初始化数据库：创建所有表并执行增量迁移。"""
    async with async_engine.begin() as conn:
        from models import Base as MBase
        await conn.run_sync(MBase.metadata.create_all)

    await _run_migrations()


async def _run_migrations():
    """执行数据库增量迁移：创建缺失的表并安全添加新字段。"""
    async with async_engine.begin() as conn:
        # --- 创建可能缺失的新表 ---
        await conn.execute(text("""
            CREATE TABLE IF NOT EXISTS agent_configs (
                id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
                name VARCHAR(200) NOT NULL,
                description TEXT,
                system_prompt TEXT,
                model VARCHAR(100) DEFAULT '',
                model_instance_id UUID,
                embedding_model_id UUID,
                temperature FLOAT DEFAULT 0.7,
                tools JSONB DEFAULT '[]',
                status VARCHAR(20) DEFAULT 'active',
                creator_id UUID REFERENCES users(id),
                created_at TIMESTAMP DEFAULT now(),
                updated_at TIMESTAMP DEFAULT now()
            )
        """))
        await conn.execute(text("""
            CREATE TABLE IF NOT EXISTS model_providers (
                id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
                name VARCHAR(100) NOT NULL,
                provider_type VARCHAR(50) NOT NULL,
                base_url VARCHAR(500),
                api_key VARCHAR(500),
                extra_config JSONB DEFAULT '{}',
                is_active BOOLEAN DEFAULT TRUE,
                created_at TIMESTAMP DEFAULT now()
            )
        """))
        await conn.execute(text("""
            CREATE TABLE IF NOT EXISTS model_instances (
                id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
                provider_id UUID NOT NULL REFERENCES model_providers(id) ON DELETE CASCADE,
                model_name VARCHAR(200) NOT NULL,
                model_type VARCHAR(20) NOT NULL,
                display_name VARCHAR(200),
                capabilities JSONB DEFAULT '{}',
                default_params JSONB DEFAULT '{}',
                is_default BOOLEAN DEFAULT FALSE,
                is_active BOOLEAN DEFAULT TRUE,
                created_at TIMESTAMP DEFAULT now()
            )
        """))
        await conn.execute(text("""
            CREATE TABLE IF NOT EXISTS custom_tools (
                id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
                name VARCHAR(100) UNIQUE NOT NULL,
                description TEXT,
                schema_json JSONB NOT NULL DEFAULT '{}',
                endpoint_url VARCHAR(500) DEFAULT '',
                method VARCHAR(10) DEFAULT 'GET',
                path VARCHAR(500) DEFAULT '',
                headers_json JSONB DEFAULT '{}',
                auth_type VARCHAR(20) DEFAULT 'none',
                auth_config JSONB DEFAULT '{}',
                is_active BOOLEAN DEFAULT TRUE,
                created_by UUID REFERENCES users(id),
                created_at TIMESTAMP DEFAULT now(),
                updated_at TIMESTAMP DEFAULT now()
            )
        """))
        await conn.execute(text("""
            CREATE TABLE IF NOT EXISTS mcp_services (
                id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
                name VARCHAR(100) NOT NULL,
                endpoint_url VARCHAR(500) NOT NULL,
                status VARCHAR(20) DEFAULT 'active',
                config JSONB DEFAULT '{}',
                last_heartbeat TIMESTAMP,
                created_at TIMESTAMP DEFAULT now(),
                updated_at TIMESTAMP DEFAULT now()
            )
        """))
        await conn.execute(text("""
            CREATE TABLE IF NOT EXISTS audit_logs (
                id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
                user_id UUID REFERENCES users(id),
                action VARCHAR(50) NOT NULL,
                resource_type VARCHAR(50) NOT NULL,
                resource_id VARCHAR(100),
                detail JSONB DEFAULT '{}',
                ip_address VARCHAR(50),
                user_agent VARCHAR(500),
                created_at TIMESTAMP DEFAULT now()
            )
        """))
        await conn.execute(text("""
            CREATE TABLE IF NOT EXISTS notification_templates (
                id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
                name VARCHAR(100) UNIQUE NOT NULL,
                channel VARCHAR(20) NOT NULL,
                event_type VARCHAR(50) NOT NULL,
                subject_template VARCHAR(500),
                body_template TEXT,
                is_active BOOLEAN DEFAULT TRUE,
                created_at TIMESTAMP DEFAULT now(),
                updated_at TIMESTAMP DEFAULT now()
            )
        """))
        await conn.execute(text("""
            CREATE TABLE IF NOT EXISTS system_metrics (
                id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
                metric_name VARCHAR(50) NOT NULL,
                metric_value FLOAT,
                tags JSONB DEFAULT '{}',
                recorded_at TIMESTAMP DEFAULT now()
            )
        """))
        await conn.execute(text("""
            CREATE TABLE IF NOT EXISTS flow_executions (
                id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
                flow_id UUID NOT NULL REFERENCES flow_definitions(id),
                session_id VARCHAR(100),
                status VARCHAR(20) DEFAULT 'running',
                input_data JSONB DEFAULT '{}',
                output_data JSONB DEFAULT '{}',
                error_message TEXT,
                started_by UUID REFERENCES users(id),
                started_at TIMESTAMP DEFAULT now(),
                finished_at TIMESTAMP,
                duration_ms INTEGER
            )
        """))
        await conn.execute(text("""
            CREATE TABLE IF NOT EXISTS memory_messages (
                id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
                flow_execution_id UUID NOT NULL REFERENCES flow_executions(id) ON DELETE CASCADE,
                node_id VARCHAR(100) NOT NULL,
                role VARCHAR(20) NOT NULL,
                content TEXT NOT NULL,
                metadata JSONB DEFAULT '{}',
                created_at TIMESTAMP DEFAULT now()
            )
        """))

        # --- 为已有表添加缺失列 ---
        await conn.execute(text(
            "ALTER TABLE flow_definitions ADD COLUMN IF NOT EXISTS published_version_id UUID REFERENCES flow_versions(id)"
        ))
        await conn.execute(text(
            "ALTER TABLE flow_definitions ADD COLUMN IF NOT EXISTS draft_definition_json JSONB"
        ))
        await conn.execute(text(
            "ALTER TABLE flow_definitions ADD COLUMN IF NOT EXISTS flow_mode VARCHAR(20) DEFAULT 'chatflow'"
        ))
        await conn.execute(text(
            "ALTER TABLE flow_definitions ADD COLUMN IF NOT EXISTS published_to_wecom BOOLEAN DEFAULT FALSE"
        ))
        await conn.execute(text(
            "ALTER TABLE flow_definitions ADD COLUMN IF NOT EXISTS published_to_web BOOLEAN DEFAULT FALSE"
        ))
        await conn.execute(text(
            "ALTER TABLE agent_configs ADD COLUMN IF NOT EXISTS model_instance_id UUID"
        ))
        await conn.execute(text(
            "ALTER TABLE agent_configs ADD COLUMN IF NOT EXISTS embedding_model_id UUID"
        ))
        # mcp_services 表可能存在旧版本，补齐所有缺失列
        await conn.execute(text(
            "ALTER TABLE mcp_services ADD COLUMN IF NOT EXISTS name VARCHAR(100) NOT NULL DEFAULT ''"
        ))
        await conn.execute(text(
            "ALTER TABLE mcp_services ADD COLUMN IF NOT EXISTS endpoint_url VARCHAR(500) NOT NULL DEFAULT ''"
        ))
        await conn.execute(text(
            "ALTER TABLE mcp_services ADD COLUMN IF NOT EXISTS status VARCHAR(20) DEFAULT 'active'"
        ))
        await conn.execute(text(
            "ALTER TABLE mcp_services ADD COLUMN IF NOT EXISTS config JSONB DEFAULT '{}'"
        ))
        await conn.execute(text(
            "ALTER TABLE mcp_services ADD COLUMN IF NOT EXISTS last_heartbeat TIMESTAMP"
        ))
        await conn.execute(text(
            "ALTER TABLE mcp_services ADD COLUMN IF NOT EXISTS created_at TIMESTAMP DEFAULT now()"
        ))
        await conn.execute(text(
            "ALTER TABLE mcp_services ADD COLUMN IF NOT EXISTS updated_at TIMESTAMP DEFAULT now()"
        ))
        # --- 其他新表的列级迁移（防止旧版本表缺少字段） ---
        # model_providers
        for col_sql in [
            "name VARCHAR(100) NOT NULL DEFAULT ''",
            "provider_type VARCHAR(50) NOT NULL DEFAULT 'openai_compatible'",
            "base_url VARCHAR(500) DEFAULT ''",
            "api_key VARCHAR(500) DEFAULT ''",
            "extra_config JSONB DEFAULT '{}'",
            "is_active BOOLEAN DEFAULT TRUE",
            "created_at TIMESTAMP DEFAULT now()",
        ]:
            await conn.execute(text(f"ALTER TABLE model_providers ADD COLUMN IF NOT EXISTS {col_sql}"))
        # model_instances
        for col_sql in [
            "id UUID PRIMARY KEY DEFAULT gen_random_uuid()",
            "provider_id UUID NOT NULL REFERENCES model_providers(id) ON DELETE CASCADE",
            "model_name VARCHAR(200) NOT NULL DEFAULT ''",
            "model_type VARCHAR(20) NOT NULL DEFAULT 'llm'",
            "display_name VARCHAR(200) DEFAULT ''",
            "capabilities JSONB DEFAULT '{}'",
            "default_params JSONB DEFAULT '{}'",
            "is_default BOOLEAN DEFAULT FALSE",
            "is_active BOOLEAN DEFAULT TRUE",
            "created_at TIMESTAMP DEFAULT now()",
        ]:
            await conn.execute(text(f"ALTER TABLE model_instances ADD COLUMN IF NOT EXISTS {col_sql}"))
        # custom_tools
        for col_sql in [
            "name VARCHAR(100) UNIQUE NOT NULL DEFAULT ''",
            "description TEXT",
            "schema_json JSONB NOT NULL DEFAULT '{}'",
            "endpoint_url VARCHAR(500) DEFAULT ''",
            "method VARCHAR(10) DEFAULT 'GET'",
            "path VARCHAR(500) DEFAULT ''",
            "headers_json JSONB DEFAULT '{}'",
            "auth_type VARCHAR(20) DEFAULT 'none'",
            "auth_config JSONB DEFAULT '{}'",
            "is_active BOOLEAN DEFAULT TRUE",
            "created_by UUID REFERENCES users(id)",
            "created_at TIMESTAMP DEFAULT now()",
            "updated_at TIMESTAMP DEFAULT now()",
        ]:
            await conn.execute(text(f"ALTER TABLE custom_tools ADD COLUMN IF NOT EXISTS {col_sql}"))
        # audit_logs
        for col_sql in [
            "user_id UUID REFERENCES users(id)",
            "action VARCHAR(50) NOT NULL DEFAULT ''",
            "resource_type VARCHAR(50) NOT NULL DEFAULT ''",
            "resource_id VARCHAR(100)",
            "detail JSONB DEFAULT '{}'",
            "ip_address VARCHAR(50)",
            "user_agent VARCHAR(500)",
            "created_at TIMESTAMP DEFAULT now()",
        ]:
            await conn.execute(text(f"ALTER TABLE audit_logs ADD COLUMN IF NOT EXISTS {col_sql}"))
        # notification_templates
        for col_sql in [
            "name VARCHAR(100) UNIQUE NOT NULL DEFAULT ''",
            "channel VARCHAR(20) NOT NULL DEFAULT ''",
            "event_type VARCHAR(50) NOT NULL DEFAULT ''",
            "subject_template VARCHAR(500)",
            "body_template TEXT",
            "is_active BOOLEAN DEFAULT TRUE",
            "created_at TIMESTAMP DEFAULT now()",
            "updated_at TIMESTAMP DEFAULT now()",
        ]:
            await conn.execute(text(f"ALTER TABLE notification_templates ADD COLUMN IF NOT EXISTS {col_sql}"))
        # system_metrics
        for col_sql in [
            "metric_name VARCHAR(50) NOT NULL DEFAULT ''",
            "metric_value FLOAT",
            "tags JSONB DEFAULT '{}'",
            "recorded_at TIMESTAMP DEFAULT now()",
        ]:
            await conn.execute(text(f"ALTER TABLE system_metrics ADD COLUMN IF NOT EXISTS {col_sql}"))
        # flow_executions
        for col_sql in [
            "flow_id UUID NOT NULL REFERENCES flow_definitions(id)",
            "session_id VARCHAR(100)",
            "status VARCHAR(20) DEFAULT 'running'",
            "input_data JSONB DEFAULT '{}'",
            "output_data JSONB DEFAULT '{}'",
            "error_message TEXT",
            "started_by UUID REFERENCES users(id)",
            "started_at TIMESTAMP DEFAULT now()",
            "finished_at TIMESTAMP",
            "duration_ms INTEGER",
        ]:
            await conn.execute(text(f"ALTER TABLE flow_executions ADD COLUMN IF NOT EXISTS {col_sql}"))
        # memory_messages
        for col_sql in [
            "flow_execution_id UUID NOT NULL REFERENCES flow_executions(id) ON DELETE CASCADE",
            "node_id VARCHAR(100) NOT NULL DEFAULT ''",
            "role VARCHAR(20) NOT NULL DEFAULT 'user'",
            "content TEXT NOT NULL DEFAULT ''",
            "metadata JSONB DEFAULT '{}'",
            "created_at TIMESTAMP DEFAULT now()",
        ]:
            await conn.execute(text(f"ALTER TABLE memory_messages ADD COLUMN IF NOT EXISTS {col_sql}"))
        await conn.execute(text(
            "ALTER TABLE roles ADD COLUMN IF NOT EXISTS role_type VARCHAR(20) DEFAULT 'position'"
        ))


async def get_db():
    """FastAPI 依赖注入函数，提供数据库会话，自动提交或回滚事务。"""
    async with AsyncSessionLocal() as session:
        try:
            yield session
            await session.commit()
        except Exception:
            await session.rollback()
            raise
        finally:
            await session.close()