backend/world.py

"""
The world of Polis: locations, the population, and the tick loop.

A "tick" is one step of simulated time. On each tick every agent:
  1. perceives nearby agents and its location,
  2. retrieves relevant memories,
  3. decides an action (and possibly a line of dialogue),
  4. moves toward a chosen location,
  5. records what happened,
  6. occasionally reflects.

The engine emits a compact list of *events* per tick that the 3D frontend
replays: movements, speech bubbles, new memories, relationship changes, and
emergent "signals" (rumors spreading, gatherings forming).
"""
from __future__ import annotations

import math
import random
import time
from dataclasses import dataclass, field
from typing import List, Optional

from .agents import Agent
from .llm import llm

# ---- map ---------------------------------------------------------------------
LOCATIONS = {
    "plaza":   {"x": 0.0,  "y": 0.0,  "label": "Central Plaza"},
    "market":  {"x": 6.0,  "y": 2.0,  "label": "Market"},
    "bakery":  {"x": -5.0, "y": 3.0,  "label": "Bakery"},
    "clinic":  {"x": 5.0,  "y": -4.0, "label": "Clinic"},
    "tavern":  {"x": -4.0, "y": -5.0, "label": "Tavern"},
    "garden":  {"x": 2.0,  "y": 6.0,  "label": "Garden"},
    "forge":   {"x": -7.0, "y": -1.0, "label": "Forge"},
    "harbor":  {"x": 8.0,  "y": -1.0, "label": "Harbor"},
}

PERSONAS = [
    ("mara",  "Mara",  34, "warm, curious, a natural connector", "baker",     "bakery", "#f6c453"),
    ("theo",  "Theo",  41, "gruff but loyal, hates gossip",       "blacksmith","forge",  "#e6704b"),
    ("ines",  "Ines",  29, "sharp, ambitious, keeps score",       "merchant",  "market", "#6ec6ff"),
    ("caleb", "Caleb", 52, "calm, wise, a good listener",         "medic",     "clinic", "#8bd17c"),
    ("nadia", "Nadia", 26, "playful, dramatic, spreads rumors",   "bard",      "tavern", "#c792ea"),
    ("otis",  "Otis",  38, "anxious, dependable, over-plans",     "harbormaster","harbor","#4dd0e1"),
    ("lena",  "Lena",  31, "idealistic, stubborn, community-minded","gardener","garden", "#a3e635"),
    ("rurik", "Rurik", 47, "quiet, observant, secretly generous", "trader",    "market", "#ff8a65"),
]


