"""PostgreSQLEvolutionStore - 基于 PostgreSQL 的异步进化存储

使用 async SQLAlchemy + asyncpg 实现完整的 EvolutionStoreProtocol，
支持进化事件、技能版本、A/B 测试结果的持久化。
"""

import asyncio
import logging
import uuid as _uuid
from datetime import datetime, timezone
from typing import Any

from sqlalchemy import select
from sqlalchemy.dialects.postgresql import JSONB
from sqlalchemy import Column, DateTime, Float, Integer, String, Text, UniqueConstraint
from sqlalchemy.exc import DBAPIError
from sqlalchemy.orm import declarative_base

from agentkit.core.protocol import EvolutionEvent

logger = logging.getLogger(__name__)

# PG 专用 Base（与 SQLite models 的 Base 隔离，避免表名冲突）
PGBase = declarative_base()


def _utcnow() -> datetime:
    return datetime.now(timezone.utc)


class PGEvolutionEventModel(PGBase):
    """PostgreSQL 进化事件 ORM 模型"""

    __tablename__ = "evolution_events"

    id = Column(String, primary_key=True, default=lambda: str(_uuid.uuid4()))
    agent_name = Column(String, index=True)
    change_type = Column(String, nullable=True)
    before = Column(JSONB, nullable=True)
    after = Column(JSONB, nullable=True)
    metrics = Column(JSONB, nullable=True)
    status = Column(String, default="active")
    created_at = Column(DateTime, default=_utcnow)


class PGSkillVersionModel(PGBase):
    """PostgreSQL 技能版本 ORM 模型"""

    __tablename__ = "skill_versions"
    __table_args__ = (UniqueConstraint("skill_name", "version"),)

    id = Column(String, primary_key=True, default=lambda: str(_uuid.uuid4()))
    skill_name = Column(String, index=True)
    version = Column(String)
    content = Column(Text)
    parent_version = Column(String, nullable=True)
    created_at = Column(DateTime, default=_utcnow)


class PGABTestResultModel(PGBase):
    """PostgreSQL A/B 测试结果 ORM 模型"""

    __tablename__ = "ab_test_results"

    id = Column(String, primary_key=True, default=lambda: str(_uuid.uuid4()))
    test_id = Column(String, index=True)
    variant = Column(String)
    score = Column(Float)
    sample_count = Column(Integer, default=0)
    created_at = Column(DateTime, default=_utcnow)


