"""ExperienceStore - 任务经验存储

提供两种后端实现：
- ExperienceStore: 基于 PostgreSQL + pgvector 的语义检索存储
- InMemoryExperienceStore: 基于内存字典的轻量存储（用于测试）

存储任务执行经验（成功路径、失败原因、耗时分布），
支持按任务类型检索和语义搜索，追踪完成率/耗时/重试率趋势。
"""

from __future__ import annotations

import logging
import math
import re
import uuid
from datetime import datetime, timedelta, timezone
from typing import Any

from sqlalchemy import text

_SAFE_TABLE_NAME_PATTERN = re.compile(r'^[a-zA-Z_][a-zA-Z0-9_]*$')

from agentkit.evolution.experience_schema import EvolutionMetrics, TaskExperience
from agentkit.memory.embedder import Embedder
from agentkit.utils.vector_math import compute_cosine_similarity

logger = logging.getLogger(__name__)


class ExperienceStore:
    """任务经验存储 - PostgreSQL + pgvector 混合存储

    基于 pgvector 向量索引 + tsvector 全文索引，
    支持精确匹配 task_type + 语义相似度排序 + 时效性衰减。

    检索策略：
    1. pgvector ``<=>`` 算符进行最近邻检索
    2. Python 侧 time_decay 重排
    3. 混合评分：alpha * cosine + (1 - alpha) * time_decay_score

    当 pgvector_enabled=False 或 embedder 不可用时，
    回退到客户端 O(N) cosine similarity。
    """

    def __init__(
        self,
        session_factory: Any,
        experience_model: Any,
        embedder: Embedder | None = None,
        decay_rate: float = 0.01,
        alpha: float = 0.7,
        retrieve_limit: int = 200,
        pgvector_enabled: bool = True,
        table_name: str = "task_experiences",
    ):
        """
        Args:
            session_factory: 返回 async context manager 的工厂
            experience_model: TaskExperience ORM 模型类
            embedder: 嵌入器，用于生成向量
            decay_rate: 时间衰减率（越大衰减越快）
            alpha: 混合评分权重，alpha * cosine + (1-alpha) * time_decay
            retrieve_limit: 客户端检索时的最大候选行数
            pgvector_enabled: 是否使用 pgvector 原生 ``<=>`` 算符检索
            table_name: pgvector 查询使用的表名
        """
        self._session_factory = session_factory
        self._experience_model = experience_model
        self._embedder = embedder
        self._decay_rate = decay_rate
        self._alpha = alpha
        self._retrieve_limit = retrieve_limit
        self._pgvector_enabled = pgvector_enabled
        self._table_name = table_name
        if not _SAFE_TABLE_NAME_PATTERN.match(self._table_name):
            raise ValueError(f"Invalid table_name: {self._table_name}. Must match [a-zA-Z_][a-zA-Z0-9_]*")

    async def record_experience(self, experience: TaskExperience) -> str:
        """记录任务经验

        如果 experience.embedding 为 None 且 embedder 可用，
        自动生成 embedding。

        Args:
            experience: 任务经验数据

        Returns:
            经验 ID
        """
        if not experience.experience_id:
            experience.experience_id = str(uuid.uuid4())

        # 自动生成 embedding
        if experience.embedding is None and self._embedder is not None:
            text = experience.text_for_embedding()
            try:
                experience.embedding = await self._embedder.embed(text)
            except Exception as e:
                logger.warning(f"Failed to generate embedding for experience {experience.experience_id}: {e}")

        async with self._session_factory() as db:
            try:
                Model = self._experience_model
                entry = Model(
                    id=experience.experience_id,
                    task_type=experience.task_type,
                    goal=experience.goal,
                    steps_summary=experience.steps_summary,
                    outcome=experience.outcome,
                    duration_seconds=experience.duration_seconds,
                    success_rate=experience.success_rate,
                    failure_reasons=experience.failure_reasons,
                    optimization_tips=experience.optimization_tips,
                    embedding=experience.embedding,
                    created_at=experience.created_at,
                )
                db.add(entry)
                await db.commit()
                logger.info(
                    f"Experience recorded: {experience.experience_id} "
                    f"task_type={experience.task_type} outcome={experience.outcome}"
                )
                return experience.experience_id
            except Exception as e:
                await db.rollback()
                logger.error(f"Failed to record experience: {e}")
                raise

    async def search(
        self,
        query: str,
        top_k: int = 5,
        task_type: str | None = None,
        search_multiplier: int = 5,
    ) -> list[TaskExperience]:
        """语义检索相似经验

        支持精确匹配 task_type + 语义相似度排序 + 时效性衰减。

        Args:
            query: 搜索查询文本
            top_k: 返回的最大结果数
            task_type: 可选的任务类型过滤
            search_multiplier: 预取行数倍数
        """
        async with self._session_factory() as db:
            try:
                if self._pgvector_enabled and self._embedder:
                    return await self._search_pgvector(db, query, top_k, task_type, search_multiplier)
                return await self._search_client_side(db, query, top_k, task_type, search_multiplier)
            except Exception as e:
                logger.error(f"Failed to search experiences: {e}")
                return []

    async def _search_pgvector(
        self,
        db: Any,
        query: str,
        top_k: int,
        task_type: str | None,
        search_multiplier: int,
    ) -> list[TaskExperience]:
        """使用 pgvector ``<=>`` 算符检索，再 Python 侧 time_decay 重排"""
        query_embedding = await self._embedder.embed(query)
        fetch_limit = top_k * search_multiplier

        where_clauses = []
        params: dict[str, Any] = {"query_vec": str(query_embedding), "lim": fetch_limit}

        if task_type:
            where_clauses.append("task_type = :task_type")
            params["task_type"] = task_type

        where_sql = (" WHERE " + " AND ".join(where_clauses)) if where_clauses else ""
        sql = text(
            f"SELECT *, embedding <=> :query_vec AS distance "
            f"FROM {self._table_name}{where_sql} "
            f"ORDER BY embedding <=> :query_vec "
            f"LIMIT :lim"
        )

        result = await db.execute(sql, params)
        rows = result.mappings().all()

        if not rows:
            return []

        # Re-rank with time_decay in Python
        items: list[tuple[float, TaskExperience]] = []
        for row in rows:
            row_embedding = row.get("embedding")
            age_hours = (
                (datetime.now(timezone.utc) - row["created_at"]).total_seconds() / 3600
                if row.get("created_at")
                else 0
            )
            decay = math.exp(-self._decay_rate * age_hours)
            time_decay_score = (row.get("success_rate") or 0.5) * decay

            if row_embedding is not None:
                cosine_sim = compute_cosine_similarity(query_embedding, row_embedding)
                score = self._alpha * cosine_sim + (1 - self._alpha) * time_decay_score
            else:
                score = time_decay_score

            exp = TaskExperience(
                experience_id=str(row.get("id", "")),
                task_type=row.get("task_type", ""),
                goal=row.get("goal", ""),
                steps_summary=row.get("steps_summary", ""),
                outcome=row.get("outcome", "success"),
                duration_seconds=row.get("duration_seconds", 0.0),
                success_rate=row.get("success_rate", 1.0),
                failure_reasons=row.get("failure_reasons") or [],
                optimization_tips=row.get("optimization_tips") or [],
                embedding=row_embedding,
                created_at=row.get("created_at") or datetime.now(timezone.utc),
            )
            items.append((score, exp))

        items.sort(key=lambda x: x[0], reverse=True)
        return [exp for _, exp in items[:top_k]]

    async def _search_client_side(
        self,
        db: Any,
        query: str,
        top_k: int,
        task_type: str | None,
        search_multiplier: int,
    ) -> list[TaskExperience]:
        """客户端 O(N) cosine similarity 检索（回退路径）"""
        Model = self._experience_model
        from sqlalchemy import select

        stmt = select(Model)
        if task_type:
            stmt = stmt.where(Model.task_type == task_type)
        stmt = stmt.order_by(Model.created_at.desc()).limit(top_k * search_multiplier)

        result = await db.execute(stmt)
        entries = result.scalars().all()

        query_embedding = None
        if self._embedder and entries:
            query_embedding = await self._embedder.embed(query)

        items: list[tuple[float, TaskExperience]] = []
        for entry in entries:
            age_hours = (
                (datetime.now(timezone.utc) - entry.created_at).total_seconds() / 3600
                if entry.created_at
                else 0
            )
            decay = math.exp(-self._decay_rate * age_hours)
            time_decay_score = (entry.success_rate or 0.5) * decay

            if self._embedder and query_embedding is not None and entry.embedding is not None:
                cosine_sim = compute_cosine_similarity(query_embedding, entry.embedding)
                score = self._alpha * cosine_sim + (1 - self._alpha) * time_decay_score
            else:
                score = time_decay_score

            exp = TaskExperience(
                experience_id=str(entry.id),
                task_type=entry.task_type,
                goal=entry.goal,
                steps_summary=entry.steps_summary,
                outcome=entry.outcome,
                duration_seconds=entry.duration_seconds,
                success_rate=entry.success_rate,
                failure_reasons=entry.failure_reasons or [],
                optimization_tips=entry.optimization_tips or [],
                embedding=entry.embedding,
                created_at=entry.created_at or datetime.now(timezone.utc),
            )
            items.append((score, exp))

        items.sort(key=lambda x: x[0], reverse=True)
        return [exp for _, exp in items[:top_k]]

    async def get_metrics(
        self,
        task_type: str | None = None,
        time_window: str = "24h",
    ) -> list[EvolutionMetrics]:
        """获取进化指标趋势

        按任务类型和时间窗口聚合完成率、平均耗时和重试率。

        Args:
            task_type: 可选的任务类型过滤，None 表示所有类型
            time_window: 时间窗口（"1h", "24h", "7d", "30d"）
        """
        window_delta = _parse_time_window(time_window)
        window_start = datetime.now(timezone.utc) - window_delta
        window_end = datetime.now(timezone.utc)

        async with self._session_factory() as db:
            try:
                where_clauses = ["created_at >= :window_start"]
                params: dict[str, Any] = {"window_start": window_start}

                if task_type:
                    where_clauses.append("task_type = :task_type")
                    params["task_type"] = task_type

                where_sql = " AND ".join(where_clauses)

                # 按任务类型聚合
                group_by = "task_type" if task_type is None else ""
                select_clause = "task_type"
                if task_type:
                    select_clause += f", '{task_type}' as filtered_task_type"

                sql = text(
                    f"SELECT task_type, "
                    f"  COUNT(*) as sample_count, "
                    f"  AVG(CASE WHEN outcome = 'success' THEN 1.0 ELSE 0.0 END) as completion_rate, "
                    f"  AVG(duration_seconds) as avg_duration, "
                    f"  AVG(CASE WHEN success_rate < 1.0 THEN 1.0 ELSE 0.0 END) as retry_rate "
                    f"FROM {self._table_name} "
                    f"WHERE {where_sql} "
                    f"GROUP BY task_type"
                )

                result = await db.execute(sql, params)
                rows = result.mappings().all()

                metrics_list = []
                for row in rows:
                    metrics_list.append(
                        EvolutionMetrics(
                            task_type=row["task_type"],
                            time_window=time_window,
                            completion_rate=row["completion_rate"] or 0.0,
                            avg_duration=row["avg_duration"] or 0.0,
                            retry_rate=row["retry_rate"] or 0.0,
                            sample_count=row["sample_count"] or 0,
                            window_start=window_start,
                            window_end=window_end,
                        )
                    )
                return metrics_list
            except Exception as e:
                logger.error(f"Failed to get metrics: {e}")
                return []


