from __future__ import annotations

import json
import random

try:
    from .prompt_builder import (
        build_cast_config_json,
        build_category_config_json,
        cast_preset_choices,
        category_preset_choices,
        subcategory_choices,
    )
    from .location_config import (
        build_composition_pool_json,
        build_location_pool_json,
        build_thematic_location_json,
        composition_pool_preset_choices,
        location_pool_preset_choices,
        location_theme_choices,
    )
except ImportError:  # Allows local smoke tests from the repository root.
    from prompt_builder import (
        build_cast_config_json,
        build_category_config_json,
        cast_preset_choices,
        category_preset_choices,
        subcategory_choices,
    )
    from location_config import (
        build_composition_pool_json,
        build_location_pool_json,
        build_thematic_location_json,
        composition_pool_preset_choices,
        location_pool_preset_choices,
        location_theme_choices,
    )


SXCP_CATEGORY_CONFIG = "SXCP_CATEGORY_CONFIG"
SXCP_LOCATION_CONFIG = "SXCP_LOCATION_CONFIG"
SXCP_COMPOSITION_CONFIG = "SXCP_COMPOSITION_CONFIG"
SXCP_CAST_CONFIG = "SXCP_CAST_CONFIG"
SXCP_SEED_CONFIG = "SXCP_SEED_CONFIG"


class SxCPCategoryPreset:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "preset": (category_preset_choices(), {"default": "auto_weighted"}),
                "subcategory": (subcategory_choices(), {"default": "random"}),
            }
        }

    RETURN_TYPES = (SXCP_CATEGORY_CONFIG, "STRING", "STRING")
    RETURN_NAMES = ("category_config", "category", "subcategory")
    FUNCTION = "build"
    CATEGORY = "prompt_builder"

    def build(self, preset, subcategory):
        config = build_category_config_json(preset=preset, subcategory=subcategory)
        parsed = json.loads(config)
        return config, parsed["category"], parsed["subcategory"]


class SxCPLocationPool:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "enabled": ("BOOLEAN", {"default": True}),
                "combine_mode": (["replace", "add"], {"default": "replace"}),
                "preset": (location_pool_preset_choices(), {"default": "custom_only"}),
                "custom_locations": ("STRING", {"default": "", "multiline": True}),
            },
            "optional": {
                "location_config": (SXCP_LOCATION_CONFIG,),
            },
        }

    RETURN_TYPES = (SXCP_LOCATION_CONFIG, "STRING")
    RETURN_NAMES = ("location_config", "summary")
    FUNCTION = "build"
    CATEGORY = "prompt_builder"

    def build(self, enabled, combine_mode, preset, custom_locations, location_config=""):
        config = build_location_pool_json(
            enabled=enabled,
            combine_mode=combine_mode,
            preset=preset,
            custom_locations=custom_locations or "",
            location_config=location_config or "",
        )
        parsed = json.loads(config)
        return config, parsed.get("summary", "")


class SxCPCompositionPool:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "enabled": ("BOOLEAN", {"default": True}),
                "combine_mode": (["replace", "add"], {"default": "replace"}),
                "preset": (composition_pool_preset_choices(), {"default": "no_outfit_check"}),
                "custom_compositions": ("STRING", {"default": "", "multiline": True}),
            },
            "optional": {
                "composition_config": (SXCP_COMPOSITION_CONFIG,),
            },
        }

    RETURN_TYPES = (SXCP_COMPOSITION_CONFIG, "STRING")
    RETURN_NAMES = ("composition_config", "summary")
    FUNCTION = "build"
    CATEGORY = "prompt_builder"

    def build(self, enabled, combine_mode, preset, custom_compositions, composition_config=""):
        config = build_composition_pool_json(
            enabled=enabled,
            combine_mode=combine_mode,
            preset=preset,
            custom_compositions=custom_compositions or "",
            composition_config=composition_config or "",
        )
        parsed = json.loads(config)
        return config, parsed.get("summary", "")


class SxCPLocationTheme:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "enabled": ("BOOLEAN", {"default": True}),
                "combine_mode": (["replace", "add"], {"default": "replace"}),
                "theme": (location_theme_choices(), {"default": "semi_public_affair"}),
                "custom_locations": ("STRING", {"default": "", "multiline": True}),
                "custom_compositions": ("STRING", {"default": "", "multiline": True}),
            },
            "optional": {
                "location_config": (SXCP_LOCATION_CONFIG,),
                "composition_config": (SXCP_COMPOSITION_CONFIG,),
            },
        }

    RETURN_TYPES = (SXCP_LOCATION_CONFIG, SXCP_COMPOSITION_CONFIG, "STRING")
    RETURN_NAMES = ("location_config", "composition_config", "summary")
    FUNCTION = "build"
    CATEGORY = "prompt_builder"

    def build(
        self,
        enabled,
        combine_mode,
        theme,
        custom_locations,
        custom_compositions,
        location_config="",
        composition_config="",
    ):
        return build_thematic_location_json(
            enabled=enabled,
            combine_mode=combine_mode,
            theme=theme,
            custom_locations=custom_locations or "",
            custom_compositions=custom_compositions or "",
            location_config=location_config or "",
            composition_config=composition_config or "",
        )


