"""Angle presets and dataset configuration for AngleForge.

The camera prompts are bilingual (Chinese + English) to match the
``dx8152/Qwen-Edit-2509-Multiple-angles`` LoRA training data, which gives
the most reliable results with Qwen Image Edit 2509.
"""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import Dict, List


# --------------------------------------------------------------------------- #
# Camera-angle presets
# --------------------------------------------------------------------------- #
# key -> (human label, prompt). An empty prompt means "keep the original".

ANGLE_PRESETS: Dict[str, tuple[str, str]] = {
    "original": ("Original", ""),
    "top_down": ("Top-down / overhead", "将镜头转为俯视 Turn the camera to a top-down view."),
    "birds_eye": ("Bird's-eye", "将相机转向鸟瞰视角 Turn the camera to a bird's-eye view."),
    "worms_eye": ("Worm's-eye", "将相机切换到仰视视角 Turn the camera to a worm's-eye view."),
    "rotate_left_45": ("Rotate left 45°", "将镜头向左旋转45度 Rotate the camera 45 degrees to the left."),
    "rotate_right_45": ("Rotate right 45°", "将镜头向右旋转45度 Rotate the camera 45 degrees to the right."),
    "rotate_left_90": ("Rotate left 90°", "将镜头向左旋转90度 Rotate the camera 90 degrees to the left."),
    "rotate_right_90": ("Rotate right 90°", "将镜头向右旋转90度 Rotate the camera 90 degrees to the right."),
    "close_up": ("Close-up", "将镜头转为特写镜头 Turn the camera to a close-up."),
    "wide_angle": ("Wide-angle", "将镜头转为广角镜头 Turn the camera to a wide-angle lens."),
    "move_left": ("Move left", "将镜头向左移动 Move the camera left."),
    "move_right": ("Move right", "将镜头向右移动 Move the camera right."),
    "move_forward": ("Move forward", "将镜头向前移动 Move the camera forward."),
    "move_down": ("Move down", "将镜头向下移动 Move the camera down."),
}

DEFAULT_ANGLES: List[str] = [
    "original",
    "top_down",
    "birds_eye",
    "worms_eye",
    "rotate_left_45",
    "rotate_right_45",
    "close_up",
    "wide_angle",
]

# Models used by the local Qwen backend (mirrors linoyts/Qwen-Image-Edit-Angles).
QWEN_BASE_MODEL = "Qwen/Qwen-Image-Edit-2509"
QWEN_RAPID_TRANSFORMER = "linoyts/Qwen-Image-Edit-Rapid-AIO"
ANGLES_LORA_REPO = "dx8152/Qwen-Edit-2509-Multiple-angles"
ANGLES_LORA_WEIGHT = "镜头转换.safetensors"

# Model id used by the serverless HF Inference Providers backend.
ANGLES_INFERENCE_MODEL = "dx8152/Qwen-Edit-2509-Multiple-angles"


def angle_prompt(angle_key: str) -> str:
    """Return the bilingual camera prompt for an angle preset key."""
    if angle_key not in ANGLE_PRESETS:
        raise ValueError(f"Unknown angle preset: {angle_key}")
    return ANGLE_PRESETS[angle_key][1]


@dataclass
class DatasetConfig:
    """All knobs controlling one image-dataset build."""

    out_dir: str = "output"
    dataset_name: str = "industrial_angles"

    # Output image geometry.
    image_size: int = 512  # longest side, snapped to a multiple of 8

    # Angle generation.
    angles: List[str] = field(default_factory=lambda: list(DEFAULT_ANGLES))
    variations_per_angle: int = 1

    # Extra, cheap (non-AI) augmentations applied to every saved image family.
    plain_augmentations_per_image: int = 0

    # Qwen inference controls.
    num_inference_steps: int = 4
    true_guidance_scale: float = 1.0

    # Split & reproducibility.
    test_ratio: float = 0.2
    seed: int = 1234