"""Add knowledge graph tables

Revision ID: f7a8b9c0de56
Revises: e5f7g9h1cd45
Create Date: 2026-05-24 12:00:00.000000

"""
from typing import Sequence, Union

import sqlalchemy as sa
from alembic import op
from sqlalchemy.dialects import postgresql
import enum

# revision identifiers, used by Alembic.
revision: str = "f7a8b9c0de56"
down_revision: Union[str, None] = "810a29804f5a"
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None


def upgrade() -> None:
    # ------------------------------------------------------------------ #
    # 1. 创建实体类型枚举
    # ------------------------------------------------------------------ #
    entity_type_enum = enum.Enum(
        'EntityType',
        [
            'ORGANIZATION', 'PRODUCT', 'PERSON', 'LOCATION',
            'TECHNOLOGY', 'BRAND', 'EVENT', 'CONCEPT', 'OTHER'
        ]
    )

    # ------------------------------------------------------------------ #
    # 2. 创建关系类型枚举
    # ------------------------------------------------------------------ #
    relation_type_enum = enum.Enum(
        'RelationType',
        [
            'COMPETES_WITH', 'PARTNERS_WITH', 'ACQUIRES', 'SUBSIDIARY_OF',
            'PRODUCES', 'USES_TECHNOLOGY', 'PART_OF',
            'LOCATED_IN', 'FOUNDED_IN',
            'CEO_OF', 'FOUNDER_OF',
            'RELATED_TO', 'MENTIONED_IN', 'ALSO_KNOWN_AS'
        ]
    )

    # ------------------------------------------------------------------ #
    # 3. knowledge_entities 表
    # ------------------------------------------------------------------ #
    op.create_table(
        "knowledge_entities",
        sa.Column(
            "id",
            postgresql.UUID(as_uuid=True),
            primary_key=True,
            nullable=False,
        ),
        sa.Column(
            "knowledge_base_id",
            postgresql.UUID(as_uuid=True),
            sa.ForeignKey("knowledge_bases.id", ondelete="CASCADE"),
            nullable=False,
        ),
        sa.Column("name", sa.String(500), nullable=False),
        sa.Column("entity_type", sa.Enum(entity_type_enum, name="entitytype"), nullable=False),
        sa.Column("description", sa.Text, nullable=True),
        sa.Column("properties", postgresql.JSONB(astext_type=sa.Text()), nullable=True, server_default="{}"),
        sa.Column(
            "source_chunk_id",
            postgresql.UUID(as_uuid=True),
            sa.ForeignKey("knowledge_chunks.id", ondelete="SET NULL"),
            nullable=True,
        ),
        sa.Column("confidence", sa.String(20), nullable=True),
        sa.Column(
            "created_at",
            sa.DateTime(timezone=True),
            server_default=sa.text("now()"),
            nullable=False,
        ),
        sa.Column(
            "updated_at",
            sa.DateTime(timezone=True),
            server_default=sa.text("now()"),
            nullable=False,
        ),
    )
    op.create_index("ix_entities_name", "knowledge_entities", ["name"])
    op.create_index("ix_entities_kb_name", "knowledge_entities", ["knowledge_base_id", "name"])
    op.create_index("ix_entities_kb_type", "knowledge_entities", ["knowledge_base_id", "entity_type"])

    # ------------------------------------------------------------------ #
    # 4. knowledge_relations 表
    # ------------------------------------------------------------------ #
    op.create_table(
        "knowledge_relations",
        sa.Column(
            "id",
            postgresql.UUID(as_uuid=True),
            primary_key=True,
            nullable=False,
        ),
        sa.Column(
            "source_entity_id",
            postgresql.UUID(as_uuid=True),
            sa.ForeignKey("knowledge_entities.id", ondelete="CASCADE"),
            nullable=False,
        ),
        sa.Column(
            "target_entity_id",
            postgresql.UUID(as_uuid=True),
            sa.ForeignKey("knowledge_entities.id", ondelete="CASCADE"),
            nullable=False,
        ),
        sa.Column("relation_type", sa.Enum(relation_type_enum, name="relationtype"), nullable=False),
        sa.Column("properties", postgresql.JSONB(astext_type=sa.Text()), nullable=True, server_default="{}"),
        sa.Column(
            "source_chunk_id",
            postgresql.UUID(as_uuid=True),
            sa.ForeignKey("knowledge_chunks.id", ondelete="SET NULL"),
            nullable=True,
        ),
        sa.Column("confidence", sa.String(20), nullable=True),
        sa.Column(
            "created_at",
            sa.DateTime(timezone=True),
            server_default=sa.text("now()"),
            nullable=False,
        ),
    )
    op.create_index("ix_relations_source", "knowledge_relations", ["source_entity_id"])
    op.create_index("ix_relations_target", "knowledge_relations", ["target_entity_id"])
    op.create_index("ix_relations_type", "knowledge_relations", ["relation_type"])


def downgrade() -> None:
    # 删除表（注意外键约束会自动处理）
    op.drop_table("knowledge_relations")
    op.drop_table("knowledge_entities")
    # 删除枚举类型
    op.execute("DROP TYPE IF EXISTS relationtype")
    op.execute("DROP TYPE IF EXISTS entitytype")