@dataclass
class World:
    seed: int = 42
    tick: int = 0
    agents: List[Agent] = field(default_factory=list)
    log: List[dict] = field(default_factory=list)   # rolling event feed
    pending_events: List[str] = field(default_factory=list)  # injected by user
    rng: random.Random = field(default_factory=lambda: random.Random(42))

    @classmethod
    def bootstrap(cls, seed: int = 42) -> "World":
        w = cls(seed=seed, rng=random.Random(seed))
        for pid, name, age, traits, role, home, color in PERSONAS:
            loc = LOCATIONS[home]
            a = Agent(id=pid, name=name, age=age, traits=traits, role=role,
                      home=home, x=loc["x"], y=loc["y"], color=color,
                      location=home)
            a.remember(0, f"My name is {name}. I work as a {role} and live near the {home}.",
                       kind="observation", importance=6.0)
            a.goal = f"do good work as the {role} and stay close to people I trust"
            w.agents.append(a)
        # seed a couple of relationships so the graph isn't empty
        w._link("mara", "nadia", 0.4)
        w._link("theo", "caleb", 0.5)
        w._link("ines", "rurik", -0.2)
        return w

    def _link(self, a: str, b: str, v: float) -> None:
        da = self.get(a); db = self.get(b)
        if da and db:
            da.relationships[b] = round(v, 2)
            db.relationships[a] = round(v, 2)

    def get(self, agent_id: str) -> Optional[Agent]:
        return next((a for a in self.agents if a.id == agent_id), None)

    # -- external events ------------------------------------------------------
    def inject_event(self, text: str) -> None:
        self.pending_events.append(text.strip())

    # -- one step of time -----------------------------------------------------
    def step(self) -> dict:
        self.tick += 1
        events: List[dict] = []

        # broadcast any injected world events as shared observations
        for ev in self.pending_events:
            for a in self.agents:
                a.remember(self.tick, f"Something happened in town: {ev}",
                           kind="observation", importance=8.0)
            events.append({"type": "world_event", "text": ev})
        self.pending_events.clear()

        # find conversation pairs (agents that are close together)
        pairs = self._nearby_pairs(radius=2.5)
        talked = set()

        for a in self.agents:
            partner = next((b for (x, b) in [(p[0], p[1]) for p in pairs] if x is a), None)
            # decide + act
            observation = self._perceive(a)
            mems = a.retrieve(observation, self.tick, k=4)
            mem_ctx = "; ".join(m.text for m in mems) or "nothing notable"

            if partner and (a.id, partner.id) not in talked and (partner.id, a.id) not in talked:
                line = self._dialogue(a, partner, mem_ctx)
                a.remember(self.tick, f"I said to {partner.name}: {line}",
                           kind="dialogue", importance=5.0)
                partner.remember(self.tick, f"{a.name} said to me: {line}",
                                 kind="dialogue", importance=5.0)
                self._adjust_relationship(a, partner, line)
                talked.add((a.id, partner.id))
                events.append({
                    "type": "dialogue", "agent": a.id, "to": partner.id,
                    "text": line, "x": a.x, "y": a.y,
                })
            else:
                action, target_loc = self._decide(a, mem_ctx)
                a.remember(self.tick, f"I {action}.", kind="observation")
                self._move_toward(a, target_loc)
                events.append({
                    "type": "action", "agent": a.id, "text": action,
                    "location": a.location, "x": round(a.x, 2), "y": round(a.y, 2),
                })

            # reflection
            refl = a.maybe_reflect(self.tick)
            if refl:
                events.append({"type": "reflection", "agent": a.id,
                               "text": refl.text})

        snapshot = {
            "tick": self.tick,
            "events": events,
            "agents": [a.snapshot() for a in self.agents],
            "budget": llm.ledger.as_dict(),
            "live": llm.live,
        }
        self.log.append(snapshot)
        if len(self.log) > 200:
            self.log = self.log[-200:]
        return snapshot

    # -- helpers --------------------------------------------------------------
    def _perceive(self, a: Agent) -> str:
        others = [b.name for b in self.agents
                  if b is not a and self._dist(a, b) < 3.0]
        who = f" I see {', '.join(others)} nearby." if others else " No one is close by."
        return f"I am at the {a.location}.{who} My goal: {a.goal}."

    def _decide(self, a: Agent, mem_ctx: str) -> tuple[str, str]:
        locs = ", ".join(LOCATIONS.keys())
        out = llm.chat(
            system=(f"You are {a.identity()} Decide your next small action in one short "
                    f"sentence, then name ONE place to go from [{locs}]. "
                    f"Format exactly: ACTION | PLACE"),
            user=f"Relevant memories: {mem_ctx}\nWhat do you do next?",
            temperature=0.9, max_tokens=50,
        )
        action, _, place = out.partition("|")
        place = place.strip().lower()
        place = place if place in LOCATIONS else self.rng.choice(list(LOCATIONS))
        action = action.strip().rstrip(".") or "wanders for a while"
        return action, place

    def _dialogue(self, a: Agent, b: Agent, mem_ctx: str) -> str:
        sentiment = a.relationships.get(b.id, 0.0)
        tone = "warmly" if sentiment > 0.2 else ("coldly" if sentiment < -0.2 else "neutrally")
        out = llm.chat(
            system=(f"You are {a.identity()} You are speaking {tone} to {b.name} "
                    f"({b.role}). Say ONE natural line of dialogue in quotes."),
            user=f"Your relevant memories: {mem_ctx}\nWhat do you say to {b.name}?",
            temperature=0.9, max_tokens=50,
        )
        return out.strip().strip('"')

    def _adjust_relationship(self, a: Agent, b: Agent, line: str) -> None:
        # cheap sentiment: warm words nudge up, harsh words nudge down
        warm = any(w in line.lower() for w in
                   ["thank", "friend", "help", "glad", "love", "welcome", "together"])
        cold = any(w in line.lower() for w in
                   ["no", "won't", "angry", "leave", "wrong", "never", "stop"])
        delta = 0.08 if warm else (-0.08 if cold else 0.02)
        for x, y in ((a, b), (b, a)):
            cur = x.relationships.get(y.id, 0.0)
            x.relationships[y.id] = round(max(-1.0, min(1.0, cur + delta)), 2)

    def _move_toward(self, a: Agent, place: str) -> None:
        loc = LOCATIONS[place]
        a.location = place
        # ease 60% of the way so movement looks gradual across ticks
        a.x += (loc["x"] - a.x) * 0.6 + self.rng.uniform(-0.4, 0.4)
        a.y += (loc["y"] - a.y) * 0.6 + self.rng.uniform(-0.4, 0.4)

    def _nearby_pairs(self, radius: float) -> List[tuple]:
        pairs = []
        seen = set()
        for a in self.agents:
            for b in self.agents:
                if a is b or (b.id, a.id) in seen:
                    continue
                if self._dist(a, b) < radius:
                    pairs.append((a, b))
                    seen.add((a.id, b.id))
        return pairs

    @staticmethod
    def _dist(a: Agent, b: Agent) -> float:
        return math.hypot(a.x - b.x, a.y - b.y)