# -*- coding: utf-8 -*-
# filename: types.py
# @Time    : 2025-01-16
# @Author  : Claude Code
# @Email   : noreply@anthropic.com
# @Software: PyCharm
"""Pydantic type definitions for image tools"""
from datetime import datetime
from typing import Literal, Optional

from pydantic import BaseModel, ConfigDict, field_validator


# Text-to-image enhancement modes
TextToImageMode = Literal[
    "technical_architecture",
    "design_mockup",
    "promotional_material",
    "office_chart",
    "custom",
]

# Image-to-image enhancement modes
ImageToImageMode = Literal[
    "content_adjustment",
    "element_replace",
    "style_transfer",
    "text_modification",
    "custom",
]


class PromptEnhancementResult(BaseModel):
    """Prompt enhancement result"""

    original_prompt: str  # Original prompt
    enhanced_prompt: str  # Enhanced prompt
    mode: str  # Mode used
    enhancement_model: str  # Enhancement model used (glm-4.6v)
    timestamp: str  # Enhancement timestamp
    token_usage: dict  # LLM token usage statistics
    # - total_tokens: int
    # - input_tokens: int
    # - output_tokens: int
    reference_image_used: bool = (
        False  # Whether reference image was used (image-to-image)
    )
    enhancement_failed: bool = (
        False  # Whether enhancement failed (fallback to original)
    )


class TextToImageParams(BaseModel):
    """Text-to-image generation tool parameters (internal use)"""

    prompt: str
    mode: TextToImageMode = "custom"
    quality: Literal["low", "medium", "high"] = "medium"
    layout: Literal["square", "landscape", "portrait"] = "square"

    model_config = ConfigDict(extra="ignore")


class ImageToImageParams(BaseModel):
    """Image-to-image generation tool parameters (internal use)"""

    prompt: str
    reference_images: list[str]  # 支持 1-8 张参考图片
    mode: ImageToImageMode = "custom"
    quality: Literal["low", "medium", "high"] = "medium"
    layout: Literal["square", "landscape", "portrait"] = "square"

    @field_validator("reference_images")
    @classmethod
    def validate_reference_images(cls, v: list[str]) -> list[str]:
        """验证参考图片数量在 1-8 张之间"""
        if not v or len(v) == 0:
            raise ValueError("reference_images 不能为空，至少需要 1 张图片")
        if len(v) > 8:
            raise ValueError(f"reference_images 最多支持 8 张图片，当前传入了 {len(v)} 张")
        return v

    model_config = ConfigDict(extra="ignore")


class ImageGenerationResult(BaseModel):
    """Image generation result"""

    # base64_data is optional - only returned when save_to_local=False
    base64_data: str | None = (
        None  # Base64-encoded image data (data:image/png;base64,...)
    )
    # file_path is optional - only returned when save_to_local=True
    file_path: str | None = None  # Absolute path to saved image file

    width: int
    height: int
    file_size: int  # Decoded byte size
    format: str
    generated_at: str
    model: str
    quality: str
    layout: str

    # Enhancement tracking
    original_prompt: str | None = None  # Original user prompt
    enhanced_prompt: str | None = None  # Enhanced prompt (if enhancement was used)


class ImageUnderstandingParams(BaseModel):
    """Image understanding tool parameters"""

    images: list[str]
    prompt_template: Literal[
        "describe",
        "ocr",
        "detect_objects",
        "analyze_chart",
        "describe_pdf",
        "describe_ppt",
        "custom",
    ] = "describe"
    custom_prompt: str | None = None

    model_config = ConfigDict(extra="ignore")


class ImageUnderstandingResult(BaseModel):
    """Image understanding result"""

    description: str
    model: str
    processed_at: str
    image_count: int
    prompt_template: str
    used_prompt: str  # The actual prompt used for analysis
    total_tokens: int
    input_tokens: int
    output_tokens: int
