"""
上下文/存储统一封装

提供统一的会话上下文读写接口：
- 私聊/单人会话 memory（优先 Redis，降级内存）
- 群聊 history（优先 Redis，降级文件）
- 持久记忆 sqlite

AIChat 只需要通过本模块读写消息，不再关心介质细节。
"""

from __future__ import annotations

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

from loguru import logger

from utils.redis_cache import get_cache


def _safe_chat_id(chat_id: str) -> str:
    return (chat_id or "").replace("@", "_").replace(":", "_")


def _extract_text_from_multimodal(content: Any) -> str:
    if isinstance(content, str):
        return content
    if isinstance(content, list):
        parts = []
        for item in content:
            if isinstance(item, dict) and item.get("type") == "text":
                parts.append(item.get("text", ""))
        return "".join(parts).strip()
    return str(content)


@dataclass
class HistoryRecord:
    role: str = "user"
    nickname: str = ""
    content: Any = ""
    timestamp: Any = None
    wxid: Optional[str] = None
    id: Optional[str] = None

    @classmethod
    def from_raw(cls, raw: Dict[str, Any]) -> "HistoryRecord":
        role = raw.get("role") or "user"
        nickname = raw.get("nickname") or raw.get("SenderNickname") or ""
        content = raw.get("content") if "content" in raw else raw.get("Content", "")
        ts = raw.get("timestamp") or raw.get("time") or raw.get("CreateTime")
        wxid = raw.get("wxid") or raw.get("SenderWxid")
        rid = raw.get("id") or raw.get("msgid")
        return cls(role=role, nickname=nickname, content=content, timestamp=ts, wxid=wxid, id=rid)

    def to_dict(self) -> Dict[str, Any]:
        d = {
            "role": self.role or "user",
            "nickname": self.nickname,
            "content": self.content,
        }
        if self.timestamp is not None:
            d["timestamp"] = self.timestamp
        if self.wxid:
            d["wxid"] = self.wxid
        if self.id:
            d["id"] = self.id
        return d


