from __future__ import annotations

import json
import random
from typing import Any, Iterable


SEED_AXIS_SALTS = {
    "category": 31,
    "subcategory": 37,
    "content": 41,
    "clothing": 41,
    "person": 43,
    "scene": 47,
    "pose": 53,
    "role": 57,
    "expression": 59,
    "composition": 61,
}

SEED_AXIS_ALIASES = {
    "category": ("category_seed", "category"),
    "subcategory": ("subcategory_seed", "subcategory"),
    "content": ("content_seed", "item_seed", "sexual_pose_seed", "content"),
    "clothing": ("clothing_seed", "outfit_seed", "wardrobe_seed", "content_seed", "content"),
    "person": ("person_seed", "appearance_seed", "cast_seed", "person"),
    "scene": ("scene_seed", "scene"),
    "pose": ("pose_seed", "sexual_pose_seed", "pose"),
    "role": ("role_seed", "role", "pose_seed", "sexual_pose_seed"),
    "expression": ("expression_seed", "face_seed", "expression"),
    "composition": ("composition_seed", "camera_seed", "composition"),
}

SEED_LOCK_AXES = (
    "category",
    "subcategory",
    "content",
    "clothing",
    "person",
    "scene",
    "pose",
    "role",
    "expression",
    "composition",
)
SEED_MODE_CHOICES = ["auto", "follow_main", "fixed", "random"]
SEED_REROLL_GROUPS = {
    "none": (),
    "category": ("category",),
    "subcategory": ("subcategory",),
    "content": ("content",),
    "clothing": ("clothing",),
    "person": ("person",),
    "scene": ("scene",),
    "pose": ("pose", "role"),
    "role": ("role",),
    "expression": ("expression",),
    "composition": ("composition",),
    "content_pose": ("content", "pose", "role"),
    "content_clothing": ("content", "clothing"),
    "clothing_pose": ("clothing", "pose", "role"),
    "scene_pose": ("scene", "pose", "role"),
}
SEED_REROLL_AXIS_CHOICES = list(SEED_REROLL_GROUPS.keys())


def _normal_key(value: Any) -> str:
    return str(value or "").strip().lower().replace("-", "_").replace(" ", "_")


def seed_mode_choices() -> list[str]:
    return list(SEED_MODE_CHOICES)


def normalize_seed_mode(value: Any) -> str:
    normalized = _normal_key(value)
    aliases = {
        "follow": "follow_main",
        "followmain": "follow_main",
        "follow_main_seed": "follow_main",
        "main": "follow_main",
        "main_seed": "follow_main",
    }
    normalized = aliases.get(normalized, normalized)
    return normalized if normalized in SEED_MODE_CHOICES else "auto"


def seed_reroll_axis_choices() -> list[str]:
    return list(SEED_REROLL_AXIS_CHOICES)


def normalize_reroll_axis(value: Any) -> str:
    normalized = _normal_key(value)
    aliases = {
        "contentpose": "content_pose",
        "contentclothing": "content_clothing",
        "clothingpose": "clothing_pose",
        "scenepose": "scene_pose",
    }
    normalized = aliases.get(normalized, normalized)
    return normalized if normalized in SEED_REROLL_GROUPS else "none"


def row_seed(seed: int, row_number: int, salt: int = 0) -> int:
    return int(seed) + int(row_number) * 1009 + salt * 9176


