#!/usr/bin/env python3
"""Generate a MazeX layer allocation from formula allocator rules.

This is a deterministic prototype allocator:
- reads MazeX_formula_allocator_rules.json
- reads MazeX_formula_assets.json
- allocates assets for every step in one layer
- writes a JSON allocation file
"""

from __future__ import annotations

import argparse
import hashlib
import json
import random
from pathlib import Path
from typing import Any


def clamp(value: int, low: int, high: int) -> int:
    return max(low, min(high, value))


def sha256_text(text: str) -> str:
    return hashlib.sha256(text.encode("utf-8")).hexdigest()


def rng_for(*parts: Any) -> random.Random:
    seed_text = "|".join(str(part) for part in parts)
    seed_int = int(sha256_text(seed_text), 16)
    return random.Random(seed_int)


def weighted_choice(items: list[dict[str, Any]], weights: dict[str, int], rng: random.Random) -> dict[str, Any]:
    total = sum(int(weights.get(item["primary_structure"], 1)) for item in items)
    pick = rng.uniform(0, total)
    cursor = 0.0
    for item in items:
        cursor += int(weights.get(item["primary_structure"], 1))
        if cursor >= pick:
            return item
    return items[-1]


def get_step_count(layer_no: int, rule: dict[str, Any]) -> int:
    return min(
        int(rule["base_steps"]) + (layer_no - 1) * int(rule["increase_per_layer"]),
        int(rule["max_steps"]),
    )