class ContextStore:
    """
    统一上下文存储。

    Args:
        config: AIChat 配置 dict
        history_dir: 历史文件目录（群聊降级）
        memory_fallback: AIChat 内存 dict（私聊降级）
        history_locks: AIChat locks dict（文件写入）
        persistent_db_path: sqlite 文件路径
    """

    def __init__(
        self,
        config: Dict[str, Any],
        history_dir: Optional[Path],
        memory_fallback: Dict[str, List[Dict[str, Any]]],
        history_locks: Dict[str, asyncio.Lock],
        persistent_db_path: Optional[Path],
    ):
        self.config = config or {}
        self.history_dir = history_dir
        self.memory_fallback = memory_fallback
        self.history_locks = history_locks
        self.persistent_db_path = persistent_db_path

    # ------------------ 私聊 memory ------------------

    def _use_redis_for_memory(self) -> bool:
        redis_config = self.config.get("redis", {})
        if not redis_config.get("use_redis_history", True):
            return False
        redis_cache = get_cache()
        return bool(redis_cache and redis_cache.enabled)

    def add_private_message(
        self,
        chat_id: str,
        role: str,
        content: Any,
        *,
        image_base64: str = None,
        nickname: str = "",
        sender_wxid: str = None,
    ) -> None:
        if not self.config.get("memory", {}).get("enabled", False):
            return

        if image_base64:
            message_content = [
                {"type": "text", "text": _extract_text_from_multimodal(content)},
                {"type": "image_url", "image_url": {"url": image_base64}},
            ]
        else:
            message_content = content

        redis_config = self.config.get("redis", {})
        if self._use_redis_for_memory():
            redis_cache = get_cache()
            ttl = redis_config.get("chat_history_ttl", 86400)
            try:
                redis_cache.add_chat_message(
                    chat_id,
                    role,
                    message_content,
                    nickname=nickname,
                    sender_wxid=sender_wxid,
                    ttl=ttl,
                )
                max_messages = self.config.get("memory", {}).get("max_messages", 20)
                redis_cache.trim_chat_history(chat_id, max_messages)
                return
            except Exception as e:
                logger.debug(f"[ContextStore] Redis private history 写入失败: {e}")

        if chat_id not in self.memory_fallback:
            self.memory_fallback[chat_id] = []
        self.memory_fallback[chat_id].append({"role": role, "content": message_content})
        max_messages = self.config.get("memory", {}).get("max_messages", 20)
        if len(self.memory_fallback[chat_id]) > max_messages:
            self.memory_fallback[chat_id] = self.memory_fallback[chat_id][-max_messages:]

    def get_private_messages(self, chat_id: str) -> List[Dict[str, Any]]:
        if not self.config.get("memory", {}).get("enabled", False):
            return []

        if self._use_redis_for_memory():
            redis_cache = get_cache()
            max_messages = self.config.get("memory", {}).get("max_messages", 20)
            try:
                history = redis_cache.get_chat_history(chat_id, max_messages)
                return [HistoryRecord.from_raw(h).to_dict() for h in history]
            except Exception as e:
                logger.debug(f"[ContextStore] Redis private history 读取失败: {e}")

        return self.memory_fallback.get(chat_id, [])

    def clear_private_messages(self, chat_id: str) -> None:
        if self._use_redis_for_memory():
            redis_cache = get_cache()
            try:
                redis_cache.clear_chat_history(chat_id)
            except Exception:
                pass
        self.memory_fallback.pop(chat_id, None)

    # ------------------ 群聊 history ------------------

    def _use_redis_for_group_history(self) -> bool:
        redis_config = self.config.get("redis", {})
        if not redis_config.get("use_redis_history", True):
            return False
        redis_cache = get_cache()
        return bool(redis_cache and redis_cache.enabled)

    def _get_history_file(self, chat_id: str) -> Optional[Path]:
        if not self.history_dir:
            return None
        return self.history_dir / f"{_safe_chat_id(chat_id)}.json"

    def _get_history_lock(self, chat_id: str) -> asyncio.Lock:
        lock = self.history_locks.get(chat_id)
        if lock is None:
            lock = asyncio.Lock()
            self.history_locks[chat_id] = lock
        return lock

    def _read_history_file(self, history_file: Path) -> List[Dict[str, Any]]:
        try:
            with open(history_file, "r", encoding="utf-8") as f:
                return json.load(f)
        except FileNotFoundError:
            return []
        except Exception as e:
            logger.error(f"读取历史记录失败: {history_file}, {e}")
            return []

    def _write_history_file(self, history_file: Path, history: List[Dict[str, Any]]) -> None:
        history_file.parent.mkdir(parents=True, exist_ok=True)
        temp_file = Path(str(history_file) + ".tmp")
        with open(temp_file, "w", encoding="utf-8") as f:
            json.dump(history, f, ensure_ascii=False, indent=2)
        temp_file.replace(history_file)

    async def load_group_history(self, chat_id: str) -> List[Dict[str, Any]]:
        if not self.config.get("history", {}).get("enabled", True):
            return []

        if self._use_redis_for_group_history():
            redis_cache = get_cache()
            max_history = self.config.get("history", {}).get("max_history", 100)
            try:
                history = redis_cache.get_group_history(chat_id, max_history)
                return [HistoryRecord.from_raw(h).to_dict() for h in history]
            except Exception as e:
                logger.debug(f"[ContextStore] Redis group history 读取失败: {e}")

        history_file = self._get_history_file(chat_id)
        if not history_file:
            return []

        lock = self._get_history_lock(chat_id)
        async with lock:
            raw_history = self._read_history_file(history_file)
        return [HistoryRecord.from_raw(h).to_dict() for h in raw_history]

    async def add_group_message(
        self,
        chat_id: str,
        nickname: str,
        content: Any,
        *,
        record_id: str = None,
        image_base64: str = None,
        role: str = "user",
        sender_wxid: str = None,
    ) -> None:
        if not self.config.get("history", {}).get("enabled", True):
            return

        if image_base64:
            message_content = [
                {"type": "text", "text": _extract_text_from_multimodal(content)},
                {"type": "image_url", "image_url": {"url": image_base64}},
            ]
        else:
            message_content = content

        if self._use_redis_for_group_history():
            redis_cache = get_cache()
            redis_config = self.config.get("redis", {})
            ttl = redis_config.get("group_history_ttl", 172800)
            try:
                redis_cache.add_group_message(
                    chat_id,
                    nickname,
                    message_content,
                    record_id=record_id,
                    role=role,
                    sender_wxid=sender_wxid,
                    ttl=ttl,
                )
                max_history = self.config.get("history", {}).get("max_history", 100)
                redis_cache.trim_group_history(chat_id, max_history)
                return
            except Exception as e:
                logger.debug(f"[ContextStore] Redis group history 写入失败: {e}")

        history_file = self._get_history_file(chat_id)
        if not history_file:
            return

        lock = self._get_history_lock(chat_id)
        async with lock:
            history = self._read_history_file(history_file)
            record = HistoryRecord(
                role=role or "user",
                nickname=nickname,
                content=message_content,
                timestamp=datetime.now().isoformat(),
                wxid=sender_wxid,
                id=record_id,
            )
            history.append(record.to_dict())
            max_history = self.config.get("history", {}).get("max_history", 100)
            if len(history) > max_history:
                history = history[-max_history:]
            self._write_history_file(history_file, history)

    async def update_group_message_by_id(self, chat_id: str, record_id: str, new_content: Any) -> None:
        if not self.config.get("history", {}).get("enabled", True):
            return

        if self._use_redis_for_group_history():
            redis_cache = get_cache()
            try:
                redis_cache.update_group_message_by_id(chat_id, record_id, new_content)
                return
            except Exception as e:
                logger.debug(f"[ContextStore] Redis group history 更新失败: {e}")

        history_file = self._get_history_file(chat_id)
        if not history_file:
            return

        lock = self._get_history_lock(chat_id)
        async with lock:
            history = self._read_history_file(history_file)
            for rec in history:
                if rec.get("id") == record_id:
                    rec["content"] = new_content
                    break
            max_history = self.config.get("history", {}).get("max_history", 100)
            if len(history) > max_history:
                history = history[-max_history:]
            self._write_history_file(history_file, history)

    # ------------------ 持久记忆 sqlite ------------------

    def init_persistent_memory_db(self) -> Optional[Path]:
        if not self.persistent_db_path:
            return None

        self.persistent_db_path.parent.mkdir(exist_ok=True, parents=True)
        conn = sqlite3.connect(self.persistent_db_path)
        cursor = conn.cursor()
        cursor.execute(
            """
            CREATE TABLE IF NOT EXISTS memories (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                chat_id TEXT NOT NULL,
                chat_type TEXT NOT NULL,
                user_wxid TEXT NOT NULL,
                user_nickname TEXT,
                content TEXT NOT NULL,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
            )
            """
        )
        cursor.execute("CREATE INDEX IF NOT EXISTS idx_chat_id ON memories(chat_id)")
        conn.commit()
        conn.close()
        logger.info(f"持久记忆数据库已初始化: {self.persistent_db_path}")
        return self.persistent_db_path

    def add_persistent_memory(
        self,
        chat_id: str,
        chat_type: str,
        user_wxid: str,
        user_nickname: str,
        content: str,
    ) -> int:
        if not self.persistent_db_path:
            return -1

        conn = sqlite3.connect(self.persistent_db_path)
        cursor = conn.cursor()
        cursor.execute(
            """
            INSERT INTO memories (chat_id, chat_type, user_wxid, user_nickname, content)
            VALUES (?, ?, ?, ?, ?)
            """,
            (chat_id, chat_type, user_wxid, user_nickname, content),
        )
        memory_id = cursor.lastrowid
        conn.commit()
        conn.close()
        return memory_id

    def get_persistent_memories(self, chat_id: str) -> List[Dict[str, Any]]:
        if not self.persistent_db_path:
            return []

        conn = sqlite3.connect(self.persistent_db_path)
        cursor = conn.cursor()
        cursor.execute(
            """
            SELECT id, user_nickname, content, created_at
            FROM memories
            WHERE chat_id = ?
            ORDER BY created_at ASC
            """,
            (chat_id,),
        )
        rows = cursor.fetchall()
        conn.close()
        return [
            {"id": r[0], "nickname": r[1], "content": r[2], "time": r[3]}
            for r in rows
        ]

    def delete_persistent_memory(self, chat_id: str, memory_id: int) -> bool:
        if not self.persistent_db_path:
            return False

        conn = sqlite3.connect(self.persistent_db_path)
        cursor = conn.cursor()
        cursor.execute(
            "DELETE FROM memories WHERE id = ? AND chat_id = ?",
            (memory_id, chat_id),
        )
        deleted = cursor.rowcount > 0
        conn.commit()
        conn.close()
        return deleted

    def clear_persistent_memories(self, chat_id: str) -> int:
        if not self.persistent_db_path:
            return 0

        conn = sqlite3.connect(self.persistent_db_path)
        cursor = conn.cursor()
        cursor.execute("DELETE FROM memories WHERE chat_id = ?", (chat_id,))
        deleted_count = cursor.rowcount
        conn.commit()
        conn.close()
        return deleted_count

    # ------------------ 旧数据扫描/清理 ------------------

    def find_legacy_group_history_keys(self) -> List[str]:
        """
        发现旧版本使用 safe_id 写入的 group_history key。

        Returns:
            legacy_keys 列表（不删除）
        """
        redis_cache = get_cache()
        if not redis_cache or not redis_cache.enabled:
            return []

        try:
            keys = redis_cache.client.keys("group_history:*")
            legacy = []
            for k in keys or []:
                # 新 key 一般包含 @chatroom；旧 safe_id 不包含 @
                if "@chatroom" not in k and "_" in k:
                    legacy.append(k)
            return legacy
        except Exception as e:
            logger.debug(f"[ContextStore] 扫描 legacy group_history keys 失败: {e}")
            return []

    def delete_legacy_group_history_keys(self, legacy_keys: List[str]) -> int:
        """删除给定 legacy key 列表"""
        redis_cache = get_cache()
        if not redis_cache or not redis_cache.enabled or not legacy_keys:
            return 0
        try:
            return redis_cache.client.delete(*legacy_keys)
        except Exception as e:
            logger.debug(f"[ContextStore] 删除 legacy group_history keys 失败: {e}")
            return 0