"""Node that updates conversation memory.

Tracks emotion trajectory, confidence, rapport, empathy quality,
NPC openness, conversation stage, and behavioral observations each turn."""

from collections import Counter

from src.schemas.state import UserMemory, WalkState


def _determine_stage(
    turn_index: int,
    max_turns: int,
    should_end: bool,
    confidence: float,
    npc_openness: float,
) -> str:
    """Determine conversation stage based on multiple signals."""
    if should_end:
        return "closing"

    # Forced closing near end
    if turn_index >= max_turns - 1:
        return "closing"

    # Early turns are always opening
    if turn_index <= 1:
        return "opening"

    # Dynamic stage based on rapport indicators
    combined = (confidence + npc_openness) / 2

    if combined >= 0.7:
        return "deepening"  # High trust — deeper conversation
    elif combined >= 0.4:
        return "developing"  # Building rapport
    else:
        return "opening"  # Still guarded — hasn't broken through yet


def _calculate_npc_openness(
    prev_openness: float,
    trust_shift: float,
    empathy_quality: str,
    npc_mood: str,
) -> float:
    """Calculate how open James is willing to be."""
    delta = trust_shift * 0.5  # trust_shift directly influences openness

    # Empathy quality bonus/penalty
    if empathy_quality == "hit":
        delta += 0.05
    elif empathy_quality == "miss":
        delta -= 0.1  # Misses hurt more than hits help

    # If already withdrawn, harder to recover
    if npc_mood == "withdrawn":
        delta = min(delta, 0.02)

    return max(0.0, min(1.0, prev_openness + delta))


def _determine_npc_mood(openness: float, empathy_quality: str, prev_mood: str) -> str:
    """Determine James's current emotional state."""
    if empathy_quality == "miss" and prev_mood in ("guarded", "cautious"):
        return "withdrawn"

    if openness >= 0.75:
        return "open"
    elif openness >= 0.55:
        return "warming"
    elif openness >= 0.35:
        return "cautious"
    else:
        return "guarded"


def _calculate_dominant_emotion(trajectory: list[dict]) -> str:
    """Find the most frequent emotion label across all turns."""
    if not trajectory:
        return "neutral"
    labels = [e.get("label", "neutral") for e in trajectory]
    counter = Counter(labels)
    return counter.most_common(1)[0][0]


def update_memory(state: WalkState) -> dict:
    """Update memory based on agent_output — multi-dimensional tracking."""

    agent_output = state.get("agent_output", {})
    memory: UserMemory = dict(state.get("memory", {}))
    turn_index = state.get("turn_index", 0)
    scenario = state.get("scenario", {})
    max_turns = scenario.get("max_turns", 10)
    user_input = state.get("user_input", "")

    # ── Emotion Trajectory ──────────────────────────────
    emotion = agent_output.get("emotion_assessment", {})
    trajectory = list(memory.get("emotion_trajectory", []))
    trajectory.append(
        {
            "turn": turn_index,
            "valence": emotion.get("valence", 0.0),
            "arousal": emotion.get("arousal", 0.0),
            "label": emotion.get("label", "neutral"),
        }
    )
    memory["emotion_trajectory"] = trajectory
    memory["dominant_emotion"] = _calculate_dominant_emotion(trajectory)

    # ── Empathy Tracking ────────────────────────────────
    empathy_quality = agent_output.get("empathy_quality", "neutral")
    hits = memory.get("empathy_hits", 0)
    misses = memory.get("empathy_misses", 0)

    if empathy_quality == "hit":
        hits += 1
    elif empathy_quality == "miss":
        misses += 1

    memory["empathy_hits"] = hits
    memory["empathy_misses"] = misses
    total = hits + misses
    memory["empathy_ratio"] = hits / total if total > 0 else 0.5

    # ── Trust / Confidence ──────────────────────────────
    trust_shift = agent_output.get("trust_shift", 0.0)
    prev_confidence = memory.get("confidence_level", 0.5)

    # Confidence is influenced by trust_shift + empathy ratio trend
    confidence_delta = trust_shift * 0.3
    if memory["empathy_ratio"] >= 0.7:
        confidence_delta += 0.02  # Consistent empathy bonus
    elif memory["empathy_ratio"] <= 0.3:
        confidence_delta -= 0.03  # Consistent insensitivity penalty

    memory["confidence_level"] = max(0.0, min(1.0, prev_confidence + confidence_delta))

    # ── NPC Internal State ──────────────────────────────
    prev_openness = memory.get("npc_openness", 0.3)
    prev_mood = memory.get("npc_mood", "guarded")
    npc_internal_mood = agent_output.get("npc_internal_mood", prev_mood)

    memory["npc_openness"] = _calculate_npc_openness(
        prev_openness, trust_shift, empathy_quality, prev_mood
    )
    memory["npc_mood"] = _determine_npc_mood(
        memory["npc_openness"], empathy_quality, prev_mood
    )

    # ── Rapport Score (composite) ───────────────────────
    memory["rapport_score"] = (
        memory["confidence_level"] * 0.4
        + memory["npc_openness"] * 0.3
        + memory["empathy_ratio"] * 0.3
    )

    # ── Engagement Tracking ─────────────────────────────
    input_len = len(user_input.strip())
    prev_avg = memory.get("avg_response_length", 0.0)
    memory["avg_response_length"] = (
        (prev_avg * turn_index + input_len) / (turn_index + 1) if turn_index >= 0 else input_len
    )

    silence_count = memory.get("silence_count", 0)
    if input_len < 5:  # Very short response = possible disengagement
        silence_count += 1
    memory["silence_count"] = silence_count

    # ── Conversation Stage ──────────────────────────────
    should_end = agent_output.get("should_end_turn", False)
    new_stage = _determine_stage(
        turn_index, max_turns, should_end,
        memory["confidence_level"], memory["npc_openness"],
    )
    memory["current_conversation_stage"] = new_stage

    stage_history = list(memory.get("stage_history", []))
    if not stage_history or stage_history[-1] != new_stage:
        stage_history.append(new_stage)
    memory["stage_history"] = stage_history

    # ── Observations ────────────────────────────────────
    observation = agent_output.get("observation", "")
    observations = list(memory.get("key_observations", []))
    if observation:
        observations.append(observation)
    memory["key_observations"] = observations

    return {
        "memory": memory,
        "turn_index": turn_index + 1,
    }