def get_layer_range(layer_no: int) -> tuple[int, int]:
    layer_min = clamp(1 + ((layer_no - 1) // 3), 1, 8)
    layer_max = clamp(layer_min + 5, 4, 10)
    return layer_min, layer_max


def get_band(progress: float) -> str:
    if progress < 0.2:
        return "warmup"
    if progress < 0.7:
        return "main"
    if progress < 0.9:
        return "advanced"
    return "boss"


def get_target_difficulty(
    layer_min: int,
    layer_max: int,
    band: str,
    step_rng: random.Random,
    boss_bonus: int,
) -> int:
    if band == "warmup":
        base = step_rng.randint(layer_min, min(layer_min + 2, layer_max))
        return clamp(base + step_rng.choice([-1, 0, 1]), layer_min, min(layer_min + 2, layer_max))
    if band == "main":
        low = min(layer_min + 2, layer_max)
        high = min(layer_min + 4, layer_max)
        base = step_rng.randint(low, high)
        return clamp(base + step_rng.choice([-1, 0, 1]), low, high)
    if band == "advanced":
        low = max(layer_min, layer_max - 2)
        high = max(layer_min, layer_max - 1)
        base = step_rng.randint(low, high)
        return clamp(base + step_rng.choice([-1, 0, 1]), low, high)
    return clamp(layer_max + boss_bonus, 1, 10)


def find_cost(target_difficulty: int, cost_rules: list[dict[str, Any]]) -> dict[str, float]:
    for rule in cost_rules:
        if int(rule["difficulty_min"]) <= target_difficulty <= int(rule["difficulty_max"]):
            return {
                "sample_cost_gas": float(rule["sample_cost_gas"]),
                "submit_cost_gas": float(rule["submit_cost_gas"]),
            }
    return {"sample_cost_gas": 0.01, "submit_cost_gas": 1.0}


def filter_candidates(
    assets: list[dict[str, Any]],
    target_difficulty: int,
    rules: dict[str, Any],
    previous_steps: list[dict[str, Any]],
) -> list[dict[str, Any]]:
    filters = rules["selection_filters"]
    allowed_states = set(filters["state_dependency_allowed"])
    allowed_verification = set(filters["verification_mode_allowed"])

    anti = rules["anti_duplication"]
    recent_structures = {
        step["primary_structure"] for step in previous_steps[-int(anti["recent_primary_structure_block"]) :]
    }
    recent_families = {step["family_id"] for step in previous_steps[-int(anti["recent_family_block"]) :]}
    structure_counts: dict[str, int] = {}
    for step in previous_steps:
        structure_counts[step["primary_structure"]] = structure_counts.get(step["primary_structure"], 0) + 1
    max_same_structure = int(anti["max_same_primary_structure_per_layer"])

    fallback: list[dict[str, Any]] = []
    for radius in range(0, 3):
        base = [
            asset
            for asset in assets
            if bool(asset.get("enabled")) == bool(filters["enabled_required"])
            and asset.get("state_dependency") in allowed_states
            and asset.get("verification_mode") in allowed_verification
            and int(asset["difficulty_min"]) <= target_difficulty + radius
            and int(asset["difficulty_max"]) >= target_difficulty - radius
            and structure_counts.get(asset["primary_structure"], 0) < max_same_structure
        ]

        if not base:
            continue
        if not fallback:
            fallback = base

        no_structure_repeat = [asset for asset in base if asset["primary_structure"] not in recent_structures]
        no_family_repeat = [asset for asset in no_structure_repeat if asset["family_id"] not in recent_families]

        if no_family_repeat:
            return no_family_repeat
        if no_structure_repeat:
            return no_structure_repeat

    if fallback:
        return fallback

    # Last-resort fallback: keep the layer generatable even if all caps are exhausted.
    return [
        asset
        for asset in assets
        if bool(asset.get("enabled")) == bool(filters["enabled_required"])
        and asset.get("state_dependency") in allowed_states
        and asset.get("verification_mode") in allowed_verification
    ]


def allocate_layer(
    layer_no: int,
    season_seed: str,
    previous_block_hash: str,
    rules: dict[str, Any],
    assets: list[dict[str, Any]],
) -> dict[str, Any]:
    step_count = get_step_count(layer_no, rules["step_count_rule"])
    layer_min, layer_max = get_layer_range(layer_no)
    boss_bonus = int(rules["layer_difficulty_rule"]["boss_bonus"])
    weights = {k: int(v) for k, v in rules["structure_weights"].items()}
    steps: list[dict[str, Any]] = []

    for step_no in range(1, step_count + 1):
        progress = 0.0 if step_count == 1 else (step_no - 1) / (step_count - 1)
        band = get_band(progress)
        step_rng = rng_for(season_seed, layer_no, step_no, previous_block_hash)
        target_difficulty = get_target_difficulty(layer_min, layer_max, band, step_rng, boss_bonus)
        candidates = filter_candidates(assets, target_difficulty, rules, steps)
        if not candidates:
            raise RuntimeError(f"No candidates for layer={layer_no} step={step_no} difficulty={target_difficulty}")

        asset = weighted_choice(candidates, weights, step_rng)
        allocator_seed = sha256_text(f"{season_seed}:{layer_no}:{step_no}:{previous_block_hash}")
        formula_seed = sha256_text(f"{allocator_seed}:{asset['asset_id']}:formula")
        parameter_seed = sha256_text(f"{formula_seed}:parameters")
        formula_hash_commit = sha256_text(
            f"{asset['asset_id']}:{asset['formula_template']}:{formula_seed}:{parameter_seed}"
        )
        theme_candidates = asset.get("theme_candidates", [])
        theme_hint = theme_candidates[step_rng.randrange(len(theme_candidates))] if theme_candidates else ""
        costs = find_cost(target_difficulty, rules["cost_rules"])
        boss_zone = band == "boss"

        steps.append(
            {
                "layer_no": layer_no,
                "step_no": step_no,
                "step_count": step_count,
                "progress_band": band,
                "target_difficulty": target_difficulty,
                "asset_id": asset["asset_id"],
                "pool_id": asset["pool_id"],
                "primary_structure": asset["primary_structure"],
                "family_id": asset["family_id"],
                "theme_hint": theme_hint,
                "allocator_seed": allocator_seed,
                "formula_seed": formula_seed,
                "parameter_seed": parameter_seed,
                "formula_hash_commit": formula_hash_commit,
                "sample_cost_gas": costs["sample_cost_gas"],
                "submit_cost_gas": costs["submit_cost_gas"],
                "reward_weight": target_difficulty * (2 if boss_zone else 1),
                "boss_zone": boss_zone,
            }
        )

    return {
        "version": rules["version"],
        "layer_no": layer_no,
        "step_count": step_count,
        "layer_difficulty_min": layer_min,
        "layer_difficulty_max": layer_max,
        "season_seed": season_seed,
        "previous_block_hash": previous_block_hash,
        "steps": steps,
    }


def main() -> None:
    parser = argparse.ArgumentParser(description="Allocate MazeX formulas for one layer.")
    parser.add_argument("--layer", type=int, required=True, help="Layer number, starting from 1.")
    parser.add_argument("--rules", default="MazeX_formula_allocator_rules.json")
    parser.add_argument("--assets", default=None, help="Override asset index path.")
    parser.add_argument("--season-seed", default="mazex-season-1")
    parser.add_argument("--previous-block-hash", default="genesis")
    parser.add_argument("--out", default=None, help="Output JSON path.")
    args = parser.parse_args()

    rules_path = Path(args.rules)
    rules = json.loads(rules_path.read_text(encoding="utf-8"))
    assets_path = Path(args.assets or rules["asset_index"])
    assets_doc = json.loads(assets_path.read_text(encoding="utf-8"))

    allocation = allocate_layer(
        layer_no=args.layer,
        season_seed=args.season_seed,
        previous_block_hash=args.previous_block_hash,
        rules=rules,
        assets=assets_doc["assets"],
    )

    out_path = Path(args.out or f"MazeX_layer_{args.layer}_allocation.json")
    out_path.write_text(json.dumps(allocation, ensure_ascii=False, indent=2), encoding="utf-8")
    print(f"Wrote {out_path} with {allocation['step_count']} steps.")


if __name__ == "__main__":
    main()
