from __future__ import annotations

from pathlib import Path

import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt
import networkx as nx

from models import VillageState
from simulation import active_rumor

GRAPH_PATH = Path("network_graph.png")


def target_choices(state: VillageState) -> list[tuple[str, str]]:
    return [(villager.name, villager.id) for villager in state.villagers.values()]


def current_rumor_text(state: VillageState) -> str:
    rumor = active_rumor(state)
    return rumor.current_version if rumor else "No rumor started yet."


def current_day_text(state: VillageState) -> str:
    return f"Day {min(state.day, 8)}"


def conversations_text(state: VillageState) -> str:
    if not state.latest_conversations:
        return "No conversations yet."
    parts = []
    for item in state.latest_conversations:
        speaker = state.villagers[item.speaker_id].name
        listener = state.villagers[item.listener_id].name
        channel = state.channels[item.channel_id].name
        parts.append(f"### {speaker} -> {listener} ({item.action}, {channel})\n{item.dialogue}")
    return "\n\n".join(parts)


def daily_summary_text(state: VillageState) -> str:
    return state.daily_summaries[-1] if state.daily_summaries else "Start a rumor, then simulate a day."


def timeline_text(state: VillageState) -> str:
    lines = []
    lines.extend(state.daily_summaries)
    for event in state.events:
        lines.append(f"Day {event.day}: {event.title} - {event.description}")
    return "\n\n".join(lines) if lines else "Timeline is waiting for the first spark."


def metrics_markdown(state: VillageState) -> str:
    rumor = active_rumor(state)
    if not rumor:
        return "No village metrics yet."
    village_trust = sum(v.trust for v in state.villagers.values()) / len(state.villagers)
    chaos = min(100, rumor.intensity * 0.7 + (1 - village_trust / 100) * 30)
    total_damage = 0
    if rumor.target_id:
        rep = state.villagers[rumor.target_id].reputation
        total_damage = sum(70 - value for value in rep.model_dump().values())
    return (
        f"**Rumor intensity:** {rumor.intensity:.1f}\n\n"
        f"**Public belief:** {rumor.public_belief:.2f}\n\n"
        f"**Village trust:** {village_trust:.1f}\n\n"
        f"**Chaos score:** {chaos:.1f}\n\n"
        f"**Total reputation damage:** {total_damage:.1f}"
    )


def evidence_markdown(state: VillageState) -> str:
    rumor = active_rumor(state)
    if not rumor:
        return "No evidence yet."
    items = [item for item in state.evidence.values() if item.rumor_id == rumor.id]
    if not items:
        return "No evidence has surfaced yet."
    lines = []
    for item in sorted(items, key=lambda evidence: evidence.reliability, reverse=True):
        discoverer = state.villagers[item.discovered_by].name if item.discovered_by in state.villagers else "Unknown"
        stance = "supports" if item.supports_rumor else "contradicts"
        visibility = "public" if item.public else "private"
        lines.append(
            f"**{item.description}**\n\n"
            f"{stance.title()} rumor | reliability {item.reliability:.2f} | {visibility} | found by {discoverer}"
        )
    return "\n\n".join(lines)


def belief_positions_markdown(state: VillageState) -> str:
    rumor = active_rumor(state)
    if not rumor:
        return "No belief positions yet."
    rows = []
    for villager in sorted(state.villagers.values(), key=lambda item: item.private_beliefs.get(rumor.id, 0), reverse=True):
        belief = villager.private_beliefs.get(rumor.id, 0.5)
        public = villager.public_positions.get(rumor.id, "neutral").replace("_", " ")
        rows.append(f"**{villager.name}:** private {belief:.2f} | public {public}")
    return "\n\n".join(rows)


def leaderboards_markdown(state: VillageState) -> str:
    villagers = list(state.villagers.values())

    def top(label: str, key):
        villager = max(villagers, key=key)
        return f"**{label}:** {villager.name} ({key(villager):.1f})"

    biggest_gossip = max(villagers, key=lambda v: v.attention + (100 - v.trust) * 0.2)
    most_damaged = min(villagers, key=lambda v: sum(v.reputation.model_dump().values()) / 5)
    lines = [
        top("Most trusted", lambda v: v.trust),
        top("Most influential", lambda v: v.influence),
        top("Most attention", lambda v: v.attention),
        f"**Most damaged reputation:** {most_damaged.name} ({sum(most_damaged.reputation.model_dump().values()) / 5:.1f})",
        f"**Biggest gossip:** {biggest_gossip.name} ({biggest_gossip.attention:.1f} attention)",
        top("Most accurate", lambda v: v.accuracy * 100),
    ]
    return "\n\n".join(lines)


