from dataclasses import dataclass
from typing import Dict, List, Any, Optional
from datetime import datetime
import sqlite3
import json


@dataclass
class Transition:
    """Transition from one state to another"""

    generation_id: str
    operation_id: str
    state_before: Dict[str, Any]
    state_after: Dict[str, Any]
    request: Dict[str, Any]
    response: Dict[str, Any]
    status_code: int
    operation_selection: Optional[List[Dict[str, Any]]] = None
    request_generation: Optional[List[Dict[str, Any]]] = None
    execution_analysis: Optional[Dict[str, Any]] = None
    target_matching: Optional[Dict[str, Any]] = None
    resource_manager: Optional[Dict[str, Any]] = None
    timestamp: Optional[datetime] = None
    error: Optional[str] = None
    error_traceback: Optional[str] = None

    def to_dict(self) -> Dict[str, Any]:
        """Convert transition to dict"""
        return {
            "generation_id": self.generation_id,
            "operation_id": self.operation_id,
            "state_before": self.state_before,
            "state_after": self.state_after,
            "request": self.request,
            "response": self.response,
            "status_code": self.status_code,
            "operation_selection": self.operation_selection,
            "request_generation": self.request_generation,
            "execution_analysis": self.execution_analysis,
            "target_matching": self.target_matching,
            "resource_manager": self.resource_manager,
            "timestamp": self.timestamp.isoformat() if self.timestamp else None,
            "error": self.error,
            "error_traceback": self.error_traceback,
        }

    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "Transition":
        """Create transition from dict"""
        if "timestamp" in data and data["timestamp"]:
            data["timestamp"] = datetime.fromisoformat(data["timestamp"])
        return cls(**data)


