"""数据中台 schema 初始化。

MVP 阶段以 SQLAlchemy metadata 建表（后续可迁移到 Alembic）。
扩展按可用性可选启用：
- btree_gist / vector：若可用则创建。
- timescaledb：若可用则把 metric_event 转为超表；不可用则保持普通表（带时间索引）。
"""

from __future__ import annotations

from sqlalchemy import text
from sqlalchemy.engine import Engine

from app.datahub import models  # noqa: F401  确保模型注册到 metadata
from app.db import Base, get_engine


def _extension_available(engine: Engine, name: str) -> bool:
    with engine.connect() as conn:
        row = conn.execute(
            text("SELECT 1 FROM pg_available_extensions WHERE name = :n"), {"n": name}
        ).first()
        return row is not None


def init_extensions(engine: Engine) -> dict[str, bool]:
    """按可用性创建扩展，返回各扩展启用状态。"""
    status: dict[str, bool] = {}
    for ext in ("btree_gist", "vector", "timescaledb"):
        available = _extension_available(engine, ext)
        status[ext] = available
        if available:
            with engine.begin() as conn:
                conn.execute(text(f"CREATE EXTENSION IF NOT EXISTS {ext}"))
    return status


def create_schema(engine: Engine | None = None) -> dict[str, bool]:
    """创建数据中台全部表，并按需启用时序超表。返回扩展状态。"""
    engine = engine or get_engine()
    status = init_extensions(engine)
    Base.metadata.create_all(engine)

    # 若 TimescaleDB 可用，将时序事件表转为超表（幂等）
    if status.get("timescaledb"):
        with engine.begin() as conn:
            conn.execute(
                text(
                    "SELECT create_hypertable('metric_event', 'event_time', "
                    "if_not_exists => TRUE, migrate_data => TRUE)"
                )
            )
    return status


if __name__ == "__main__":
    st = create_schema()
    print("数据中台 schema 初始化完成。扩展状态：", st)