"""媒体管理器。

负责图片和表情包的识别、存储和管理。

功能：
- 使用 VLM 识别图片和表情包内容
- 缓存识别结果到数据库，避免重复识别
- 管理媒体文件的存储和检索
- 支持按哈希值去重，节省存储和计算资源

设计原则：
- 优先从缓存读取，减少 VLM 调用
- 使用哈希值标识图片，避免重复处理
- 异步处理，不阻塞主流程
- 异常友好，识别失败不影响消息流转
"""

from __future__ import annotations

import asyncio
import hashlib
import time
from pathlib import Path
from collections.abc import Iterable
from typing import Any
from sqlalchemy import select

from src.kernel.logger import get_logger
from src.app.plugin_system.api.llm_api import get_model_set_by_task
from src.kernel.llm.model_client.registry import ModelClientRegistry
from src.core.prompt import PromptTemplate, get_prompt_manager
from src.core.config import get_core_config
from src.core.utils.base64_helper import base64_decode_to_bytes
from src.kernel.scheduler import get_unified_scheduler, TriggerType
from src.kernel.db import QueryBuilder, invalidate_model_cache
from src.kernel.db.core.session import get_db_session
from src.core.models.sql_alchemy import Images, ImageDescriptions
from src.kernel.llm import LLMContextManager, LLMPayload, ROLE, Text, Image

logger = get_logger("media_manager")

# 单例实例
_media_manager: "MediaManager | None" = None