class TransitionStore:
    def __init__(self, db_path: str = "transitions.db"):
        self.db_path = db_path
        # 对于内存数据库，保持连接
        if db_path == ":memory:":
            self.conn = sqlite3.connect(db_path)
        else:
            self.conn = None
        self._init_db()

    def _get_conn(self) -> sqlite3.Connection:
        """获取数据库连接"""
        if self.db_path == ":memory:":
            return self.conn
        return sqlite3.connect(self.db_path)

    def _init_db(self):
        """Initialize database"""
        with self._get_conn() as conn:
            cursor = conn.cursor()
            cursor.execute(
                """
                CREATE TABLE IF NOT EXISTS transitions (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    generation_id TEXT NOT NULL,
                    operation_id TEXT NOT NULL,
                    state_before TEXT NOT NULL,
                    state_after TEXT NOT NULL,
                    -- Operation Selection
                    operation_selection TEXT,
                    -- Request Generation
                    request_generation TEXT,
                    -- Execution
                    request TEXT NOT NULL,
                    response TEXT NOT NULL,
                    status_code INTEGER NOT NULL,
                    -- Execution Analysis
                    execution_analysis TEXT,
                    -- Target Matching
                    target_matching TEXT,
                    -- Resource Manager
                    resource_manager TEXT,
                    -- Metadata
                    timestamp DATETIME NOT NULL,
                    error TEXT,
                    error_traceback TEXT,
                    -- Indexes
                    UNIQUE(generation_id, operation_id, timestamp)
                )
                """
            )
            cursor.execute(
                """
                CREATE INDEX IF NOT EXISTS idx_transitions_generation_id ON transitions(generation_id)
                """
            )
            cursor.execute(
                """
                CREATE INDEX IF NOT EXISTS idx_transitions_operation_id ON transitions(operation_id)
                """
            )

    def store(self, transition: Transition):
        """Store a transition record"""
        data = transition.to_dict()
        with self._get_conn() as conn:
            cursor = conn.cursor()
            cursor.execute(
                """
                INSERT INTO transitions (
                    generation_id, operation_id,
                    state_before, state_after,
                    operation_selection, request_generation,
                    request, response, status_code,
                    execution_analysis, target_matching, resource_manager,
                    timestamp, error, error_traceback
                ) VALUES (
                    ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?
                )
                """,
                (
                    data["generation_id"],
                    data["operation_id"],
                    json.dumps(data["state_before"]),
                    json.dumps(data["state_after"]),
                    json.dumps(data["operation_selection"])
                    if data["operation_selection"]
                    else None,
                    json.dumps(data["request_generation"])
                    if data["request_generation"]
                    else None,
                    json.dumps(data["request"]),
                    json.dumps(data["response"]),
                    data["status_code"],
                    json.dumps(data["execution_analysis"])
                    if data["execution_analysis"]
                    else None,
                    json.dumps(data["target_matching"])
                    if data["target_matching"]
                    else None,
                    json.dumps(data["resource_manager"])
                    if data["resource_manager"]
                    else None,
                    data["timestamp"],
                    data["error"],
                    data["error_traceback"],
                ),
            )
            conn.commit()

    def get_transitions_by_operation(self, operation_id: str) -> List[Transition]:
        """Get all transitions for a specific operation"""
        with self._get_conn() as conn:
            cursor = conn.cursor()
            cursor.execute(
                """
                SELECT 
                    generation_id, operation_id,
                    state_before, state_after,
                    operation_selection, request_generation,
                    request, response, status_code,
                    execution_analysis, target_matching,
                    resource_manager,
                    timestamp, error, error_traceback
                FROM transitions
                WHERE operation_id = ?
                ORDER BY timestamp ASC
                """,
                (operation_id,),
            )
            return [
                Transition.from_dict(
                    {
                        "generation_id": row[0],
                        "operation_id": row[1],
                        "state_before": json.loads(row[2]),
                        "state_after": json.loads(row[3]),
                        "operation_selection": json.loads(row[4]) if row[4] else None,
                        "request_generation": json.loads(row[5]) if row[5] else None,
                        "request": json.loads(row[6]),
                        "response": json.loads(row[7]),
                        "status_code": row[8],
                        "execution_analysis": json.loads(row[9]) if row[9] else None,
                        "target_matching": json.loads(row[10]) if row[10] else None,
                        "resource_manager": json.loads(row[11]) if row[11] else None,
                        "timestamp": row[12],
                        "error": row[13],
                        "error_traceback": row[14],
                    }
                )
                for row in cursor.fetchall()
            ]

    def get_transitions_by_generation(self, generation_id: str) -> List[Transition]:
        """Get all transitions for a specific generation task"""
        with self._get_conn() as conn:
            cursor = conn.cursor()
            cursor.execute(
                """
                SELECT 
                    generation_id, operation_id,
                    state_before, state_after,
                    operation_selection, request_generation,
                    request, response, status_code,
                    execution_analysis, target_matching,
                    resource_manager,
                    timestamp, error, error_traceback
                FROM transitions
                WHERE generation_id = ?
                ORDER BY timestamp ASC
                """,
                (generation_id,),
            )
            return [
                Transition.from_dict(
                    {
                        "generation_id": row[0],
                        "operation_id": row[1],
                        "state_before": json.loads(row[2]),
                        "state_after": json.loads(row[3]),
                        "operation_selection": json.loads(row[4]) if row[4] else None,
                        "request_generation": json.loads(row[5]) if row[5] else None,
                        "request": json.loads(row[6]),
                        "response": json.loads(row[7]),
                        "status_code": row[8],
                        "execution_analysis": json.loads(row[9]) if row[9] else None,
                        "target_matching": json.loads(row[10]) if row[10] else None,
                        "resource_manager": json.loads(row[11]) if row[11] else None,
                        "timestamp": row[12],
                        "error": row[13],
                        "error_traceback": row[14],
                    }
                )
                for row in cursor.fetchall()
            ]

    def get_success_rate(self, operation_id: str) -> float:
        """Get the historical success rate of an operation"""
        with self._get_conn() as conn:
            cursor = conn.cursor()
            cursor.execute(
                """
                SELECT COUNT(*) as total,
                       SUM(CASE WHEN status_code >= 200 AND status_code < 300 THEN 1 ELSE 0 END) as success
                FROM transitions
                WHERE operation_id = ?
                """,
                (operation_id,),
            )
            total, success = cursor.fetchone()
            return success / total if total > 0 else 0.0

    def get_average_similarity_improvement(
        self,
        operation_id: str,
        current_state: Dict[str, Any],
        target_operation: Dict[str, Any],
    ) -> float:
        """计算操作的平均状态相似度提升

        通过比较历史记录中，执行该操作前后状态与目标状态的相似度变化，
        来评估该操作对达到目标状态的帮助程度。

        Args:
            operation_id: 操作ID
            current_state: 当前状态
            target_operation: 目标状态

        Returns:
            平均相似度提升，范围 [-1, 1]
        """
        conn = self._get_conn()
        try:
            cursor = conn.execute(
                """
                SELECT state_before, state_after
                FROM transitions
                WHERE operation_id = ?
                ORDER BY timestamp DESC
                LIMIT 10
            """,
                (operation_id,),
            )

            improvements = []
            for row in cursor.fetchall():
                state_before = json.loads(row[0])
                state_after = json.loads(row[1])

                # 计算状态相似度的变化
                before_similarity = self._calculate_state_similarity(
                    state_before, target_operation
                )
                after_similarity = self._calculate_state_similarity(
                    state_after, target_operation
                )
                improvement = after_similarity - before_similarity
                improvements.append(improvement)

            # 如果没有历史记录，返回中性值
            if not improvements:
                return 0.0

            # 返回平均提升
            return sum(improvements) / len(improvements)
        finally:
            if self.db_path != ":memory:":
                conn.close()

    def _calculate_state_similarity(
        self, state1: Dict[str, Any], state2: Dict[str, Any]
    ) -> float:
        """计算两个状态的相似度

        目前使用一个简单的实现：
        1. 比较资源类型和数量
        2. 比较资源状态
        3. 比较关键属性

        Returns:
            相似度分数，范围 [0, 1]
        """
        try:
            # 验证状态格式
            if "resources" not in state1 or "resources" not in state2:
                return 0.0

            # 1. 资源类型和数量的相似度
            resources1 = {
                r["type"]: r["state"]["count"] for r in state1.get("resources", [])
            }
            resources2 = {
                r["type"]: r["state"]["count"] for r in state2.get("resources", [])
            }

            # 计算资源类型的交集和并集
            common_types = set(resources1.keys()) & set(resources2.keys())
            all_types = set(resources1.keys()) | set(resources2.keys())

            if not all_types:
                return 1.0  # 如果都没有资源，认为完全相似

            # 如果没有共同的资源类型，认为完全不相似
            if not common_types:
                return 0.0

            # 计算资源数量的差异
            count_diff = sum(
                abs(resources1.get(t, 0) - resources2.get(t, 0)) for t in all_types
            )
            max_count = sum(
                max(resources1.get(t, 0), resources2.get(t, 0)) for t in all_types
            )

            # 综合考虑类型匹配度和数量差异
            type_similarity = len(common_types) / len(all_types)
            count_similarity = 1 - (
                count_diff / (2 * max_count) if max_count > 0 else 0
            )

            return 0.7 * type_similarity + 0.3 * count_similarity

        except Exception:
            return 0.0  # 如果计算出错，返回最低相似度

    def get_recent_success_rate(self, operation_id: str, n: int = 10) -> float:
        """获取最近 N 次操作的成功率"""
        conn = self._get_conn()
        try:
            cursor = conn.execute(
                """
                WITH recent_transitions AS (
                    SELECT *
                    FROM transitions
                    WHERE operation_id = ?
                    ORDER BY timestamp DESC
                    LIMIT ?
                )
                SELECT COUNT(*) as total,
                       SUM(CASE WHEN status_code >= 200 AND status_code < 300
                           THEN 1 ELSE 0 END) as success
                FROM recent_transitions
            """,
                (operation_id, n),
            )
            row = cursor.fetchone()
            total, success = row
            return success / total if total > 0 else 0.0
        finally:
            if self.db_path != ":memory:":
                conn.close()

    def get_best_operations_for_state(
        self, state: Dict[str, Any], limit: int = 3
    ) -> List[Dict[str, Any]]:
        """获取特定状态下最成功的操作"""
        state_json = json.dumps(state)
        conn = self._get_conn()
        try:
            cursor = conn.execute(
                """
                WITH state_transitions AS (
                    SELECT operation_id,
                           COUNT(*) as total,
                           SUM(CASE WHEN status_code >= 200 AND status_code < 300 THEN 1 ELSE 0 END) as success
                    FROM transitions
                    WHERE state_before = ?
                    GROUP BY operation_id
                    HAVING total >= 3  -- 至少有3次尝试
                )
                SELECT operation_id,
                       CAST(success AS FLOAT) / total as success_rate,
                       total as attempts
                FROM state_transitions
                ORDER BY success_rate DESC, attempts DESC
                LIMIT ?
            """,
                (state_json, limit),
            )
            return [
                {"operation_id": row[0], "success_rate": row[1], "attempts": row[2]}
                for row in cursor.fetchall()
            ]
        finally:
            if self.db_path != ":memory:":
                conn.close()

    def get_operation_avg_response_time(
        self, operation_id: str, recent_n: Optional[int] = None
    ) -> float:
        """获取操作的平均响应时间（毫秒）"""
        conn = self._get_conn()
        try:
            # 添加时间计算函数
            def time_diff_ms(t1: str, t2: str) -> float:
                try:
                    if not t1 or not t2:
                        return 0.0
                    dt1 = datetime.fromisoformat(t1)
                    dt2 = datetime.fromisoformat(t2)
                    return (dt2 - dt1).total_seconds() * 1000
                except (ValueError, TypeError):
                    return 0.0

            conn.create_function("time_diff_ms", 2, time_diff_ms)

            if recent_n:
                query = """
                    WITH recent_ops AS (
                        SELECT *
                        FROM transitions
                        WHERE operation_id = ?
                        ORDER BY timestamp DESC
                        LIMIT ?
                    ),
                    operation_times AS (
                        SELECT time_diff_ms(
                            LAG(timestamp) OVER (PARTITION BY generation_id ORDER BY timestamp),
                            timestamp
                        ) as response_time
                        FROM recent_ops
                    )
                    SELECT AVG(response_time)
                    FROM operation_times
                    WHERE response_time IS NOT NULL
                      AND response_time > 0
                """
                cursor = conn.execute(query, (operation_id, recent_n))
            else:
                query = """
                    WITH operation_times AS (
                        SELECT time_diff_ms(
                            LAG(timestamp) OVER (PARTITION BY generation_id ORDER BY timestamp),
                            timestamp
                        ) as response_time
                        FROM transitions
                        WHERE operation_id = ?
                    )
                    SELECT AVG(response_time)
                    FROM operation_times
                    WHERE response_time IS NOT NULL
                      AND response_time > 0
                """
                cursor = conn.execute(query, (operation_id,))

            result = cursor.fetchone()[0]
            return result if result is not None else 0.0
        finally:
            if self.db_path != ":memory:":
                conn.close()

    def get_operation_stats(self, operation_id: str) -> Dict[str, Any]:
        """Get statistics for a specific operation"""
        with self._get_conn() as conn:
            cursor = conn.cursor()
            cursor.execute(
                """
                SELECT COUNT(*) as total,
                       SUM(CASE WHEN status_code >= 200 AND status_code < 300 THEN 1 ELSE 0 END) as success_count,
                       MIN(timestamp) as first_seen,
                       MAX(timestamp) as last_seen
                FROM transitions
                WHERE operation_id = ?
                """,
                (operation_id,),
            )
            row = cursor.fetchone()
            return {
                "total_attempts": row[0],
                "success_count": row[1],
                "success_rate": row[1] / row[0] if row[0] > 0 else 0.0,
                "first_seen": row[2],
                "last_seen": row[3],
                "avg_response_time": self.get_operation_avg_response_time(operation_id),
            }