"""
角色卡格式转换器
支持 SillyTavern V2/V3 格式与内部格式的双向转换
"""
import json
import base64
from typing import Optional, Dict, Any
from pathlib import Path
from PIL import Image
import io

try:
    from backend.models.internal import CharacterCard
except ImportError:
    from models.internal import CharacterCard


class CharacterCardConverter:
    """角色卡格式转换器"""
    
    @staticmethod
    def st_to_internal(st_data: dict, avatar_path: Optional[str] = None) -> CharacterCard:
        """
        SillyTavern 格式 → 内部格式
        
        Args:
            st_data: SillyTavern 角色卡数据（V2/V3）
            avatar_path: 头像路径（可选）
        
        Returns:
            CharacterCard 对象
        """
        import uuid
        from datetime import datetime
        
        # 兼容两种传入方式：完整ST格式或直接data
        if 'spec' in st_data:
            data = st_data.get('data', {})
        else:
            data = st_data
        
        extensions = data.get('extensions', {})
        
        return CharacterCard(
            id=str(uuid.uuid4()),
            name=data['name'],
            description=data.get('description', ''),
            personality=data.get('personality', ''),
            scenario=data.get('scenario', ''),
            first_mes=data.get('first_mes', ''),
            mes_example=data.get('mes_example', ''),
            categories=[],  # ST没有categories
            tableHeaders=[],  # ST没有tableHeaders
            worldInfoId=extensions.get('world'),
            outputSchema=None,  # ST不支持结构化输出
            avatarPath=avatar_path,
            alternate_greetings=data.get('alternate_greetings', []),
            tags=data.get('tags', []),
            createdAt=int(datetime.now().timestamp()),
            updatedAt=int(datetime.now().timestamp()),
            lastChatAt=None,
            isFavorite=extensions.get('fav', False),
            version=1
        )
    
    @staticmethod
    def internal_to_st(character: CharacterCard) -> dict:
        """
        内部格式 → SillyTavern V3 格式
        
        Args:
            character: CharacterCard 对象
        
        Returns:
            SillyTavern V3 格式字典
        """
        return {
            "spec": "chara_card_v3",
            "spec_version": "3.0",
            "data": {
                "name": character.name,
                "description": character.description,
                "personality": character.personality,
                "scenario": character.scenario,
                "first_mes": character.first_mes,
                "mes_example": character.mes_example,
                "alternate_greetings": character.alternate_greetings or [],
                "tags": character.tags or [],
                "creator_notes": "",
                "system_prompt": "",
                "post_history_instructions": "",
                "extensions": {
                    "world": character.worldInfoId,
                    "talkativeness": 0.5,
                    "fav": character.isFavorite
                }
            }
        }
    
    @staticmethod
    def export_as_png(character: CharacterCard, avatar_path: Optional[str] = None, use_default_avatar: bool = False) -> bytes:
        """
        导出为 SillyTavern PNG 格式
        
        Args:
            character: CharacterCard 对象
            avatar_path: 头像图片路径（可选）
            use_default_avatar: 是否使用默认头像（不嵌入JSON数据）
        
        Returns:
            PNG 文件的二进制数据
        """
        # 1. 创建/加载图片
        if avatar_path and Path(avatar_path).exists():
            img = Image.open(avatar_path)
        else:
            # 创建默认图片（400x600像素，灰色背景）
            img = Image.new('RGB', (400, 600), color=(73, 109, 137))
        
        # 确保是 RGBA 模式
        if img.mode != 'RGBA':
            img = img.convert('RGBA')
        
        # 2. 如果不是默认头像，才嵌入 JSON 数据
        if not use_default_avatar:
            st_data = CharacterCardConverter.internal_to_st(character)
            json_str = json.dumps(st_data, ensure_ascii=False)
            base64_data = base64.b64encode(json_str.encode('utf-8')).decode('ascii')
            img.text['ccv3'] = base64_data
        
        # 3. 保存到字节流
        buffer = io.BytesIO()
        img.save(buffer, format='PNG')
        buffer.seek(0)
        
        return buffer.read()
    
    @staticmethod
    def extract_from_png(png_data: bytes) -> Optional[dict]:
        """
        从 PNG 文件中提取嵌入的角色数据
        
        Args:
            png_data: PNG 文件的二进制数据
        
        Returns:
            SillyTavern 格式字典，如果没有嵌入数据则返回 None
        """
        try:
            img = Image.open(io.BytesIO(png_data))
            
            # 尝试 V3 格式 (ccv3)
            if 'ccv3' in img.text:
                json_str = base64.b64decode(img.text['ccv3']).decode('utf-8')
                return json.loads(json_str)
            
            # 尝试 V2 格式 (chara)
            elif 'chara' in img.text:
                json_str = base64.b64decode(img.text['chara']).decode('utf-8')
                return json.loads(json_str)
            
            # 没有嵌入数据
            return None
        
        except Exception as e:
            print(f"解析PNG失败: {e}")
            return None