class MediaManager:
    """媒体管理器。
    
    管理图片、表情包等媒体资源的识别、存储和检索。
    
    主要功能：
    1. VLM 识别：调用 VLM 模型识别图片/表情包内容
    2. 缓存管理：使用哈希值缓存识别结果
    3. 数据库存储：持久化媒体信息
    4. 去重优化：相同内容的图片只识别一次
    
    Examples:
        >>> manager = get_media_manager()
        >>> description = await manager.recognize_image(base64_data, "image")
        >>> await manager.save_media_info(...)
    """

    def __init__(self):
        """初始化媒体管理器。"""
        self._vlm_model_set = None
        # stream_id -> 跳过的媒体类型集合；值为 None 表示跳过所有类型
        self._skip_vlm_streams: dict[str, frozenset[str] | None] = {}
        self._initialize_vlm()
        self._initialize_asr()
        self._register_prompts()
        self._setup_media_folders()
        self._cleanup_task_id = None
        self._start_cleanup_scheduler()

    def _initialize_vlm(self) -> None:
        """初始化 VLM 模型配置。"""
        try:
            self._vlm_model_set = get_model_set_by_task("vlm")
            self._vlm_available = self._vlm_model_set is not None
            
            if self._vlm_available:
                logger.info("VLM 模型已加载，媒体识别功能可用")
            else:
                logger.info("未配置 VLM 模型，媒体识别功能不可用")
        except Exception as e:
            logger.error(f"初始化 VLM 模型失败: {e}")

    def _initialize_asr(self) -> None:
        """初始化 ASR 模型配置。"""
        try:
            self._asr_model_set = get_model_set_by_task("voice")
            self._asr_available = self._asr_model_set is not None

            if self._asr_available:
                logger.info("ASR 模型已加载，语音识别功能可用")
            else:
                logger.info("未配置 ASR 模型，语音识别功能不可用")
        except Exception as e:
            self._asr_model_set = None
            self._asr_available = False
            logger.error(f"初始化 ASR 模型失败: {e}")

    def _register_prompts(self) -> None:
        """注册媒体识别相关的提示词模板。"""
        try:
            manager = get_prompt_manager()
            
            # 注册图片识别提示词
            custom_prompt = get_core_config().chat.image_recognition_prompt
            default_template = "描述这张图片的内容，包含主题、主要元素。若有文字或代码，完整转述。"
            image_prompt = PromptTemplate(
                name="media.image_recognition",
                template=custom_prompt if custom_prompt else default_template
            )
            manager.register_template(image_prompt)
            
            # 注册表情包识别提示词
            emoji_prompt = PromptTemplate(
                name="media.emoji_recognition",
                template="请简要描述这个表情包的内容和含义，用一句话概括。"
            )
            manager.register_template(emoji_prompt)
            
            logger.debug("媒体识别提示词模板已注册")
        except Exception as e:
            logger.warning(f"注册提示词模板失败: {e}")

    def _setup_media_folders(self) -> None:
        """设置媒体文件夹结构。"""
        try:
            # 媒体根目录
            self.media_root = Path("data/media_cache")
            
            # 子文件夹
            self.pending_folder = self.media_root / "pending"  # 待识别
            self.images_folder = self.media_root / "images"    # 识别完成的图片
            self.emojis_folder = self.media_root / "emojis"    # 识别完成的表情包
            
            # 创建所有必要的文件夹
            for folder in [self.pending_folder, self.images_folder, self.emojis_folder]:
                folder.mkdir(parents=True, exist_ok=True)
            
            logger.info(f"媒体文件夹已初始化: {self.media_root}")
        except Exception as e:
            logger.error(f"创建媒体文件夹失败: {e}")

    def _start_cleanup_scheduler(self) -> None:
        """启动定时清理任务（每5分钟清理一次缓存）。"""
        try:
            # 延迟导入，避免循环依赖
            # 确保在异步上下文中创建任务
            asyncio.create_task(self._register_cleanup_task())
            
            logger.info("媒体缓存清理调度器已启动(每5分钟)")
        except Exception as e:
            logger.error(f"启动清理调度器失败: {e}")

    async def _register_cleanup_task(self) -> None:
        """注册定时清理任务到调度器。"""
        try:
            scheduler = get_unified_scheduler()
            
            # 创建周期性清理任务（每5分钟 = 300秒）
            schedule_id = await scheduler.create_schedule(
                callback=self._cleanup_pending_folder,
                trigger_type=TriggerType.TIME,
                trigger_config={"delay_seconds": 300},  # 5分钟
                is_recurring=True,
                task_name="media_cache_cleanup"
            )
            
            self._cleanup_task_id = schedule_id
            logger.info(f"媒体缓存清理任务已注册: {schedule_id}")
        except Exception as e:
            logger.error(f"注册清理任务失败: {e}")

    async def _cleanup_pending_folder(self) -> None:
        """清理待识别文件夹中的陈旧文件。"""
        try:
            if not self.pending_folder.exists():
                return
            
            current_time = time.time()
            cleanup_count = 0
            
            # 遍历所有待识别文件
            for file_path in self.pending_folder.iterdir():
                if not file_path.is_file():
                    continue
                
                # 获取文件修改时间
                file_mtime = file_path.stat().st_mtime
                
                # 如果文件超过5分钟未处理，删除它
                if current_time - file_mtime >= 300:  # 5分钟 = 300秒
                    try:
                        file_path.unlink()
                        cleanup_count += 1
                    except Exception as e:
                        logger.warning(f"删除文件失败 {file_path.name}: {e}")
            
            if cleanup_count > 0:
                logger.info(f"媒体缓存清理完成，删除了 {cleanup_count} 个陈旧文件")
        except Exception as e:
            logger.error(f"清理待识别文件夹失败: {e}")

    # ──────────────────────────────────────────
    # 公共 API：VLM 识别控制
    # ──────────────────────────────────────────

    def skip_vlm_for_stream(
        self,
        stream_id: str,
        media_types: Iterable[str] | None = None,
    ) -> None:
        """注册指定聊天流跳过 VLM 识别。

        调用后，该 stream_id 的消息在 MessageConverter 中将不再触发
        VLM 图片/表情包识别，媒体数据仅保留原始 base64。
        适用于聊天流程自行处理多模态内容的场景。

        Args:
            stream_id: 要跳过 VLM 识别的聊天流 ID
            media_types: 要跳过的媒体类型集合（如 ``("image",)``）。为 ``None``
                时表示跳过所有类型，保持与旧调用兼容。
        """
        if media_types is None:
            self._skip_vlm_streams[stream_id] = None
            logger.debug(f"已注册跳过 VLM 识别: stream_id={stream_id[:8]} (全部类型)")
            return
        types = frozenset(media_types)
        self._skip_vlm_streams[stream_id] = types
        logger.debug(
            f"已注册跳过 VLM 识别: stream_id={stream_id[:8]} (类型={sorted(types)})"
        )

    def unskip_vlm_for_stream(self, stream_id: str) -> None:
        """取消指定聊天流的 VLM 识别跳过。

        Args:
            stream_id: 要恢复 VLM 识别的聊天流 ID
        """
        self._skip_vlm_streams.pop(stream_id, None)
        logger.debug(f"已取消跳过 VLM 识别: stream_id={stream_id[:8]}")

    def should_skip_vlm(
        self,
        stream_id: str,
        media_type: str | None = None,
    ) -> bool:
        """查询指定聊天流是否应跳过 VLM 识别。

        Args:
            stream_id: 聊天流 ID
            media_type: 待识别媒体的类型；省略时表示
                ""该流是否对任意类型注册了跳过""，用于保留旧的整流粒度语义。

        Returns:
            True 表示该聊天流（针对给定媒体类型）应跳过 VLM 识别
        """
        if stream_id not in self._skip_vlm_streams:
            return False
        types = self._skip_vlm_streams[stream_id]
        if types is None:
            return True
        if media_type is None:
            return True
        return media_type in types

    # ──────────────────────────────────────────
    # 公共 API：媒体识别
    # ──────────────────────────────────────────

    async def recognize_media(
        self, 
        base64_data: str, 
        media_type: str,
        use_cache: bool = True
    ) -> str | None:
        """识别媒体内容（图片或表情包）。
        
        Args:
            base64_data: base64 编码的媒体数据
            media_type: 媒体类型，"image" 或 "emoji"
            use_cache: 是否使用缓存（默认 True）
            
        Returns:
            媒体的文字描述，识别失败返回 None
        """
        try:
            # 计算哈希值
            media_hash = self._compute_hash(base64_data)
            
            # 尝试从缓存读取
            if use_cache:
                cached_description = await self._get_cached_description(
                    media_hash, 
                    media_type
                )
                if cached_description:
                    logger.debug(f"从缓存获取{media_type}描述: {media_hash[:8]}...")
                    return cached_description

            if not self._vlm_model_set:
                logger.debug(f"VLM 模型不可用，跳过{media_type}识别")
                return None
            
            # 保存到待识别文件夹
            pending_file_path = await self._save_to_pending(
                base64_data,
                media_hash,
                media_type
            )
            
            # VLM 识别
            description = await self._recognize_with_vlm(base64_data, media_type)
            
            if description:
                # 保存到缓存
                await self._save_description_cache(
                    media_hash,
                    media_type,
                    description
                )
                logger.info(f"成功识别{media_type}: {description[:50]}...")
                
                # 移动到对应的分类文件夹
                await self._move_to_category_folder(
                    pending_file_path,
                    media_type,
                    media_hash
                )
                
                # 保存媒体信息到数据库
                target_folder = self.images_folder if media_type == "image" else self.emojis_folder
                target_file_path = target_folder / pending_file_path.name
                await self.save_media_info(
                    media_hash=media_hash,
                    media_type=media_type,
                    file_path=str(target_file_path),
                    description=description,
                    vlm_processed=True
                )
            else:
                # 识别失败，保持在待识别文件夹，等待定时清理
                logger.warning(f"识别失败，文件保留在待识别文件夹: {pending_file_path.name}")
            
            return description
            
        except Exception as e:
            logger.error(f"识别{media_type}失败: {e}", exc_info=True)
            return None

    async def recognize_voice(self, audio_base64: str) -> str | None:
        """使用 ASR 识别语音内容。

        Args:
            audio_base64: base64 编码的 WAV 音频数据。

        Returns:
            识别出的文字，失败返回 None。
        """
        try:
            if not self._asr_available or not self._asr_model_set:
                logger.debug("ASR 模型不可用")
                return None
            return await self._recognize_with_asr(audio_base64)
        except Exception as e:
            logger.error(f"语音识别失败: {e}", exc_info=True)
            return None

    async def recognize_batch(
        self,
        media_list: list[tuple[str, str]],
        use_cache: bool = True
    ) -> list[tuple[int, str | None]]:
        """批量识别多个媒体。
        
        Args:
            media_list: [(base64_data, media_type), ...] 列表
            use_cache: 是否使用缓存
            
        Returns:
            [(index, description), ...] 列表，description 为 None 表示识别失败
        """
        results = []
        for idx, (base64_data, media_type) in enumerate(media_list):
            description = await self.recognize_media(
                base64_data,
                media_type,
                use_cache=use_cache
            )
            results.append((idx, description))
        return results

    # ──────────────────────────────────────────
    # 公共 API：数据库操作
    # ──────────────────────────────────────────

    async def save_media_info(
        self,
        media_hash: str,
        media_type: str,
        file_path: str | None = None,
        description: str | None = None,
        vlm_processed: bool = False
    ) -> None:
        """保存媒体信息到数据库。
        
        Args:
            media_hash: 媒体哈希值（作为唯一标识）
            media_type: 媒体类型（image/emoji）
            file_path: 文件路径（可选）
            description: 描述文本（可选）
            vlm_processed: 是否已经过 VLM 处理
        """
        try:
            changed = False
            async with get_db_session() as session:
                # 查找现有记录（使用 image_id 作为唯一标识）
                # 这里使用 scalars().first() 来避免数据库中存在多条重复记录导致的 MultipleResultsFound 错误
                stmt = (
                    select(Images)
                    .where(Images.image_id == media_hash)
                    .order_by(Images.timestamp.desc())
                    .limit(1)
                )
                result = await session.execute(stmt)
                existing = result.scalars().first()

                if existing:
                    # 更新现有记录
                    existing.count += 1
                    if description:
                        existing.description = description
                    if vlm_processed:
                        existing.vlm_processed = True
                    changed = True
                    logger.debug(f"更新媒体记录: {media_hash[:8]}... count={existing.count}")
                else:
                    # 创建新记录
                    new_image = Images(
                        image_id=media_hash,
                        path=file_path or media_hash,  # 如果没有路径，用哈希值
                        type=media_type,
                        description=description,
                        timestamp=time.time(),
                        vlm_processed=vlm_processed,
                        count=1
                    )
                    session.add(new_image)
                    changed = True
                    logger.debug(f"创建新媒体记录: {media_hash[:8]}...")

                await session.commit()
            if changed:
                invalidate_model_cache(Images)

        except Exception as e:
            logger.error(f"保存媒体信息失败: {e}", exc_info=True)

    async def get_media_info(self, media_hash: str) -> dict[str, Any] | None:
        """根据哈希值获取媒体信息。
        
        Args:
            media_hash: 媒体哈希值
            
        Returns:
            媒体信息字典，不存在返回 None
        """
        try:
            media = await (
                QueryBuilder(Images)
                .filter(image_id=media_hash)
                .order_by("-timestamp")
                .first()
            )
            if media is not None:
                return {
                    "id": media.id,
                    "image_id": media.image_id,
                    "path": media.path,
                    "type": media.type,
                    "description": media.description,
                    "count": media.count,
                    "timestamp": media.timestamp,
                    "vlm_processed": media.vlm_processed,
                }
            return None

        except Exception as e:
            logger.error(f"查询媒体信息失败: {e}", exc_info=True)
            return None

    # ──────────────────────────────────────────
    # 内部方法
    # ──────────────────────────────────────────

    async def _recognize_with_vlm(
        self, 
        base64_data: str, 
        media_type: str
    ) -> str | None:
        """使用 VLM 识别单个媒体。
        
        Args:
            base64_data: base64 编码的媒体数据
            media_type: 媒体类型（image 或 emoji）
            
        Returns:
            识别结果文本，失败返回 None
        """
        try:
            from src.app.plugin_system.api.llm_api import create_llm_request
            
            # 检查 VLM 模型是否可用
            if not self._vlm_model_set:
                logger.debug("VLM 模型不可用")
                return None

            # 创建 VLM 请求
            context_manager = LLMContextManager()
            request = create_llm_request(
                self._vlm_model_set,
                "image_recognition",
                context_manager=context_manager,
            )

            # 从提示词管理器获取提示词模板
            prompt_manager = get_prompt_manager()
            if media_type == "emoji":
                template = prompt_manager.get_template("media.emoji_recognition")
            else:
                template = prompt_manager.get_template("media.image_recognition")
            
            # 构建提示词（模板不需要参数，直接build）
            if template:
                prompt = await template.build()

            # 处理 base64 数据：提取纯净的 base64 内容
            clean_base64 = self._extract_clean_base64(base64_data)
            mime_type = self._extract_image_mime_type(base64_data)
            
            # 使用标准的 data URL 格式（大多数 VLM API 都支持）
            image_value = f"data:{mime_type};base64,{clean_base64}"

            # 添加 payload 并发送请求
            request.add_payload(LLMPayload(ROLE.USER, [Text(prompt), Image(image_value)]))
            response = await request.send(stream=False)
            await response

            # 提取并处理描述
            description = response.message.strip() if response.message else ""
            
            # 限制长度
            if len(description) > 100:
                description = description[:97] + "..."

            return description if description else None

        except Exception as e:
            logger.error(f"VLM 识别失败: {e}", exc_info=True)
            return None

    async def _recognize_with_asr(self, audio_base64: str) -> str | None:
        """调用 ASR 客户端执行语音转文字。

        Args:
            audio_base64: base64 编码的 WAV 音频数据。

        Returns:
            识别出的文字，失败返回 None。
        """
        try:
            registry = ModelClientRegistry()
            model_set = self._asr_model_set
            # model_set 是 list[dict]，每个元素即一个 ModelEntry
            if not isinstance(model_set, list) or not model_set:
                logger.debug("ASR model_set 中无可用模型")
                return None

            model_entry = model_set[0]
            client = registry.get_asr_client_for_model(model_entry)
            model_name = model_entry.get("model_identifier") if isinstance(model_entry, dict) else str(model_entry)

            clean_b64 = self._extract_clean_base64(audio_base64)
            audio_bytes = await asyncio.to_thread(
                base64_decode_to_bytes,
                clean_b64,
            )

            text = await client.create_transcription(
                model_name=model_name,
                audio_bytes=audio_bytes,
                request_name="voice_recognition",
                model_set=model_entry,
            )
            return text.strip() if text else None
        except Exception as e:
            logger.error(f"ASR 请求失败: {e}", exc_info=True)
            return None

    async def _get_cached_description(
        self,
        media_hash: str,
        media_type: str
    ) -> str | None:
        """从数据库缓存获取描述。
        
        Args:
            media_hash: 媒体哈希值
            media_type: 媒体类型
            
        Returns:
            缓存的描述，不存在返回 None
        """
        try:
            desc = await (
                QueryBuilder(ImageDescriptions)
                .filter(image_description_hash=media_hash, type=media_type)
                .order_by("-timestamp")
                .first()
            )
            return desc.description if desc else None

        except Exception as e:
            logger.debug(f"查询缓存失败: {e}")
            return None

    async def _save_description_cache(
        self,
        media_hash: str,
        media_type: str,
        description: str
    ) -> None:
        """保存描述到缓存。
        
        Args:
            media_hash: 媒体哈希值
            media_type: 媒体类型
            description: 描述文本
        """
        try:
            created = False
            async with get_db_session() as session:
                # 检查是否已存在（避免重复记录导致 MultipleResultsFound）
                stmt = (
                    select(ImageDescriptions)
                    .where(
                        ImageDescriptions.image_description_hash == media_hash,
                        ImageDescriptions.type == media_type
                    )
                    .order_by(ImageDescriptions.timestamp.desc())
                    .limit(1)
                )
                result = await session.execute(stmt)
                # 使用 scalars().first() 避免 MultipleResultsFound 错误
                existing = result.scalars().first()

                if not existing:
                    # 创建新缓存记录
                    new_desc = ImageDescriptions(
                        image_description_hash=media_hash,
                        type=media_type,
                        description=description,
                        timestamp=time.time()
                    )
                    session.add(new_desc)
                    created = True
                    await session.commit()
                    logger.debug(f"保存描述缓存: {media_hash[:8]}...")
            if created:
                invalidate_model_cache(ImageDescriptions)

        except Exception as e:
            logger.error(f"保存描述缓存失败: {e}", exc_info=True)

    @staticmethod
    def _extract_clean_base64(data: str) -> str:
        """提取纯净的 base64 数据（移除前缀和多余字符）。
        
        Args:
            data: 可能包含前缀的 base64 字符串
            
        Returns:
            纯净的 base64 字符串
        """
        # 移除可能的 data URL 前缀
        if data.startswith("data:"):
            # 提取 base64 部分
            if "base64," in data:
                data = data.split("base64,", 1)[1]
        elif data.startswith("base64|"):
            data = data[7:]
        
        # 移除可能的换行符和空格
        data = data.replace("\n", "").replace("\r", "").replace(" ", "")
        
        return data

    @staticmethod
    def _extract_image_mime_type(data: str) -> str:
        """从 data URL 中提取图片 MIME 类型。"""
        if data.startswith("data:") and ";base64," in data:
            mime_type = data.split(";", 1)[0][len("data:"):].strip().lower()
            if mime_type.startswith("image/"):
                return mime_type
        return "image/png"
    
    async def _save_to_pending(
        self,
        base64_data: str,
        media_hash: str,
        media_type: str
    ) -> Path:
        """保存媒体文件到待识别文件夹。
        
        Args:
            base64_data: base64 编码的媒体数据
            media_hash: 媒体哈希值
            media_type: 媒体类型
            
        Returns:
            保存的文件路径
        """
        try:
            # 提取纯净的 base64 数据
            clean_base64 = self._extract_clean_base64(base64_data)
            
            # 解码为二进制数据
            binary_data = await asyncio.to_thread(
                base64_decode_to_bytes,
                clean_base64,
            )
            
            # 根据类型确定文件扩展名
            ext = ".jpg" if media_type == "image" else ".png"
            
            # 生成文件名（哈希值前16位 + 类型标记 + 扩展名）
            filename = f"{media_hash[:16]}_{media_type}{ext}"
            file_path = self.pending_folder / filename
            
            # 写入文件
            await asyncio.to_thread(file_path.write_bytes, binary_data)
            logger.debug(f"媒体已保存到待识别文件夹: {filename}")
            
            return file_path
        except Exception as e:
            logger.error(f"保存到待识别文件夹失败: {e}")
            # 返回一个虚拟路径，不影响后续流程
            return self.pending_folder / f"{media_hash[:16]}_error.tmp"

    async def _move_to_category_folder(
        self,
        source_path: Path,
        media_type: str,
        media_hash: str
    ) -> None:
        """将识别完成的文件移动到对应的分类文件夹。
        
        Args:
            source_path: 源文件路径（待识别文件夹中的文件）
            media_type: 媒体类型
            media_hash: 媒体哈希值
        """
        try:
            if not await asyncio.to_thread(source_path.exists):
                logger.debug(f"源文件不存在，跳过移动: {source_path.name}")
                return
            
            # 确定目标文件夹
            target_folder = self.images_folder if media_type == "image" else self.emojis_folder
            
            # 确定目标文件名
            target_path = target_folder / source_path.name
            
            # 如果目标文件已存在，删除源文件即可（去重）
            if await asyncio.to_thread(target_path.exists):
                await asyncio.to_thread(source_path.unlink)
                logger.debug(f"目标文件已存在，删除源文件: {source_path.name}")
                return
            
            # 移动文件
            await asyncio.to_thread(source_path.rename, target_path)
            logger.debug(f"文件已移动到 {media_type} 文件夹: {target_path.name}")
        except Exception as e:
            logger.error(f"移动文件失败: {e}")


    @staticmethod
    def _compute_hash(data: str) -> str:
        """计算数据的 SHA256 哈希值。
        
        Args:
            data: 待哈希的数据（base64 字符串）
            
        Returns:
            十六进制哈希字符串
        """
        # 使用提取的纯净 base64 数据计算哈希
        clean_data = MediaManager._extract_clean_base64(data)
        return hashlib.sha256(clean_data.encode()).hexdigest()


# ──────────────────────────────────────────
# 单例访问
# ──────────────────────────────────────────


def get_media_manager() -> MediaManager:
    """获取媒体管理器单例。
    
    Returns:
        MediaManager 实例
    """
    global _media_manager
    if _media_manager is None:
        _media_manager = MediaManager()
    return _media_manager


def initialize_media_manager() -> MediaManager:
    """初始化媒体管理器（用于显式初始化）。
    
    Returns:
        MediaManager 实例
    """
    global _media_manager
    _media_manager = MediaManager()
    return _media_manager