def memories_markdown(state: VillageState) -> str:
    villagers = sorted(
        state.villagers.values(),
        key=lambda villager: max([abs(value) for value in villager.strategy_biases.values()] or [0]),
        reverse=True,
    )
    lines = []
    for villager in villagers:
        if not villager.memories and not villager.strategy_biases:
            continue
        biases = sorted(villager.strategy_biases.items(), key=lambda item: abs(item[1]), reverse=True)[:3]
        bias_text = ", ".join(f"{action}: {value:+.1f}" for action, value in biases) if biases else "no bias yet"
        memory_text = " / ".join(memory.summary for memory in villager.memories[-2:]) if villager.memories else "No memories yet."
        lines.append(f"**{villager.name}**\n\nBiases: {bias_text}\n\n{memory_text}")
    return "\n\n".join(lines) if lines else "No strategy memories yet."


def judge_mode_markdown(state: VillageState) -> str:
    rows = [item for item in state.trace if "action_scores" in item]
    if not rows:
        return (
            "**Judge Mode**\n\n"
            "Start a rumor and simulate a day to see the decision market. "
            "Python will choose actions through scores; the LLM only narrates dialogue, rumor mutations, and summaries."
        )
    days = sorted({row["day"] for row in rows})
    visible_days = days
    sections = [
        "**Judge Mode: Decision Market**",
        "Python controls bids, incentives, state, belief, evidence, memory, and reputation. The LLM only adds language after winners are chosen.",
    ]
    for day in visible_days:
        day_rows = [row for row in rows if row["day"] == day]
        sections.append(f"### Day {day}")
        for rank, row in enumerate(day_rows, start=1):
            scores = row.get("action_scores", {})
            top_scores = sorted(scores.items(), key=lambda item: item[1], reverse=True)[:4]
            score_text = ", ".join(f"{action} {score:.1f}" for action, score in top_scores)
            memory = row.get("memory") or {}
            economy = row.get("economy_changes") or {}
            biases = row.get("strategy_biases") or {}
            action = row.get("chosen_action", "")
            bias = biases.get(action, 0)
            reputation_delta = row.get("reputation_damage") or {}
            rep_text = ", ".join(f"{key} {value:+.2f}" for key, value in reputation_delta.items()) or "none"
            sections.append(
                f"**{rank}. {row['villager']} won `{action}`** "
                f"with score **{row['bid_score']:.1f}** against listener **{row['listener']}** via **{row['channel']}**.\n\n"
                f"Top scores: {score_text}\n\n"
                f"Learned bias on chosen action: {bias:+.1f}\n\n"
                f"Economy delta: trust {economy.get('trust_delta', 0):+.1f}, "
                f"influence {economy.get('influence_delta', 0):+.1f}, "
                f"attention {economy.get('attention_delta', 0):+.1f}, "
                f"payoff {economy.get('payoff', 0):+.1f}\n\n"
                f"Belief: {row['belief_before']:.2f} -> {row['belief_after']:.2f}; reputation delta: {rep_text}\n\n"
                f"Memory formed: {memory.get('summary', 'none')}\n\n"
                f"Reason: {row['reason']}"
            )
    return "\n\n".join(sections)


def target_reputation_markdown(state: VillageState) -> str:
    rumor = active_rumor(state)
    if not rumor or not rumor.target_id:
        return "No target selected yet."
    target = state.villagers[rumor.target_id]
    values = target.reputation.model_dump()
    return "\n".join([f"**{name.title()}:** {score:.1f}" for name, score in values.items()])


