src/world_simulator/simulation/spawning.py

from __future__ import annotations

import random

from world_simulator.config import GameConfig
from world_simulator.domain import (
    Beast,
    CountryState,
    House,
    MemoryEntry,
    Npc,
    ResourceNode,
    Terrain,
    TreasuryState,
    Vec3,
    WorldState,
)
from world_simulator.simulation.mechanics import BLOCK_SIZE, GRID_BLOCKS
from world_simulator.simulation.roles import normalize_role

_NAMES = [
    "Ada",
    "Boris",
    "Cora",
    "Dima",
    "Elena",
    "Farid",
    "Gita",
    "Hana",
    "Ivan",
    "Juno",
    "Kira",
    "Lena",
    "Mira",
    "Niko",
    "Oleg",
    "Pavel",
    "Raya",
    "Sofia",
    "Toma",
    "Vera",
]

_ROLES = ["citizen"]

_QWEN_NAMES = ["Kira", "Lena", "Mira", "Niko", "Oleg", "Pavel"]

def create_world(config: GameConfig) -> WorldState:
    rng = random.Random(config.world.seed)
    terrain = Terrain(
        kind=config.world.terrain,
        width=config.world.width,
        depth=config.world.depth,
    )
    village_npcs = [
        _spawn_npc(index=index, config=config, rng=rng, country_id="nemotron")
        for index in range(config.npcs.count)
    ]
    qwen_npcs = _spawn_qwen_npcs(config, count=config.npcs.count) if config.world.survival else []
    houses = _spawn_houses(config)

    if config.world.survival:
        _seed_survival_relationships(village_npcs + qwen_npcs)
        _assign_starting_homes(village_npcs + qwen_npcs, houses)
        _initialize_importance(village_npcs + qwen_npcs)

    all_npcs = village_npcs + qwen_npcs
    population = sum(1 for npc in all_npcs if npc.health > 0)
    world = WorldState(
        tick=0,
        seed=config.world.seed,
        terrain=terrain,
        npcs=all_npcs,
        resource_nodes=_spawn_resource_nodes(config),
        beasts=_spawn_beasts(config),
        houses=houses,
        countries=_spawn_countries(config, all_npcs),
        population=population,
        population_cap=max(24, population + 8),
        peak_population=population,
        overseer_mode="off",
        overseer_cycle_ticks=0,
    )
    _assign_initial_rulers(world)
    return world


def _spawn_npc(*, index: int, config: GameConfig, rng: random.Random, country_id: str) -> Npc:
    half_width = config.world.width / 2
    half_depth = config.world.depth / 2

    name = _NAMES[index % len(_NAMES)]
    role = normalize_role(_ROLES[index % len(_ROLES)])
    if config.world.survival:
        # Village is anchored on the left side of the map.
        village_x = -(half_width * 2 / 3)
        ring = index % 6
        offsets = [
            (0.0, 0.0),
            (1.2, 0.8),
            (-1.0, -0.9),
            (3.4, 0.0),
            (-3.4, 0.0),
            (0.0, 3.4),
        ]
        dx, z = offsets[ring]
        x = village_x + dx
    else:
        x = _random_block_center(rng, half_width=half_width)
        z = _random_block_center(rng, half_width=half_depth)

    npc = Npc(
        id=f"npc-{index + 1:03d}",
        name=name,
        role=role,
        position=Vec3(x=x, y=0.0, z=z),
        country_id=country_id,
        attack_damage=rng.randint(5, 15),
        memory=[MemoryEntry(tick=0, text=f"{name} arrived as a {role}.")],
        age=60 if config.world.survival else 0,
        max_age=_random_max_age(config.world.seed, index),
    )

    if config.world.survival:
        # Independent per-NPC RNG so survival inventory never perturbs the shared
        # position/damage stream (keeps non-survival spawning deterministic).
        inv_rng = random.Random(config.world.seed + index)
        npc.hunger = 0.0 if index == 0 else 20.0 + inv_rng.randint(0, 10)
        npc.safety = 100.0 if index == 0 else 0.0
        npc.inventory_food = 2 if index == 0 else inv_rng.randint(0, 2)
        npc.inventory_herbs = inv_rng.randint(0, 1)
        # Weapons stay scarce: citizens start unarmed and must gather them from the
        # (limited) weapon nodes during play.
        npc.inventory_weapon = 0
        if index < 2:
            npc.inventory_wood = 5

    return npc