class PostgreSQLEvolutionStore:
    """PostgreSQL 异步进化存储

    使用 async SQLAlchemy session 实现完整的 EvolutionStoreProtocol。
    支持进化事件、技能版本、A/B 测试结果的 CRUD 操作。

    用法：
        store = PostgreSQLEvolutionStore(
            database_url="postgresql+asyncpg://user:pass@localhost/dbname"
        )
        await store.ensure_tables()
        event_id = await store.record(event)
    """

    def __init__(self, database_url: str) -> None:
        self._database_url = database_url
        self._engine: Any = None
        self._session_factory: Any = None
        self._initialized = False
        self._init_lock = asyncio.Lock()

    async def _ensure_initialized(self) -> None:
        """延迟初始化 async engine 和 session factory（带锁防并发）"""
        if self._initialized:
            return

        async with self._init_lock:
            # Double-check after acquiring lock
            if self._initialized:
                return

            from sqlalchemy.ext.asyncio import AsyncSession, create_async_engine
            from sqlalchemy.orm import sessionmaker

            self._engine = create_async_engine(self._database_url, echo=False)
            self._session_factory = sessionmaker(
                self._engine, class_=AsyncSession, expire_on_commit=False
            )
            self._initialized = True

    async def ensure_tables(self) -> None:
        """创建所有表（如果不存在）

        安全的启动调用 — 使用 CREATE TABLE IF NOT EXISTS。
        """
        await self._ensure_initialized()
        async with self._engine.begin() as conn:
            await conn.run_sync(PGBase.metadata.create_all)

    async def close(self) -> None:
        """关闭 engine，释放所有连接"""
        if self._engine is not None:
            await self._engine.dispose()
            self._engine = None
            self._session_factory = None
            self._initialized = False

    async def __aenter__(self) -> "PostgreSQLEvolutionStore":
        return self

    async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
        await self.close()

    # ── 进化事件 ──────────────────────────────────────────

    async def record(self, event: EvolutionEvent) -> str:
        """记录进化事件"""
        await self._ensure_initialized()
        async with self._session_factory() as db:
            try:
                event_id = str(_uuid.uuid4())
                entry = PGEvolutionEventModel(
                    id=event_id,
                    agent_name=event.agent_name,
                    change_type=event.change_type,
                    before=event.before,
                    after=event.after,
                    metrics=event.metrics,
                    status="active",
                )
                db.add(entry)
                await db.commit()
                event.event_id = event_id
                logger.info(f"Evolution event recorded: {event_id} for agent '{event.agent_name}'")
                return event_id
            except (DBAPIError, ValueError, KeyError, RuntimeError) as e:
                await db.rollback()
                logger.error(f"Failed to record evolution event: {e}")
                raise

    async def rollback(self, event_id: str) -> bool:
        """回滚进化事件"""
        await self._ensure_initialized()
        async with self._session_factory() as db:
            try:
                stmt = select(PGEvolutionEventModel).where(
                    PGEvolutionEventModel.id == event_id
                )
                result = await db.execute(stmt)
                entry = result.scalar_one_or_none()

                if not entry:
                    logger.error(f"Evolution event {event_id} not found")
                    return False

                entry.status = "rolled_back"
                await db.commit()
                logger.info(f"Evolution event {event_id} rolled back")
                return True
            except (DBAPIError, ValueError, KeyError, RuntimeError) as e:
                await db.rollback()
                logger.error(f"Failed to rollback evolution event {event_id}: {e}")
                return False

    async def list_events(
        self,
        agent_name: str | None = None,
        change_type: str | None = None,
        status: str | None = None,
        limit: int = 100,
    ) -> list[dict]:
        """列出进化事件"""
        await self._ensure_initialized()
        async with self._session_factory() as db:
            try:
                stmt = select(PGEvolutionEventModel)
                if agent_name:
                    stmt = stmt.where(PGEvolutionEventModel.agent_name == agent_name)
                if change_type:
                    stmt = stmt.where(PGEvolutionEventModel.change_type == change_type)
                if status:
                    stmt = stmt.where(PGEvolutionEventModel.status == status)
                stmt = stmt.order_by(PGEvolutionEventModel.created_at.desc()).limit(limit)
                result = await db.execute(stmt)
                entries = result.scalars().all()

                return [
                    {
                        "id": e.id,
                        "agent_name": e.agent_name,
                        "change_type": e.change_type,
                        "before": e.before,
                        "after": e.after,
                        "metrics": e.metrics,
                        "status": e.status,
                        "created_at": e.created_at.isoformat() if e.created_at else None,
                    }
                    for e in entries
                ]
            except (DBAPIError, ValueError, KeyError, RuntimeError) as e:
                logger.error(f"Failed to list evolution events: {e}")
                return []

    # ── 技能版本 ──────────────────────────────────────────

    async def record_skill_version(
        self, skill_name: str, version: str, content: str, parent_version: str | None = None
    ) -> str:
        """记录技能版本"""
        await self._ensure_initialized()
        async with self._session_factory() as db:
            try:
                vid = str(_uuid.uuid4())
                entry = PGSkillVersionModel(
                    id=vid,
                    skill_name=skill_name,
                    version=version,
                    content=content,
                    parent_version=parent_version,
                )
                db.add(entry)
                await db.commit()
                return vid
            except (DBAPIError, ValueError, KeyError, RuntimeError) as e:
                await db.rollback()
                logger.error(f"Failed to record skill version: {e}")
                raise

    async def list_skill_versions(self, skill_name: str, limit: int = 100) -> list[dict]:
        """列出技能版本历史"""
        await self._ensure_initialized()
        async with self._session_factory() as db:
            try:
                stmt = (
                    select(PGSkillVersionModel)
                    .where(PGSkillVersionModel.skill_name == skill_name)
                    .order_by(PGSkillVersionModel.created_at.desc())
                    .limit(limit)
                )
                result = await db.execute(stmt)
                entries = result.scalars().all()
                return [
                    {
                        "id": e.id,
                        "skill_name": e.skill_name,
                        "version": e.version,
                        "content": e.content,
                        "parent_version": e.parent_version,
                        "created_at": e.created_at.isoformat() if e.created_at else None,
                    }
                    for e in entries
                ]
            except (DBAPIError, ValueError, KeyError, RuntimeError) as e:
                logger.error(f"Failed to list skill versions: {e}")
                return []

    # ── A/B 测试结果 ──────────────────────────────────────

    async def record_ab_test_result(
        self, test_id: str, variant: str, score: float, sample_count: int = 0
    ) -> str:
        """记录 A/B 测试结果"""
        await self._ensure_initialized()
        async with self._session_factory() as db:
            try:
                rid = str(_uuid.uuid4())
                entry = PGABTestResultModel(
                    id=rid,
                    test_id=test_id,
                    variant=variant,
                    score=score,
                    sample_count=sample_count,
                )
                db.add(entry)
                await db.commit()
                return rid
            except (DBAPIError, ValueError, KeyError, RuntimeError) as e:
                await db.rollback()
                logger.error(f"Failed to record A/B test result: {e}")
                raise

    async def get_ab_test_results(self, test_id: str, limit: int = 100) -> list[dict]:
        """获取 A/B 测试结果"""
        await self._ensure_initialized()
        async with self._session_factory() as db:
            try:
                stmt = (
                    select(PGABTestResultModel)
                    .where(PGABTestResultModel.test_id == test_id)
                    .order_by(PGABTestResultModel.created_at.desc())
                    .limit(limit)
                )
                result = await db.execute(stmt)
                entries = result.scalars().all()
                return [
                    {
                        "id": e.id,
                        "test_id": e.test_id,
                        "variant": e.variant,
                        "score": e.score,
                        "sample_count": e.sample_count,
                        "created_at": e.created_at.isoformat() if e.created_at else None,
                    }
                    for e in entries
                ]
            except (DBAPIError, ValueError, KeyError, RuntimeError) as e:
                logger.error(f"Failed to get A/B test results: {e}")
                return []
