"""Episodic Memory - 基于 pgvector + PostgreSQL 的任务经验记忆"""

import json
import logging
import math
from datetime import datetime, timezone
from typing import Any

from sqlalchemy import text

from agentkit.memory.base import Memory, MemoryItem
from agentkit.memory.embedder import Embedder
from agentkit.utils.vector_math import compute_cosine_similarity

logger = logging.getLogger(__name__)


class EpisodicMemory(Memory):
    """Episodic Memory - 记录每次任务的输入/输出/效果/反思

    基于 pgvector + PostgreSQL 实现，支持语义检索和时间衰减。
    生命周期：永久（可配置衰减）。

    当 pgvector_enabled=True 且 session_factory 可用时，search/retrieve
    使用 pgvector 原生 ``<=>`` 算符进行最近邻检索，再在 Python 侧做
    time_decay 重排；否则回退到客户端 O(N) cosine similarity。
    """

    def __init__(
        self,
        session_factory: Any,
        episodic_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 = "episodic_memories",
    ):
        """
        Args:
            session_factory: 返回 async context manager 的工厂
            episodic_model: EpisodicMemory ORM 模型类
            embedder: 嵌入器，用于生成向量
            decay_rate: 时间衰减率（越大衰减越快）
            alpha: 混合评分权重，alpha * cosine + (1-alpha) * time_decay
            retrieve_limit: retrieve() 时的最大候选行数（默认 200）
            pgvector_enabled: 是否使用 pgvector 原生 ``<=>`` 算符检索
            table_name: pgvector 查询使用的表名（默认 ``episodic_memories``）
        """
        self._session_factory = session_factory
        self._episodic_model = episodic_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

    async def store(self, key: str, value: Any, metadata: dict[str, Any] | None = None) -> None:
        """存储任务经验"""
        async with self._session_factory() as db:
            try:
                Model = self._episodic_model
                meta = metadata or {}

                # 生成 embedding
                embedding = None
                if self._embedder:
                    if isinstance(value, dict):
                        text = value.get("output_summary", "") or value.get("input_summary", "") or json.dumps(value, ensure_ascii=False)[:500]
                    else:
                        text = str(value)
                    embedding = await self._embedder.embed(text)

                entry = Model(
                    agent_name=meta.get("agent_name", ""),
                    task_type=meta.get("task_type", ""),
                    input_summary=str(value)[:500] if value else "",
                    output_summary=meta.get("output_summary", ""),
                    outcome=meta.get("outcome", "success"),
                    quality_score=meta.get("quality_score", 0.5),
                    reflection=meta.get("reflection", ""),
                    embedding=embedding,
                )
                db.add(entry)
                await db.commit()
            except Exception as e:
                await db.rollback()
                logger.error(f"Failed to store episodic memory: {e}")
                raise

    async def retrieve(self, key: str) -> MemoryItem | None:
        """按 key 语义检索（使用 embedding 相似度）"""
        if not self._embedder:
            return None

        query_embedding = await self._embedder.embed(key)

        async with self._session_factory() as db:
            try:
                if self._pgvector_enabled:
                    return await self._retrieve_pgvector(db, query_embedding)
                return await self._retrieve_client_side(db, query_embedding)
            except Exception as e:
                logger.error(f"Failed to retrieve episodic memory: {e}")
                return None

    async def _retrieve_pgvector(self, db: Any, query_embedding: list[float]) -> MemoryItem | None:
        """使用 pgvector ``<=>`` 算符检索最相似条目"""
        sql = text(
            f"SELECT * FROM {self._table_name} "
            f"ORDER BY embedding <=> :query_vec "
            f"LIMIT :lim"
        )
        result = await db.execute(sql, {"query_vec": str(query_embedding), "lim": 1})
        row = result.mappings().first()

        if row is None:
            return None

        # Compute cosine similarity for the returned row
        row_embedding = row.get("embedding")
        if row_embedding is None:
            return None

        cosine = compute_cosine_similarity(query_embedding, row_embedding)
        if cosine < 0.1:
            return None

        return MemoryItem(
            key=str(row.get("id", "")),
            value={
                "input_summary": row.get("input_summary", ""),
                "output_summary": row.get("output_summary", ""),
                "outcome": row.get("outcome", "success"),
                "quality_score": row.get("quality_score", 0.5),
                "reflection": row.get("reflection", ""),
            },
            metadata={
                "agent_name": row.get("agent_name", ""),
                "task_type": row.get("task_type", ""),
                "created_at": row["created_at"].isoformat() if row.get("created_at") else None,
                "cosine_similarity": cosine,
            },
            score=cosine,
            created_at=row.get("created_at") or datetime.now(timezone.utc),
        )

    async def _retrieve_client_side(self, db: Any, query_embedding: list[float]) -> MemoryItem | None:
        """客户端 O(N) cosine similarity 检索（回退路径）"""
        Model = self._episodic_model
        from sqlalchemy import select

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

        if not entries:
            return None

        best_item = None
        best_score = -1.0

        for entry in entries:
            entry_embedding = entry.embedding
            if entry_embedding is None:
                continue
            cosine = compute_cosine_similarity(query_embedding, entry_embedding)
            if cosine > best_score:
                best_score = cosine
                best_item = entry

        if best_item is None or best_score < 0.1:
            return None

        return MemoryItem(
            key=str(best_item.id),
            value={
                "input_summary": best_item.input_summary,
                "output_summary": best_item.output_summary,
                "outcome": best_item.outcome,
                "quality_score": best_item.quality_score,
                "reflection": best_item.reflection,
            },
            metadata={
                "agent_name": best_item.agent_name,
                "task_type": best_item.task_type,
                "created_at": best_item.created_at.isoformat() if best_item.created_at else None,
                "cosine_similarity": best_score,
            },
            score=best_score,
            created_at=best_item.created_at or datetime.now(timezone.utc),
        )

    async def search(self, query: str, top_k: int = 5, filters: dict[str, Any] | None = None, search_multiplier: int = 5) -> list[MemoryItem]:
        """语义检索相似历史案例

        Args:
            query: 搜索查询文本。
            top_k: 返回的最大结果数。
            filters: 可选过滤条件（agent_name, task_type, outcome）。
            search_multiplier: 预取行数倍数（fetch top_k * 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, filters, search_multiplier)
                return await self._search_client_side(db, query, top_k, filters, search_multiplier)
            except Exception as e:
                logger.error(f"Failed to search episodic memory: {e}")
                return []

    async def _search_pgvector(
        self,
        db: Any,
        query: str,
        top_k: int,
        filters: dict[str, Any] | None,
        search_multiplier: int,
    ) -> list[MemoryItem]:
        """使用 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}

        filters = filters or {}
        if filters.get("agent_name"):
            where_clauses.append("agent_name = :agent_name")
            params["agent_name"] = filters["agent_name"]
        if filters.get("task_type"):
            where_clauses.append("task_type = :task_type")
            params["task_type"] = filters["task_type"]
        if filters.get("outcome"):
            where_clauses.append("outcome = :outcome")
            params["outcome"] = filters["outcome"]

        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 = []
        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("quality_score") 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

            items.append(MemoryItem(
                key=str(row.get("id", "")),
                value={
                    "input_summary": row.get("input_summary", ""),
                    "output_summary": row.get("output_summary", ""),
                    "outcome": row.get("outcome", "success"),
                    "quality_score": row.get("quality_score", 0.5),
                    "reflection": row.get("reflection", ""),
                },
                metadata={
                    "agent_name": row.get("agent_name", ""),
                    "task_type": row.get("task_type", ""),
                    "created_at": row["created_at"].isoformat() if row.get("created_at") else None,
                },
                score=score,
                created_at=row.get("created_at") or datetime.now(timezone.utc),
            ))

        items.sort(key=lambda x: x.score, reverse=True)
        return items[:top_k]

    async def _search_client_side(
        self,
        db: Any,
        query: str,
        top_k: int,
        filters: dict[str, Any] | None,
        search_multiplier: int,
    ) -> list[MemoryItem]:
        """客户端 O(N) cosine similarity 检索（回退路径）"""
        Model = self._episodic_model
        filters = filters or {}

        from sqlalchemy import select
        stmt = select(Model)

        if filters.get("agent_name"):
            stmt = stmt.where(Model.agent_name == filters["agent_name"])
        if filters.get("task_type"):
            stmt = stmt.where(Model.task_type == filters["task_type"])
        if filters.get("outcome"):
            stmt = stmt.where(Model.outcome == filters["outcome"])

        stmt = stmt.order_by(Model.created_at.desc()).limit(top_k * search_multiplier)

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

        # 如果有 embedder，生成 query embedding
        query_embedding = None
        if self._embedder and entries:
            query_embedding = await self._embedder.embed(query)

        # 计算得分并构建 MemoryItem
        items = []
        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.quality_score or 0.5) * decay

            # 混合评分：alpha * cosine + (1 - alpha) * time_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

            items.append(MemoryItem(
                key=str(entry.id),
                value={
                    "input_summary": entry.input_summary,
                    "output_summary": entry.output_summary,
                    "outcome": entry.outcome,
                    "quality_score": entry.quality_score,
                    "reflection": entry.reflection,
                },
                metadata={
                    "agent_name": entry.agent_name,
                    "task_type": entry.task_type,
                    "created_at": entry.created_at.isoformat() if entry.created_at else None,
                },
                score=score,
                created_at=entry.created_at or datetime.now(timezone.utc),
            ))

        items.sort(key=lambda x: x.score, reverse=True)
        if len(items) < top_k:
            logger.warning(
                "EpisodicMemory.search returned %d results after scoring (top_k=%d). "
                "Consider increasing search_multiplier (current=%d) to avoid missing relevant entries.",
                len(items), top_k, search_multiplier,
            )
        return items[:top_k]

    async def delete(self, key: str) -> bool:
        """删除指定经验"""
        async with self._session_factory() as db:
            try:
                from sqlalchemy import select, delete as sql_delete
                import uuid
                Model = self._episodic_model

                stmt = sql_delete(Model).where(Model.id == uuid.UUID(key))
                await db.execute(stmt)
                await db.commit()
                return True
            except Exception as e:
                await db.rollback()
                logger.error(f"Failed to delete episodic memory: {e}")
                return False