def build_seed_config_json(
    category_seed: int = -1,
    subcategory_seed: int = -1,
    content_seed: int = -1,
    person_seed: int = -1,
    scene_seed: int = -1,
    pose_seed: int = -1,
    role_seed: int = -1,
    expression_seed: int = -1,
    composition_seed: int = -1,
    category_seed_mode: str = "auto",
    subcategory_seed_mode: str = "auto",
    content_seed_mode: str = "auto",
    person_seed_mode: str = "auto",
    scene_seed_mode: str = "auto",
    pose_seed_mode: str = "auto",
    role_seed_mode: str = "auto",
    expression_seed_mode: str = "auto",
    composition_seed_mode: str = "auto",
    clothing_seed: int = -1,
    clothing_seed_mode: str = "auto",
) -> str:
    rng = random.SystemRandom()

    def axis_seed(value: int, mode: str) -> int:
        mode = normalize_seed_mode(mode)
        if mode == "auto":
            return int(value)
        if mode == "random":
            return rng.randint(0, 0xFFFFFFFF)
        if mode == "fixed":
            return max(0, int(value))
        return -1

    return json.dumps(
        {
            "category_seed": axis_seed(category_seed, category_seed_mode),
            "subcategory_seed": axis_seed(subcategory_seed, subcategory_seed_mode),
            "content_seed": axis_seed(content_seed, content_seed_mode),
            "clothing_seed": axis_seed(clothing_seed, clothing_seed_mode),
            "person_seed": axis_seed(person_seed, person_seed_mode),
            "scene_seed": axis_seed(scene_seed, scene_seed_mode),
            "pose_seed": axis_seed(pose_seed, pose_seed_mode),
            "role_seed": axis_seed(role_seed, role_seed_mode),
            "expression_seed": axis_seed(expression_seed, expression_seed_mode),
            "composition_seed": axis_seed(composition_seed, composition_seed_mode),
        },
        ensure_ascii=True,
        sort_keys=True,
    )


def build_seed_lock_config_json(
    base_seed: int = 20260614,
    reroll_axis: str = "none",
    reroll_seed: int = -1,
) -> str:
    base_seed = int(base_seed)
    reroll_seed = int(reroll_seed)
    reroll = set(SEED_REROLL_GROUPS[normalize_reroll_axis(reroll_axis)])
    config: dict[str, int] = {}
    for axis in SEED_LOCK_AXES:
        config[f"{axis}_seed"] = reroll_seed if axis in reroll else base_seed
    return json.dumps(config, ensure_ascii=True, sort_keys=True)


def parse_seed_config(seed_config: str | dict[str, Any] | None) -> dict[str, int]:
    if not seed_config:
        return {}
    if isinstance(seed_config, dict):
        raw = seed_config
    else:
        try:
            raw = json.loads(str(seed_config))
        except json.JSONDecodeError as exc:
            raise ValueError(f"Invalid seed_config JSON: {exc}") from exc
    if not isinstance(raw, dict):
        raise ValueError("seed_config must be a JSON object")
    parsed: dict[str, int] = {}
    for key, value in raw.items():
        try:
            parsed[str(key)] = int(value)
        except (TypeError, ValueError):
            continue
    return parsed


def configured_axis_seed(seed_config: dict[str, int], axis: str) -> int | None:
    for key in SEED_AXIS_ALIASES.get(axis, (axis,)):
        value = seed_config.get(key)
        if value is not None and value >= 0:
            return value
    return None


def configured_seed_from_axes(
    seed_config: str | dict[str, Any] | None,
    axes: Iterable[str],
    *,
    extra_keys: Iterable[str] = (),
) -> int | None:
    try:
        parsed = parse_seed_config(seed_config)
    except ValueError:
        return None
    for axis in axes:
        value = configured_axis_seed(parsed, axis)
        if value is not None:
            return value
    for key in extra_keys:
        value = parsed.get(str(key))
        if value is not None and value >= 0:
            return value
    return None


def axis_rng(seed_config: dict[str, int], axis: str, base_seed: int, row_number: int) -> random.Random:
    configured = configured_axis_seed(seed_config, axis)
    salt = SEED_AXIS_SALTS.get(axis, 0)
    if configured is None:
        return random.Random(row_seed(base_seed, row_number, salt))
    return random.Random(row_seed(configured, row_number, salt))


def axis_seed_trace(
    seed_config: str | dict[str, Any] | None,
    base_seed: int,
    row_number: int,
    axes: Iterable[str] = SEED_LOCK_AXES,
) -> dict[str, dict[str, int | str]]:
    parsed = parse_seed_config(seed_config)
    trace: dict[str, dict[str, int | str]] = {}
    for axis in axes:
        configured = configured_axis_seed(parsed, axis)
        seed_value = int(configured) if configured is not None else int(base_seed)
        source = "configured" if configured is not None else "main"
        salt = SEED_AXIS_SALTS.get(axis, 0)
        trace[axis] = {
            "source": source,
            "seed": seed_value,
            "rng_seed": row_seed(seed_value, row_number, salt),
        }
    return trace