def _spawn_resource_nodes(config: GameConfig) -> list[ResourceNode]:
    if not config.world.survival:
        return []
    rng = random.Random(f"{config.world.seed}:initial_resources")
    # Weapons are NOT placed randomly: one node is mirrored near each base so both
    # nations have equal, but still limited, access to arming up (see below).
    counts = {"food": 4, "herbs": 2, "wood": 3}
    max_amounts = {"food": 5, "herbs": 4, "wood": 6, "weapon": 2}
    nodes: list[ResourceNode] = []

    village_x = -(config.world.width / 2 * 2 / 3)
    village_food_offsets = [
        (2 * BLOCK_SIZE, BLOCK_SIZE),
        (-2 * BLOCK_SIZE, 2 * BLOCK_SIZE),
        (BLOCK_SIZE, -2 * BLOCK_SIZE),
    ]
    for i, (dx, dz) in enumerate(village_food_offsets, start=1):
        nodes.append(
            ResourceNode(
                id=f"res_food_village_{i}",
                resource_type="food",
                position=Vec3(x=round(village_x + dx, 3), y=0.0, z=round(dz, 3)),
                amount=5,
                max_amount=5,
            )
        )

    for resource_type, count in counts.items():
        for index in range(1, count + 1):
            max_amount = max_amounts[resource_type]
            nodes.append(
                ResourceNode(
                    id=f"res_{resource_type}_{index}",
                    resource_type=resource_type,
                    position=_random_world_position(
                        rng,
                        half_width=config.world.width / 2,
                        half_depth=config.world.depth / 2,
                    ),
                    amount=rng.randint(max(1, max_amount // 2), max_amount),
                    max_amount=max_amount,
                )
            )

    # One weapon node per flank, mirrored across the x-axis so neither nation is
    # favoured. Offset off the base centre so it does not sit on the home/treasury.
    qwen_x = config.world.width / 2 * 2 / 3
    weapon_max = max_amounts["weapon"]
    weapon_z = -2 * BLOCK_SIZE
    for index, base_x in enumerate((-qwen_x, qwen_x), start=1):
        nodes.append(
            ResourceNode(
                id=f"res_weapon_{index}",
                resource_type="weapon",
                position=Vec3(x=round(base_x, 3), y=0.0, z=round(weapon_z, 3)),
                amount=weapon_max,
                max_amount=weapon_max,
            )
        )
    return nodes


def _spawn_beasts(config: GameConfig) -> list[Beast]:
    if not config.world.survival:
        return []
    # One beast on each flank so both nations face an equal threat: beast_1 near
    # the Qwen side (+x), beast_2 mirrored near the Nemotron side (-x).
    edge_x = config.world.width / 2 - BLOCK_SIZE
    return [
        Beast(
            "beast_1",
            Vec3(x=edge_x, y=0.0, z=0.0),
            health=60.0,
            damage=9.0,
        ),
        Beast(
            "beast_2",
            Vec3(x=-edge_x, y=0.0, z=0.0),
            health=60.0,
            damage=9.0,
        ),
    ]


def _spawn_houses(config: GameConfig) -> list[House]:
    if not config.world.survival:
        return []
    village_x = -(config.world.width / 2 * 2 / 3)
    qwen_x = config.world.width / 2 * 2 / 3
    return [
        House(
            id="house_nemotron_001",
            position=Vec3(x=village_x, y=0.0, z=0.0),
            hp=60.0,
            max_hp=60.0,
            state="completed",
            build_progress=10,
            capacity=3,
        ),
        House(
            id="house_qwen_001",
            position=Vec3(x=qwen_x, y=0.0, z=0.0),
            hp=60.0,
            max_hp=60.0,
            state="completed",
            build_progress=10,
            capacity=3,
        ),
    ]


def _seed_survival_relationships(npcs: list[Npc]) -> None:
    for npc in npcs:
        for other in npcs:
            if other.id == npc.id:
                continue
            npc.relationships[other.id] = 0.45


def _assign_starting_homes(npcs: list[Npc], houses: list[House]) -> None:
    if not houses:
        return
    for npc in npcs:
        home = min(houses, key=lambda house: abs(house.position.x - npc.position.x))
        npc.home_house_id = home.id
    for house in houses:
        citizens = [npc.id for npc in npcs if npc.home_house_id == house.id]
        house.occupant_ids = citizens[: house.capacity]


def _initialize_importance(npcs: list[Npc]) -> None:
    counts: dict[str, int] = {}
    for npc in npcs:
        role = normalize_role(npc.role)
        counts[role] = counts.get(role, 0) + 1
    for npc in npcs:
        role_count = max(1, counts[normalize_role(npc.role)])
        npc.importance = round(1.0 + 2.0 * (1 / role_count), 3)


def _spawn_qwen_npcs(config: GameConfig, *, count: int) -> list[Npc]:
    half_width = config.world.width / 2
    qwen_x = half_width * 2 / 3
    offsets = [
        (0.0, 0.0),
        (BLOCK_SIZE, 0.0),
        (-BLOCK_SIZE, 0.0),
        (0.0, BLOCK_SIZE),
        (BLOCK_SIZE, BLOCK_SIZE),
        (-BLOCK_SIZE, BLOCK_SIZE),
    ]
    npcs = []
    inv_rng = random.Random(f"{config.world.seed}:qwen_inventory")
    for i in range(count):
        dx, dz = offsets[i % len(offsets)]
        name = _QWEN_NAMES[i % len(_QWEN_NAMES)]
        npc = Npc(
            id=f"qwen-{i + 1:03d}",
            name=name,
            role="citizen",
            position=Vec3(x=round(qwen_x + dx, 3), y=0.0, z=round(dz, 3)),
            attack_damage=inv_rng.randint(5, 12),
            memory=[MemoryEntry(tick=0, text=f"{name} joined Qwen as a free spirit.")],
            age=25,
            max_age=550,
            unrestricted_actions=True,
            connector_id="qwen",
            country_id="qwen",
            personality="curious and resourceful",
            hunger=15.0,
            safety=60.0,
            inventory_food=inv_rng.randint(1, 3),
            inventory_herbs=inv_rng.randint(0, 1),
        )
        npcs.append(npc)
    return npcs


def _spawn_countries(config: GameConfig, npcs: list[Npc]) -> list[CountryState]:
    if not config.world.survival:
        return []
    half_width = config.world.width / 2
    nemotron_x = -(half_width * 2 / 3)
    qwen_x = half_width * 2 / 3
    citizens_by_country = {
        "nemotron": [npc.id for npc in npcs if npc.country_id == "nemotron"],
        "qwen": [npc.id for npc in npcs if npc.country_id == "qwen"],
    }
    return [
        CountryState(
            id="nemotron",
            name="Nemotron",
            color="#4f8cff",
            badge="N",
            citizen_ids=citizens_by_country["nemotron"],
            treasury=TreasuryState(
                id="treasury_nemotron",
                position=Vec3(x=nemotron_x, y=0.0, z=BLOCK_SIZE * 1.5),
                resources={"food": 2, "herbs": 1, "wood": 20, "coins": 25},
            ),
        ),
        CountryState(
            id="qwen",
            name="Qwen",
            color="#d86bff",
            badge="Q",
            citizen_ids=citizens_by_country["qwen"],
            treasury=TreasuryState(
                id="treasury_qwen",
                position=Vec3(x=qwen_x, y=0.0, z=BLOCK_SIZE * 1.5),
                resources={"food": 2, "herbs": 1, "wood": 20, "coins": 25},
            ),
        ),
    ]


def _assign_initial_rulers(world: WorldState) -> None:
    for country in world.countries:
        ruler = next((npc for npc in world.living_npcs() if npc.id in country.citizen_ids), None)
        if ruler is None:
            continue
        country.ruler_id = ruler.id
        ruler.special_status = "ruler"


def _random_max_age(seed: int, index: int) -> int:
    return random.Random(f"{seed}:{index}:max_age").randint(320, 480)


def _random_block_center(rng: random.Random, *, half_width: float) -> float:
    first_center = -half_width + (BLOCK_SIZE / 2)
    num_blocks = round(half_width * 2 / BLOCK_SIZE)
    return first_center + (rng.randrange(num_blocks) * BLOCK_SIZE)


def _random_world_position(
    rng: random.Random,
    *,
    half_width: float,
    half_depth: float,
) -> Vec3:
    return Vec3(
        x=round(rng.uniform(-half_width + BLOCK_SIZE, half_width - BLOCK_SIZE), 3),
        y=0.0,
        z=round(rng.uniform(-half_depth + BLOCK_SIZE, half_depth - BLOCK_SIZE), 3),
    )