"""
Generative agents for Polis.

Each agent owns a *memory stream* — a time-ordered list of observations,
reflections and plans. When the agent needs to act, it retrieves the most
relevant memories using the scoring function from Park et al. 2023
("Generative Agents: Interactive Simulacra of Human Behavior"):

    score = alpha_recency * recency
          + alpha_importance * importance
          + alpha_relevance * relevance

Reflection periodically synthesises low-level memories into higher-level
insights, which are themselves written back into the stream. That feedback
loop is what makes behaviour compound into something that looks like a
personality over time.
"""
from __future__ import annotations

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

from .llm import llm, cosine

# retrieval weights
A_RECENCY = 1.0
A_IMPORTANCE = 1.0
A_RELEVANCE = 1.0
RECENCY_DECAY = 0.985          # per-tick exponential decay
REFLECT_EVERY = 12            # ticks of accumulated importance before reflecting
REFLECT_THRESHOLD = 40.0     # sum-of-importance trigger


@dataclass
class Memory:
    id: int
    tick: int
    kind: str                # "observation" | "reflection" | "plan" | "dialogue"
    text: str
    importance: float        # 1..10
    embedding: List[float] = field(default_factory=list)
    last_access_tick: int = 0
    causes: List[int] = field(default_factory=list)  # ids of parent memories

    def to_public(self) -> dict:
        return {
            "id": self.id, "tick": self.tick, "kind": self.kind,
            "text": self.text, "importance": round(self.importance, 1),
            "causes": self.causes,
        }


@dataclass
class Agent:
    id: str
    name: str
    age: int
    traits: str
    role: str                # e.g. "baker", "medic", "gossip"
    home: str
    x: float
    y: float
    color: str
    # runtime state
    mood: str = "neutral"
    goal: str = ""
    location: str = "home"
    memories: List[Memory] = field(default_factory=list)
    relationships: dict = field(default_factory=dict)  # other_id -> sentiment
    _mem_counter: int = 0
    _importance_since_reflect: float = 0.0

    # -- persona prompt -------------------------------------------------------
    def identity(self) -> str:
        return (f"{self.name}, age {self.age}. Role: {self.role}. "
                f"Personality: {self.traits}. Currently feeling {self.mood}.")

    # -- memory ---------------------------------------------------------------
    def remember(self, tick: int, text: str, kind: str = "observation",
                 importance: Optional[float] = None,
                 causes: Optional[List[int]] = None) -> Memory:
        if importance is None:
            importance = self._rate_importance(text)
        m = Memory(
            id=self._mem_counter, tick=tick, kind=kind, text=text,
            importance=importance, embedding=llm.embed(text),
            last_access_tick=tick, causes=causes or [],
        )
        self._mem_counter += 1
        self.memories.append(m)
        self._importance_since_reflect += importance
        return m

    def _rate_importance(self, text: str) -> float:
        """Ask the model how poignant a memory is (1 mundane .. 10 pivotal)."""
        try:
            out = llm.chat(
                system="You rate how important a memory is to a person's life on a 1-10 scale. Reply with ONLY a number.",
                user=f"Memory: {text}\nImportance (1-10):",
                temperature=0.0, max_tokens=4,
            )
            n = float("".join(c for c in out if c.isdigit() or c == ".") or "3")
            return max(1.0, min(10.0, n))
        except Exception:
            return 3.0

    def retrieve(self, query: str, tick: int, k: int = 5) -> List[Memory]:
        if not self.memories:
            return []
        q = llm.embed(query)
        scored = []
        for m in self.memories:
            recency = RECENCY_DECAY ** max(0, tick - m.last_access_tick)
            importance = m.importance / 10.0
            relevance = cosine(q, m.embedding) if m.embedding else 0.0
            score = (A_RECENCY * recency + A_IMPORTANCE * importance
                     + A_RELEVANCE * relevance)
            scored.append((score, m))
        scored.sort(key=lambda t: t[0], reverse=True)
        top = [m for _, m in scored[:k]]
        for m in top:
            m.last_access_tick = tick
        return top

    def maybe_reflect(self, tick: int) -> Optional[Memory]:
        if self._importance_since_reflect < REFLECT_THRESHOLD:
            return None
        self._importance_since_reflect = 0.0
        recent = self.memories[-15:]
        context = "\n".join(f"- {m.text}" for m in recent)
        insight = llm.chat(
            system=f"You are {self.identity()} Form ONE concise first-person insight about your life or relationships based on recent memories.",
            user=f"Recent memories:\n{context}\n\nInsight:",
            temperature=0.7, max_tokens=60,
        )
        return self.remember(tick, insight, kind="reflection", importance=7.0,
                             causes=[m.id for m in recent])

    # -- public snapshot ------------------------------------------------------
    def snapshot(self) -> dict:
        return {
            "id": self.id, "name": self.name, "age": self.age,
            "role": self.role, "traits": self.traits, "color": self.color,
            "mood": self.mood, "goal": self.goal, "location": self.location,
            "x": round(self.x, 2), "y": round(self.y, 2),
            "memory_count": len(self.memories),
            "relationships": self.relationships,
            "recent_memories": [m.to_public() for m in self.memories[-6:]],
        }