class InMemoryExperienceStore:
    """基于内存字典的任务经验存储（用于测试和轻量场景）

    无需数据库，纯 dict-based 实现，支持与 ExperienceStore 相同的接口。
    """

    def __init__(
        self,
        embedder: Embedder | None = None,
        decay_rate: float = 0.01,
        alpha: float = 0.7,
    ):
        self._embedder = embedder
        self._decay_rate = decay_rate
        self._alpha = alpha
        self._experiences: dict[str, TaskExperience] = {}

    async def record_experience(self, experience: TaskExperience) -> str:
        """记录任务经验"""
        if not experience.experience_id:
            experience.experience_id = str(uuid.uuid4())

        # 自动生成 embedding
        if experience.embedding is None and self._embedder is not None:
            text = experience.text_for_embedding()
            try:
                experience.embedding = await self._embedder.embed(text)
            except Exception as e:
                logger.warning(f"Failed to generate embedding for experience {experience.experience_id}: {e}")

        # 存储副本，避免外部修改影响内部状态
        self._experiences[experience.experience_id] = TaskExperience(
            experience_id=experience.experience_id,
            task_type=experience.task_type,
            goal=experience.goal,
            steps_summary=experience.steps_summary,
            outcome=experience.outcome,
            duration_seconds=experience.duration_seconds,
            success_rate=experience.success_rate,
            failure_reasons=list(experience.failure_reasons),
            optimization_tips=list(experience.optimization_tips),
            embedding=experience.embedding,
            created_at=experience.created_at,
        )
        logger.info(
            f"Experience recorded: {experience.experience_id} "
            f"task_type={experience.task_type} outcome={experience.outcome}"
        )
        return experience.experience_id

    async def search(
        self,
        query: str,
        top_k: int = 5,
        task_type: str | None = None,
        search_multiplier: int = 5,
    ) -> list[TaskExperience]:
        """语义检索相似经验"""
        # 生成 query embedding
        query_embedding = None
        if self._embedder:
            try:
                query_embedding = await self._embedder.embed(query)
            except Exception as e:
                logger.warning(f"Failed to generate query embedding: {e}")

        # 筛选候选
        candidates = list(self._experiences.values())
        if task_type:
            candidates = [e for e in candidates if e.task_type == task_type]

        # 计算得分
        items: list[tuple[float, TaskExperience]] = []
        for exp in candidates:
            age_hours = (
                (datetime.now(timezone.utc) - exp.created_at).total_seconds() / 3600
                if exp.created_at
                else 0
            )
            decay = math.exp(-self._decay_rate * age_hours)
            time_decay_score = exp.success_rate * decay

            if query_embedding is not None and exp.embedding is not None:
                cosine_sim = compute_cosine_similarity(query_embedding, exp.embedding)
                score = self._alpha * cosine_sim + (1 - self._alpha) * time_decay_score
            else:
                score = time_decay_score

            items.append((score, exp))

        items.sort(key=lambda x: x[0], reverse=True)
        return [exp for _, exp in items[:top_k]]

    async def get_metrics(
        self,
        task_type: str | None = None,
        time_window: str = "24h",
    ) -> list[EvolutionMetrics]:
        """获取进化指标趋势"""
        window_delta = _parse_time_window(time_window)
        window_start = datetime.now(timezone.utc) - window_delta
        window_end = datetime.now(timezone.utc)

        # 筛选时间窗口内的经验
        candidates = [
            e for e in self._experiences.values()
            if e.created_at >= window_start
        ]
        if task_type:
            candidates = [e for e in candidates if e.task_type == task_type]

        # 按 task_type 分组聚合
        groups: dict[str, list[TaskExperience]] = {}
        for exp in candidates:
            groups.setdefault(exp.task_type, []).append(exp)

        metrics_list = []
        for tt, exps in groups.items():
            n = len(exps)
            if n == 0:
                continue
            completion_rate = sum(1 for e in exps if e.outcome == "success") / n
            avg_duration = sum(e.duration_seconds for e in exps) / n
            retry_rate = sum(1 for e in exps if e.success_rate < 1.0) / n

            metrics_list.append(
                EvolutionMetrics(
                    task_type=tt,
                    time_window=time_window,
                    completion_rate=completion_rate,
                    avg_duration=avg_duration,
                    retry_rate=retry_rate,
                    sample_count=n,
                    window_start=window_start,
                    window_end=window_end,
                )
            )
        return metrics_list


# ── 辅助函数 ──────────────────────────────────────────────


def _parse_time_window(window: str) -> timedelta:
    """解析时间窗口字符串为 timedelta

    支持格式: "1h", "24h", "7d", "30d"
    """
    unit = window[-1].lower()
    try:
        value = int(window[:-1])
    except ValueError:
        return timedelta(hours=24)
    if unit == "h":
        return timedelta(hours=value)
    elif unit == "d":
        return timedelta(days=value)
    else:
        logger.warning(f"Unknown time window unit '{unit}', defaulting to 24h")
        return timedelta(hours=24)