def explainability_markdown(state: VillageState) -> str:
    if not state.trace:
        return "No reputation changes to explain yet."
    rows = [item for item in state.trace if "reputation_damage" in item]
    if not rows:
        return "No reputation changes to explain yet."
    row = rows[-1]
    rumor = active_rumor(state)
    channel = row["channel"]
    target = state.villagers[rumor.target_id].name if rumor and rumor.target_id else "target"
    deltas = ", ".join(f"{k}: {v}" for k, v in row["reputation_damage"].items()) or "none"
    return (
        f"**Why did {target}'s reputation change?**\n\n"
        f"Rumor severity: {rumor.severity if rumor else 0:.1f}\n\n"
        f"Chosen action: {row['chosen_action']}\n\n"
        f"Speaker: {row['villager']}\n\n"
        f"Listener belief: {row['belief_before']} -> {row['belief_after']}\n\n"
        f"Channel: {channel}\n\n"
        f"Reputation delta: {deltas}"
    )


def network_graph_path(state: VillageState) -> str | None:
    villagers = list(state.villagers.values())
    if not villagers:
        return None
    rumor = active_rumor(state)
    graph = nx.Graph()
    for villager in villagers:
        belief = villager.private_beliefs.get(rumor.id, 0.5) if rumor else 0.5
        reputation = sum(villager.reputation.model_dump().values()) / 5
        graph.add_node(
            villager.id,
            label=villager.name.replace(" ", "\n"),
            influence=villager.influence,
            belief=belief,
            reputation=reputation,
            target=bool(rumor and rumor.target_id == villager.id),
        )
    for villager in villagers:
        for other_id, score in villager.relationships.items():
            if other_id in graph and abs(score) >= 10:
                graph.add_edge(villager.id, other_id, weight=abs(score), score=score)

    if graph.number_of_edges() == 0:
        ids = [villager.id for villager in villagers]
        graph.add_edges_from(zip(ids, ids[1:]))

    fig, ax = plt.subplots(figsize=(9, 6), dpi=140)
    fig.patch.set_facecolor("#10151f")
    ax.set_facecolor("#10151f")
    pos = nx.spring_layout(graph, seed=7, k=0.7)
    node_sizes = [420 + graph.nodes[node]["influence"] * 8 for node in graph.nodes]
    node_colors = []
    for node in graph.nodes:
        data = graph.nodes[node]
        if data["target"]:
            node_colors.append("#f97316")
        elif data["belief"] >= 0.65:
            node_colors.append("#ef4444")
        elif data["belief"] <= 0.35:
            node_colors.append("#22c55e")
        else:
            node_colors.append("#38bdf8")
    edge_colors = ["#ef4444" if graph.edges[edge]["score"] < 0 else "#94a3b8" for edge in graph.edges]
    edge_widths = [1 + graph.edges[edge]["weight"] / 25 for edge in graph.edges]

    nx.draw_networkx_edges(graph, pos, edge_color=edge_colors, width=edge_widths, alpha=0.45, ax=ax)
    nx.draw_networkx_nodes(graph, pos, node_size=node_sizes, node_color=node_colors, linewidths=1.5, edgecolors="#f8fafc", ax=ax)
    nx.draw_networkx_labels(
        graph,
        pos,
        labels={node: graph.nodes[node]["label"] for node in graph.nodes},
        font_size=7,
        font_color="#f8fafc",
        ax=ax,
    )
    ax.set_title("Trust, Belief, and Influence Network", color="#f8fafc", fontsize=13, pad=14)
    ax.text(
        0.01,
        0.01,
        "Orange target | Red believers | Green skeptics | Blue undecided",
        transform=ax.transAxes,
        color="#cbd5e1",
        fontsize=8,
    )
    ax.axis("off")
    fig.tight_layout()
    fig.savefig(GRAPH_PATH, bbox_inches="tight", facecolor=fig.get_facecolor())
    plt.close(fig)
    return str(GRAPH_PATH.resolve())


def ui_snapshot(state: VillageState) -> tuple:
    return (
        current_day_text(state),
        current_rumor_text(state),
        conversations_text(state),
        daily_summary_text(state),
        timeline_text(state),
        state.baraza_output or "The baraza will be called after Day 7.",
        metrics_markdown(state),
        evidence_markdown(state),
        belief_positions_markdown(state),
        leaderboards_markdown(state),
        memories_markdown(state),
        judge_mode_markdown(state),
        target_reputation_markdown(state),
        explainability_markdown(state),
        network_graph_path(state),
    )