"""EvolutionStore - 进化日志存储

提供三种后端实现：
- EvolutionStore: 基于外部注入的异步 SQLAlchemy session（原有实现）
- PersistentEvolutionStore: 基于 SQLite 的持久化存储
- InMemoryEvolutionStore: 基于内存字典的轻量存储（用于测试）
"""

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

from sqlalchemy import create_engine, event as sa_event, select
from sqlalchemy.exc import OperationalError
from sqlalchemy.orm import sessionmaker

from agentkit.core.protocol import EvolutionEvent
from agentkit.evolution.models import (
    ABTestResultModel,
    Base,
    EvolutionEventModel,
    SkillVersionModel,
)

logger = logging.getLogger(__name__)


class EvolutionStore:
    """进化日志存储

    记录 Agent 的自我进化变更，支持回滚。
    """

    def __init__(self, session_factory: Any, evolution_model: Any):
        self._session_factory = session_factory
        self._evolution_model = evolution_model

    async def record(self, event: EvolutionEvent) -> str:
        """记录进化事件"""
        async with self._session_factory() as db:
            try:
                import uuid
                Model = self._evolution_model
                entry = Model(
                    id=uuid.uuid4(),
                    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()
                await db.refresh(entry)
                event_id = str(entry.id)
                event.event_id = event_id
                logger.info(f"Evolution event recorded: {event_id} for agent '{event.agent_name}'")
                return event_id
            except Exception as e:
                await db.rollback()
                logger.error(f"Failed to record evolution event: {e}")
                raise

    async def rollback(self, event_id: str) -> bool:
        """回滚进化事件"""
        async with self._session_factory() as db:
            try:
                import uuid
                from sqlalchemy import select
                Model = self._evolution_model

                stmt = select(Model).where(Model.id == uuid.UUID(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 Exception 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,
    ) -> list[dict]:
        """列出进化事件"""
        async with self._session_factory() as db:
            try:
                from sqlalchemy import select
                Model = self._evolution_model
                stmt = select(Model)

                if agent_name:
                    stmt = stmt.where(Model.agent_name == agent_name)
                if change_type:
                    stmt = stmt.where(Model.change_type == change_type)
                if status:
                    stmt = stmt.where(Model.status == status)

                stmt = stmt.order_by(Model.created_at.desc())
                result = await db.execute(stmt)
                entries = result.scalars().all()

                return [
                    {
                        "id": str(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 Exception as e:
                logger.error(f"Failed to list evolution events: {e}")
                return []


class PersistentEvolutionStore:
    """SQLite 持久化进化存储

    使用同步 SQLAlchemy + SQLite 实现持久化，通过 run_in_executor
    提供异步接口兼容性。
    """

    def __init__(self, db_path: str = "~/.agentkit/evolution.db"):
        self._db_path = os.path.expanduser(db_path)
        os.makedirs(os.path.dirname(self._db_path), exist_ok=True)
        self._engine = create_engine(f"sqlite:///{self._db_path}", echo=False)

        # Enable WAL mode for better concurrent read/write performance
        @sa_event.listens_for(self._engine, "connect")
        def _set_sqlite_pragma(dbapi_connection, connection_record):
            cursor = dbapi_connection.cursor()
            cursor.execute("PRAGMA journal_mode=WAL")
            cursor.close()

        Base.metadata.create_all(self._engine)
        self._Session = sessionmaker(bind=self._engine)

    # ── 内部辅助 ──────────────────────────────────────────

    def _run_sync(self, func: Any) -> Any:
        loop = asyncio.get_running_loop()
        return loop.run_in_executor(None, func)

    async def close(self) -> None:
        """Dispose the SQLAlchemy engine, releasing all pooled connections."""
        if self._engine is not None:
            await self._run_sync(self._engine.dispose)
            self._engine = None

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

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

    @staticmethod
    def _retry_locked(func, *args, max_retries: int = 5, base_delay: float = 0.05, **kwargs):
        """Retry a function on SQLite 'database is locked' OperationalError."""
        for attempt in range(max_retries):
            try:
                return func(*args, **kwargs)
            except OperationalError as exc:
                if "database is locked" not in str(exc).lower():
                    raise
                if attempt == max_retries - 1:
                    raise
                delay = base_delay * (2 ** attempt)
                time.sleep(delay)

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

    def _record_sync(self, event: EvolutionEvent) -> str:
        with self._Session() as session:
            event_id = str(_uuid.uuid4())
            entry = EvolutionEventModel(
                id=event_id,
                agent_name=event.agent_name,
                change_type=event.change_type,
                before=json.dumps(event.before, ensure_ascii=False),
                after=json.dumps(event.after, ensure_ascii=False),
                metrics=json.dumps(event.metrics, ensure_ascii=False) if event.metrics else None,
                status="active",
            )
            session.add(entry)
            session.commit()
            event.event_id = event_id
            logger.info(f"Evolution event recorded: {event_id} for agent '{event.agent_name}'")
            return event_id

    async def record(self, event: EvolutionEvent) -> str:
        """记录进化事件"""
        return await self._run_sync(lambda: self._retry_locked(self._record_sync, event))

    def _rollback_sync(self, event_id: str) -> bool:
        with self._Session() as session:
            stmt = select(EvolutionEventModel).where(EvolutionEventModel.id == event_id)
            entry = session.execute(stmt).scalar_one_or_none()
            if not entry:
                logger.error(f"Evolution event {event_id} not found")
                return False
            entry.status = "rolled_back"
            session.commit()
            logger.info(f"Evolution event {event_id} rolled back")
            return True

    async def rollback(self, event_id: str) -> bool:
        """回滚进化事件"""
        return await self._run_sync(lambda: self._retry_locked(self._rollback_sync, event_id))

    def _list_events_sync(
        self,
        agent_name: str | None = None,
        change_type: str | None = None,
        status: str | None = None,
    ) -> list[dict]:
        with self._Session() as session:
            stmt = select(EvolutionEventModel)
            if agent_name:
                stmt = stmt.where(EvolutionEventModel.agent_name == agent_name)
            if change_type:
                stmt = stmt.where(EvolutionEventModel.change_type == change_type)
            if status:
                stmt = stmt.where(EvolutionEventModel.status == status)
            stmt = stmt.order_by(EvolutionEventModel.created_at.desc())
            entries = session.execute(stmt).scalars().all()
            return [
                {
                    "id": e.id,
                    "agent_name": e.agent_name,
                    "change_type": e.change_type,
                    "before": json.loads(e.before) if e.before else None,
                    "after": json.loads(e.after) if e.after else None,
                    "metrics": json.loads(e.metrics) if e.metrics else None,
                    "status": e.status,
                    "created_at": e.created_at.isoformat() if e.created_at else None,
                }
                for e in entries
            ]

    async def list_events(
        self,
        agent_name: str | None = None,
        change_type: str | None = None,
        status: str | None = None,
    ) -> list[dict]:
        """列出进化事件"""
        return await self._run_sync(lambda: self._retry_locked(self._list_events_sync, agent_name, change_type, status))

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

    def _record_skill_version_sync(
        self, skill_name: str, version: str, content: str, parent_version: str | None = None
    ) -> str:
        with self._Session() as session:
            vid = str(_uuid.uuid4())
            entry = SkillVersionModel(
                id=vid,
                skill_name=skill_name,
                version=version,
                content=content,
                parent_version=parent_version,
            )
            session.add(entry)
            session.commit()
            return vid

    async def record_skill_version(
        self, skill_name: str, version: str, content: str, parent_version: str | None = None
    ) -> str:
        """记录技能版本"""
        return await self._run_sync(
            lambda: self._retry_locked(self._record_skill_version_sync, skill_name, version, content, parent_version)
        )

    def _list_skill_versions_sync(self, skill_name: str) -> list[dict]:
        with self._Session() as session:
            stmt = (
                select(SkillVersionModel)
                .where(SkillVersionModel.skill_name == skill_name)
                .order_by(SkillVersionModel.created_at.desc())
            )
            entries = session.execute(stmt).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
            ]

    async def list_skill_versions(self, skill_name: str) -> list[dict]:
        """列出技能版本历史"""
        return await self._run_sync(lambda: self._retry_locked(self._list_skill_versions_sync, skill_name))

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

    def _record_ab_test_result_sync(
        self, test_id: str, variant: str, score: float, sample_count: int = 0
    ) -> str:
        with self._Session() as session:
            rid = str(_uuid.uuid4())
            entry = ABTestResultModel(
                id=rid,
                test_id=test_id,
                variant=variant,
                score=score,
                sample_count=sample_count,
            )
            session.add(entry)
            session.commit()
            return rid

    async def record_ab_test_result(
        self, test_id: str, variant: str, score: float, sample_count: int = 0
    ) -> str:
        """记录 A/B 测试结果"""
        return await self._run_sync(
            lambda: self._retry_locked(self._record_ab_test_result_sync, test_id, variant, score, sample_count)
        )

    def _get_ab_test_results_sync(self, test_id: str) -> list[dict]:
        with self._Session() as session:
            stmt = select(ABTestResultModel).where(ABTestResultModel.test_id == test_id)
            entries = session.execute(stmt).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
            ]

    async def get_ab_test_results(self, test_id: str) -> list[dict]:
        """获取 A/B 测试结果"""
        return await self._run_sync(lambda: self._retry_locked(self._get_ab_test_results_sync, test_id))


class InMemoryEvolutionStore:
    """基于内存字典的进化存储（用于测试和轻量场景）"""

    def __init__(self) -> None:
        self._events: dict[str, dict] = {}
        self._skill_versions: dict[str, list[dict]] = {}
        self._ab_results: dict[str, list[dict]] = {}

    async def record(self, event: EvolutionEvent) -> str:
        """记录进化事件"""
        event_id = str(_uuid.uuid4())
        event.event_id = event_id
        self._events[event_id] = {
            "id": event_id,
            "agent_name": event.agent_name,
            "change_type": event.change_type,
            "before": event.before,
            "after": event.after,
            "metrics": event.metrics,
            "status": "active",
            "created_at": datetime.now(timezone.utc).isoformat(),
        }
        logger.info(f"Evolution event recorded: {event_id} for agent '{event.agent_name}'")
        return event_id

    async def rollback(self, event_id: str) -> bool:
        """回滚进化事件"""
        if event_id not in self._events:
            logger.error(f"Evolution event {event_id} not found")
            return False
        self._events[event_id]["status"] = "rolled_back"
        logger.info(f"Evolution event {event_id} rolled back")
        return True

    async def list_events(
        self,
        agent_name: str | None = None,
        change_type: str | None = None,
        status: str | None = None,
    ) -> list[dict]:
        """列出进化事件"""
        results = []
        for e in self._events.values():
            if agent_name and e["agent_name"] != agent_name:
                continue
            if change_type and e["change_type"] != change_type:
                continue
            if status and e["status"] != status:
                continue
            results.append(e)
        results.sort(key=lambda x: x["created_at"], reverse=True)
        return results

    async def record_skill_version(
        self, skill_name: str, version: str, content: str, parent_version: str | None = None
    ) -> str:
        """记录技能版本"""
        vid = str(_uuid.uuid4())
        entry = {
            "id": vid,
            "skill_name": skill_name,
            "version": version,
            "content": content,
            "parent_version": parent_version,
            "created_at": datetime.now(timezone.utc).isoformat(),
        }
        self._skill_versions.setdefault(skill_name, []).append(entry)
        return vid

    async def list_skill_versions(self, skill_name: str) -> list[dict]:
        """列出技能版本历史"""
        versions = self._skill_versions.get(skill_name, [])
        return sorted(versions, key=lambda x: x["created_at"], reverse=True)

    async def record_ab_test_result(
        self, test_id: str, variant: str, score: float, sample_count: int = 0
    ) -> str:
        """记录 A/B 测试结果"""
        rid = str(_uuid.uuid4())
        entry = {
            "id": rid,
            "test_id": test_id,
            "variant": variant,
            "score": score,
            "sample_count": sample_count,
            "created_at": datetime.now(timezone.utc).isoformat(),
        }
        self._ab_results.setdefault(test_id, []).append(entry)
        return rid

    async def get_ab_test_results(self, test_id: str) -> list[dict]:
        """获取 A/B 测试结果"""
        return self._ab_results.get(test_id, [])


def create_evolution_store(
    backend: str = "memory",
    db_path: str = "~/.agentkit/evolution.db",
    session_factory: Any = None,
    evolution_model: Any = None,
) -> EvolutionStore | PersistentEvolutionStore | InMemoryEvolutionStore:
    """工厂函数：创建进化存储实例

    Args:
        backend: 存储后端类型 - "memory" | "sqlite" | "sql"
        db_path: SQLite 数据库路径（仅 backend="sqlite" 时使用）
        session_factory: 异步 SQLAlchemy session 工厂（仅 backend="sql" 时使用）
        evolution_model: SQLAlchemy ORM 模型类（仅 backend="sql" 时使用）

    Returns:
        对应后端的进化存储实例
    """
    if backend == "sqlite":
        return PersistentEvolutionStore(db_path=db_path)
    elif backend == "sql" and session_factory and evolution_model:
        return EvolutionStore(session_factory=session_factory, evolution_model=evolution_model)
    else:
        return InMemoryEvolutionStore()
