SillyTavern_replica/backend/api/routes/charactersRoute.py

"""
角色卡 API 路由
"""
from fastapi import APIRouter, HTTPException, UploadFile, File
from fastapi.responses import FileResponse, StreamingResponse
from typing import List
from pathlib import Path
import io

try:
    from backend.services.character_service import CharacterService
except ImportError:
    from services.character_service import CharacterService

router = APIRouter(prefix="/characters", tags=["characters"])
character_service = CharacterService()


@router.get("/", response_model=List[dict])
async def list_characters():
    """
    获取所有角色卡列表
    
    Returns:
        按最后聊天时间排序的角色卡列表
    """
    characters = character_service.scan_all_characters()
    return [c.dict() for c in characters]


@router.get("/{name}", response_model=dict)
async def get_character(name: str):
    """
    获取指定角色卡
    
    Args:
        name: 角色名（URL编码）
    """
    character = character_service.get_character_by_name(name)
    if not character:
        raise HTTPException(status_code=404, detail=f"角色 '{name}' 不存在")
    
    return character.dict()


@router.post("/", response_model=dict)
async def create_character(character_data: dict):
    """
    创建新角色卡
    
    Request Body:
        {
            "name": "角色名",
            "description": "描述",
            "personality": "性格",
            "scenario": "场景",
            "first_mes": "开场白",
            "categories": ["分类1", "分类2"],
            "tags": ["tag1", "tag2"]
        }
    """
    try:
        character = character_service.create_character(character_data)
        return {
            "success": True,
            "character": character.dict()
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@router.put("/{name}", response_model=dict)
async def update_character(name: str, updates: dict):
    """
    更新角色卡
    
    Args:
        name: 角色名
        updates: 要更新的字段
    """
    try:
        character = character_service.update_character(name, updates)
        return {
            "success": True,
            "character": character.dict()
        }
    except FileNotFoundError:
        raise HTTPException(status_code=404, detail=f"角色 '{name}' 不存在")
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@router.delete("/{name}")
async def delete_character(name: str):
    """
    删除角色卡及其所有聊天记录
    """
    success = character_service.delete_character(name)
    if not success:
        raise HTTPException(status_code=404, detail=f"角色 '{name}' 不存在")
    
    return {"success": True, "message": f"角色 '{name}' 已删除"}


@router.get("/{name}/avatar")
async def get_avatar(name: str):
    """
    获取角色头像
    
    Returns:
        PNG 图片文件或 404
    """
    char_folder = character_service.characters_dir / name
    avatar_file = char_folder / "avatar.png"
    
    if not avatar_file.exists():
        # 返回默认头像
        default_avatar = Path("data/images/avatars/fallback.png")
        if default_avatar.exists():
            return FileResponse(default_avatar, media_type="image/png")
        raise HTTPException(status_code=404, detail="头像不存在")
    
    return FileResponse(avatar_file, media_type="image/png")


@router.post("/{name}/avatar")
async def upload_avatar(name: str, file: UploadFile = File(...)):
    """
    上传角色头像
    
    Args:
        name: 角色名
        file: PNG 图片文件
    """
    # 验证文件类型
    if not file.content_type.startswith('image/'):
        raise HTTPException(status_code=400, detail="只支持图片文件")
    
    # 检查角色是否存在
    character = character_service.get_character_by_name(name)
    if not character:
        raise HTTPException(status_code=404, detail=f"角色 '{name}' 不存在")
    
    # 保存图片
    image_data = await file.read()
    avatar_path = character_service.save_avatar(name, image_data)
    
    return {
        "success": True,
        "avatar_path": avatar_path
    }


@router.get("/{name}/chats")
async def list_chats(name: str):
    """
    获取角色的所有聊天列表
    
    Returns:
        聊天文件列表（包含最后一条消息预览）
    """
    char_folder = character_service.characters_dir / name
    chats_dir = char_folder / "chats"
    
    if not chats_dir.exists():
        return {"chats": []}
    
    import json
    from datetime import datetime
    
    chats = []
    for chat_file in chats_dir.glob("*.jsonl"):
        try:
            with open(chat_file, 'r', encoding='utf-8') as f:
                lines = f.readlines()
            
            if not lines:
                continue
            
            # 第一行是header
            header = json.loads(lines[0])
            
            # 计算消息数量（排除header）
            message_count = len(lines) - 1
            
            # 获取最后修改时间
            last_modified = datetime.fromtimestamp(
                chat_file.stat().st_mtime
            ).isoformat()
            
            # 获取最后一条消息预览
            last_message = ""
            if message_count > 0:
                try:
                    last_msg_data = json.loads(lines[-1])
                    last_message = last_msg_data.get("mes", "")
                except:
                    pass
            
            chats.append({
                "chat_name": chat_file.stem,
                "user_name": header.get("user_name", "User"),
                "character_name": header.get("character_name", ""),
                "last_modified": last_modified,
                "message_count": message_count,
                "last_message": last_message
            })
        except Exception as e:
            # 如果解析失败，使用基本信息
            chats.append({
                "chat_name": chat_file.stem,
                "last_modified": datetime.fromtimestamp(chat_file.stat().st_mtime).isoformat(),
                "message_count": 0,
                "last_message": ""
            })
    
    # 按修改时间排序
    chats.sort(key=lambda c: c.get('last_modified', ''), reverse=True)
    
    return {"chats": chats}


@router.post("/import")
async def import_character(file: UploadFile = File(...)):
    """
    导入角色卡（支持 PNG 或 JSON）
    
    - PNG: 自动提取嵌入数据，创建文件夹
    - JSON: 创建文件夹并保存
    """
    content = await file.read()
    filename = file.filename
    
    if filename.endswith('.png'):
        # 导入 PNG
        try:
            character = character_service.import_from_png(content, filename)
            
            return {
                "success": True,
                "character": character.dict(),
                "format": "png_embedded"
            }
        
        except ValueError as e:
            raise HTTPException(status_code=400, detail=str(e))
        except Exception as e:
            raise HTTPException(status_code=500, detail=f"导入失败: {str(e)}")
    
    elif filename.endswith('.json'):
        # 导入 JSON
        try:
            import json
            data = json.loads(content.decode('utf-8'))
            
            character = character_service.create_character(data)
            
            return {
                "success": True,
                "character": character.dict(),
                "format": "json"
            }
        
        except Exception as e:
            raise HTTPException(status_code=500, detail=f"导入失败: {str(e)}")
    
    else:
        raise HTTPException(status_code=400, detail="不支持的文件格式")


@router.post("/{name}/export/png")
async def export_character_as_png(name: str):
    """
    导出角色为 SillyTavern PNG 格式
    
    Returns:
        PNG 文件下载
    """
    try:
        png_data = character_service.export_as_png(name)
        
        return StreamingResponse(
            io.BytesIO(png_data),
            media_type="image/png",
            headers={
                "Content-Disposition": f"attachment; filename={name}.png"
            }
        )
    except FileNotFoundError:
        raise HTTPException(status_code=404, detail=f"角色 '{name}' 不存在")
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"导出失败: {str(e)}")