from dataclasses import dataclass, field
from PIL import Image


# -------------------- #
# -- System Configs -- #
# -------------------- #
@dataclass
class AgentConfig:

    # -- Projects
    output_dir: str = "agent_output"

    # -- Loggings
    log_dir: str = "agent_log"

    # -- Inference
    seed: int = 42
    width: int = 1024
    height: int = 1024
    max_new_tokens: int = 1024
    max_refine_times: int = 1
    num_inference_steps: int = 25

    # -- LoRAs
    lora_box: str = "lora_box.json"


# ---------------------------- #
# -- Default System Prompts -- #
# ---------------------------- #
SYS_TASK_ROUTER = r"""Your task is to split the user prompt into ordered sub-tasks based on the Picture references mentioned (e.g., “Picture 1”, “Picture 2”).
Each sub-task must preserve the sequence implied in the user prompt.

For every sub-task, output an object in the form:
{"ref_id": <picture_number>, "ref_prompt": <the instruction related to that picture>}

Output a JSON list of these objects in the same order as the instructions appear in the user prompt.

Example:
User prompt: “Using the same style of Picture 1 firstly, then apply the colors and textures of Picture 2. Finally add some local design patterns from Picture 1.”
Output:

[
  {"ref_id": 1, "ref_prompt": "The same style."},
  {"ref_id": 2, "ref_prompt": "Colors and textures."},
  {"ref_id": 1, "ref_prompt": "Local design patterns."}
]"""

SYS_STYLE_ROUTER = r"""You are an image style analysis expert.
Your job is to analyze a given style image from two perspectives: style_type and style_value.

1. style_type

"semantic": the artistic style category of the image (e.g., Van Gogh, LEGO, etc.).

"pixel": the color/texture distribution of the image, focusing on colors rather than high-level artistic semantics.

2. style_value

If style_type = "semantic", then style_value can be one of (if no suitable style_value, assign to null):
{style_value_and_descriptions}

If style_type = "pixel", then style_value must be a float in [0, 1], inferred from the user prompt:

0-0.3 → slight, subtle, minimal…

0.3-0.7 → moderate, balanced, typical…

0.7-1.0 → drastic, heavy, intense…

3. Output format

Always output valid JSON, for example:

{{"style_type": "semantic", "style_value": "Van_Gogh"}}

{{"style_type": "pixel", "style_value": 0.7}}

You will be given Picture 1 and a user prompt."""

SYS_ANALYSIS_SEMANTIC = r"""You are an expert image-style analyst.
Describe the style characteristics of the given style image (Picture 1) using clear, accurate, and concrete stylistic terms.

Requirements

Output only a JSON object.

Escape inner double quotes with backslashes (\\") to ensure valid JSON.

Do not mention any content, subjects, people, objects, or specific logos. Focus solely on stylistic attributes such as color use, texture, line quality, rendering methods, mood, and artistic conventions.

The entire output must be under 250 tokens.

Output Format
{
    "description": "..."
}

Examples

Example 1:

{
    "description": "This style features bold outlines, exaggerated expressions, vibrant flat colors, and dynamic compositional energy rooted in classic American cartoon aesthetics."
}


Example 2:

{
    "description": "This style features hand-drawn painterly rendering, soft watercolor textures, lush atmospheric backgrounds, and a warm whimsical tone typical of Ghibli-inspired animation."
}"""

SYS_ANALYSIS_PIXEL = r"""You are an expert in pixel-level image style analysis.
Describe the style characteristics of the given style image (Picture 2), focusing only on low-level visual attributes.

Requirements

Emphasize color distribution, texture patterns, brightness, contrast, saturation, sharpness, gradients, and other pixel-level properties.

Do not describe any subjects, objects, scenes, or logos.

Output only a JSON object.

Escape inner double quotes with backslashes (\\").

Keep the entire output under 250 tokens.

Output Format
{
    "description": "Colors: ... ,\nTextures: ... ,\nBrightness: ... "
}"""

SYS_TRANSFER = r"""Describe the key features of the style features of Picture 2 (if given) (color, shape, size, texture, objects, background), then explain how the user's text instruction should alter or modify the Picture 1. Generate a new image that meets the user's requirements while maintaining consistency with the original input where appropriate."""

SYS_CRITERIA_CS = r"""You are an expert in evaluating content-level consistency between images.
Given a reference image (Picture 1) and a stylized target image (Picture 2), your tasks are:

1. **Ignore all style, artistic rendering, colors, and textures. Focus only on semantic content, object presence, and spatial layout.**

2. Detect and compare the objects in both images.

3. Evaluate whether the object lists and their relative positions are consistent.

4. Provide a consistency score in the range [0, 10].

5. Provide clear instructions on how Picture 2 can be modified to better match Picture 1 at the semantic and layout level. If no instructions, output empty string.

Output a valid JSON object:

{"score": <float>, "suggestion": "<instruction>"}"""