class SxCPCastControl:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "cast_mode": (cast_preset_choices(), {"default": "mixed_couple"}),
                "women_count": ("INT", {"default": 1, "min": 0, "max": 12, "step": 1}),
                "men_count": ("INT", {"default": 1, "min": 0, "max": 12, "step": 1}),
            }
        }

    RETURN_TYPES = (SXCP_CAST_CONFIG, "INT", "INT", "STRING")
    RETURN_NAMES = ("cast_config", "women_count", "men_count", "cast_summary")
    FUNCTION = "build"
    CATEGORY = "prompt_builder"

    def build(self, cast_mode, women_count, men_count):
        config = build_cast_config_json(cast_mode=cast_mode, women_count=women_count, men_count=men_count)
        parsed = json.loads(config)
        summary = f"{parsed['women_count']} women, {parsed['men_count']} men"
        return config, parsed["women_count"], parsed["men_count"], summary


class SxCPCastBias:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "seed": ("INT", {"default": -1, "min": -1, "max": 0xFFFFFFFF, "step": 1}),
                "row_number": ("INT", {"default": 1, "min": 1, "max": 1000000, "step": 1}),
                "women_weights": ("STRING", {"default": "0.60,0.25,0.10,0.05"}),
                "women_start_count": ("INT", {"default": 1, "min": 0, "max": 12, "step": 1}),
                "men_weights": ("STRING", {"default": "0.45,0.40,0.10,0.05"}),
                "men_start_count": ("INT", {"default": 0, "min": 0, "max": 12, "step": 1}),
                "empty_behavior": (["force_one_woman", "force_one_man", "allow_empty"], {"default": "force_one_woman"}),
            },
            "optional": {
                "seed_config": (SXCP_SEED_CONFIG,),
            },
        }

    RETURN_TYPES = (SXCP_CAST_CONFIG, "INT", "INT", "STRING")
    RETURN_NAMES = ("cast_config", "women_count", "men_count", "cast_summary")
    FUNCTION = "build"
    CATEGORY = "prompt_builder"

    @staticmethod
    def _configured_cast_seed(seed_config):
        if not seed_config:
            return None
        if isinstance(seed_config, dict):
            raw = seed_config
        else:
            try:
                raw = json.loads(str(seed_config))
            except (TypeError, ValueError, json.JSONDecodeError):
                return None
        if not isinstance(raw, dict):
            return None
        for key in ("category_seed", "content_seed", "role_seed", "seed", "global_seed"):
            try:
                value = int(raw.get(key))
            except (TypeError, ValueError):
                continue
            if value >= 0:
                return value
        return None

    @staticmethod
    def _weight_pairs(weights_text, start_count):
        pairs = []
        start = max(0, min(12, int(start_count)))
        parts = str(weights_text or "").replace("\n", ",").split(",")
        for offset, raw in enumerate(parts):
            count = start + offset
            if count > 12:
                break
            try:
                weight = float(raw.strip())
            except (TypeError, ValueError):
                continue
            if weight > 0:
                pairs.append((count, weight))
        return pairs or [(start, 1.0)]

    @staticmethod
    def _weighted_count(rng, pairs):
        total = sum(weight for _count, weight in pairs)
        point = rng.random() * total
        upto = 0.0
        for count, weight in pairs:
            upto += weight
            if point <= upto:
                return int(count)
        return int(pairs[-1][0])

    @classmethod
    def IS_CHANGED(cls, *args, **kwargs):
        seed_value = kwargs.get("seed")
        if seed_value is None and args:
            seed_value = args[0]
        seed_config = kwargs.get("seed_config", "")
        if not seed_config and len(args) > 7:
            seed_config = args[7]
        try:
            seed = int(seed_value)
        except (TypeError, ValueError):
            seed = -1
        if seed < 0 and cls._configured_cast_seed(seed_config) is None:
            return random.random()
        return tuple(args), tuple(sorted(kwargs.items()))

    def build(
        self,
        seed,
        row_number,
        women_weights,
        women_start_count,
        men_weights,
        men_start_count,
        empty_behavior,
        seed_config="",
    ):
        configured_seed = self._configured_cast_seed(seed_config)
        if configured_seed is None and int(seed) < 0:
            rng = random.Random(random.getrandbits(64))
        else:
            cast_seed = configured_seed if configured_seed is not None else int(seed)
            rng = random.Random(f"sxcp_cast_bias:{cast_seed}:{int(row_number)}")
        women_pairs = self._weight_pairs(women_weights, women_start_count)
        men_pairs = self._weight_pairs(men_weights, men_start_count)
        women_count = self._weighted_count(rng, women_pairs)
        men_count = self._weighted_count(rng, men_pairs)
        if women_count + men_count == 0:
            if empty_behavior == "force_one_man":
                men_count = 1
            elif empty_behavior != "allow_empty":
                women_count = 1
        config = build_cast_config_json(cast_mode="custom_counts", women_count=women_count, men_count=men_count)
        parsed = json.loads(config)
        summary = f"weighted cast: {parsed['women_count']} women, {parsed['men_count']} men"
        return config, parsed["women_count"], parsed["men_count"], summary


NODE_CLASS_MAPPINGS = {
    "SxCPCategoryPreset": SxCPCategoryPreset,
    "SxCPLocationPool": SxCPLocationPool,
    "SxCPCompositionPool": SxCPCompositionPool,
    "SxCPLocationTheme": SxCPLocationTheme,
    "SxCPCastControl": SxCPCastControl,
    "SxCPCastBias": SxCPCastBias,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "SxCPCategoryPreset": "SxCP Category Preset",
    "SxCPLocationPool": "SxCP Location Pool",
    "SxCPCompositionPool": "SxCP Composition Pool",
    "SxCPLocationTheme": "SxCP Location Theme",
    "SxCPCastControl": "SxCP Cast Control",
    "SxCPCastBias": "SxCP Cast Bias",
}