SYS_CRITERIA_RS = r"""You are an expert in evaluating style consistency between images.
Given a reference style image (Picture 1) and a target image (Picture 2), your tasks are:

Assess how closely the style of Picture 2 matches Picture 1 and provide a score in the range [0, 10].

Evaluate whether the stylization strength of Picture 2 is appropriate.

If Picture 2 needs stronger stylization to match Picture 1, output "increase".

If it is overly stylized, output "decrease".

If the stylization strength is appropriate, output "unchange".

Output a valid JSON object, for example:

{"score": 9.0, "suggestion": "unchange"}
{"score": 4.5, "suggestion": "increase"}
{"score": 3.0, "suggestion": "decrease"}"""

SYS_CRITERIA_DS = r"""You are an expert in evaluating instruction–image consistency.
Given an image editing instruction and an edited image (Picture 1), determine how well the image fulfills the instruction.

Your tasks are:

1. Provide a consistency score in the range [0, 10], where higher means Picture 1 matches the instruction more closely.

2. Provide a revised instruction in imperative form that would make Picture 1 better align with the intended edit.

Output a valid JSON object:

{"score": <float>, "suggestion": "<imperative instruction>"}


Example:

{"score": 8.0, "suggestion": "Make the tree taller than the man. Use light blue colors for the leaves."}"""


# ----------------------------------- #
# -- Inputs and Outputs structures -- #
# ----------------------------------- #
@dataclass
class UserInput:
    """
    Used for Agent conversation user interfaces
    """

    prompt: str = ""
    cnt_image_path: str = ""
    ref_image_paths: str | list[str] = ""
    ref_dict: dict[str, str] = field(default_factory=lambda: {})

    def __post_init__(self):
        ref_dict = {}
        if isinstance(self.ref_image_paths, str):
            self.ref_image_paths = [self.ref_image_paths]
        for ref_id, ref_image_path in enumerate(self.ref_image_paths):
            ref_dict[f"Picture {ref_id + 1}"] = ref_image_path
        self.ref_dict = ref_dict


@dataclass
class SubTask:
    """
    Each sub task responds to a style reference image
    """

    ref_id: int = 1
    ref_image_path: str = ""
    style_type: str = ""
    style_value: str | float = ""


@dataclass
class AgentInput:

    cnt_image_path: str = ""
    sub_tasks: list[SubTask] = field(default_factory=lambda: [])


@dataclass
class AnalysisInput:
    """
    Analysis for sub task
    """

    ref_id: int = 1
    ref_image_or_path: str = ""
    style_type: str = ""
    style_value: str | float = ""
    suggestion: str = ""

    sys_prompt: str = "Analysis the style of the given image."


@dataclass
class AnalysisOutput:

    semantic: dict[str, str] = field(default_factory=lambda: {})
    pixel: dict[str, str] = field(default_factory=lambda: {})


@dataclass
class TransferInput:

    prompt: str = ""
    cnt_image_or_path: str | Image.Image = ""
    ref_image_or_path: str | Image.Image = ""
    prompt_prefix: str = "Transfer the style of Picture 1 into the style of Picture 2."

    sys_prompt: str = "Make the style of Picture into another style."


@dataclass
class TransferOutput:

    instruct: str = ""
    cnt_image_or_path: str | Image.Image = ""
    ref_image_or_path: str | Image.Image = ""
    sty_image: Image.Image = None


@dataclass
class CriteriaInput:

    instruction: str = ""
    cnt_image_or_path: str | Image.Image = ""
    ref_image_or_path: str | Image.Image = ""
    sty_image_or_path: str | Image.Image = ""

    sys_prompts: dict[str, str] = field(
        default_factory=lambda: {
            "cs": "Score the semantic alignment degree between the two images.",
            "rs": "Score the style features between the two images.",
            "ds": "Score the instruction following ability of the image.",
        }
    )


@dataclass
class CriteriaOutput:

    cs_score: dict[str, str] = field(default_factory=lambda: {})
    rs_score: dict[str, str] = field(default_factory=lambda: {})
    ds_score: dict[str, str] = field(default_factory=lambda: {})


@dataclass
class SubTaskOutput:

    cnt_image_path: str = ""
    ref_image_path: str = ""
    analysis_output: AnalysisOutput = None
    transfer_output: TransferOutput = None
    criteria_output: CriteriaOutput = None


@dataclass
class AgentOutput:

    sty_image_path: str = ""
    sty_image: Image.Image = None
    sub_tasks: list[SubTask] = field(default_factory=lambda: [])
    analysis_outputs: list[AnalysisOutput] = field(default_factory=lambda: [])
    transfer_outputs: list[TransferOutput] = field(default_factory=lambda: [])
    criteria_outputs: list[CriteriaOutput] = field(default_factory=lambda: [])