#!/usr/bin/env python3
"""
Hermes Swarm Chat — reusable multi-agent discussion over Redis Pub/Sub.

Usage:
    python3 swarm_chat.py \
        --topic "Design a Python function that caches HTTP responses" \
        --agents architect critic coder \
        --rounds 3 \
        --room myproject:swarm:demo

Each agent is a separate Hermes profile. The dispatcher:
  1. seeds a Redis room with the topic,
  2. asks every agent in turn to read room history and respond,
  3. limits discussion to N rounds,
  4. asks the first agent (synthesizer) to write the final answer.

Environment:
    REDIS_URL           Redis URI (default: redis://localhost:6379/0)
    HERMES_CMD          Hermes binary (default: hermes)
    HERMES_WORKDIR      Working directory passed to hermes chat (default: current dir)
    SWARM_OUTPUT_DIR    Where to save transcripts/final answer (default: ./swarm_outputs)
"""

import argparse
import json
import os
import re
import subprocess
import sys
import time
import uuid
from dataclasses import dataclass, field
from pathlib import Path
from typing import List, Optional

import redis


DEFAULT_ROOM_PREFIX = "hermes:swarm"
REDIS_URL = os.getenv("REDIS_URL", "redis://localhost:6379/0")
HERMES_CMD = os.getenv("HERMES_CMD", "hermes")
HERMES_WORKDIR = os.getenv("HERMES_WORKDIR", os.getcwd())
OUTPUT_DIR = Path(os.getenv("SWARM_OUTPUT_DIR", "./swarm_outputs")).resolve()


@dataclass
class SwarmMessage:
    id: str
    round: int
    agent: str
    role: str
    content: str
    timestamp: float = field(default_factory=time.time)
    reply_to: Optional[str] = None

    def to_json(self) -> str:
        return json.dumps({
            "id": self.id,
            "round": self.round,
            "agent": self.agent,
            "role": self.role,
            "content": self.content,
            "timestamp": self.timestamp,
            "reply_to": self.reply_to,
        }, ensure_ascii=False, default=str)

    @classmethod
    def from_dict(cls, d: dict) -> "SwarmMessage":
        return cls(
            id=d["id"],
            round=d["round"],
            agent=d["agent"],
            role=d["role"],
            content=d["content"],
            timestamp=d.get("timestamp", time.time()),
            reply_to=d.get("reply_to"),
        )


@dataclass
class AgentDef:
    name: str           # Hermes profile name, e.g. swarm-architect
    role: str           # human-readable role, e.g. architect
    prompt: str         # system role instruction
    model: str = ""    # informational only


class RedisRoom:
    def __init__(self, room: str, redis_url: str = REDIS_URL):
        self.room = room
        self.control_key = f"{room}:control"
        self.r = redis.from_url(redis_url, decode_responses=True)
        self.pubsub = self.r.pubsub()
        self.pubsub.subscribe(room)

    def post(self, msg: SwarmMessage) -> None:
        self.r.publish(self.room, msg.to_json())
        # Also keep durable history sorted by timestamp in a Redis sorted set.
        self.r.zadd(f"{self.room}:history", {msg.to_json(): msg.timestamp})

    def history(self, since: float = 0) -> List[SwarmMessage]:
        items = self.r.zrangebyscore(f"{self.room}:history", since, "+inf")
        messages = []
        for raw in items:
            try:
                messages.append(SwarmMessage.from_dict(json.loads(raw)))
            except Exception:
                continue
        return messages

    def clear(self) -> None:
        self.r.delete(f"{self.room}:history")
        self.r.delete(self.control_key)

    def signal_stop(self, reason: str = "rounds_exhausted") -> None:
        self.r.set(self.control_key, json.dumps({"stop": True, "reason": reason}))

    def should_stop(self) -> bool:
        raw = self.r.get(self.control_key)
        if isinstance(raw, str) and raw:
            try:
                return json.loads(raw).get("stop", False)
            except Exception:
                pass
        return False


class HermesAgentRunner:
    """Runs a single Hermes agent via CLI and captures the text reply."""

    def __init__(self, profile: str, timeout: int = 120, workdir: str = HERMES_WORKDIR, hermes_cmd: str = HERMES_CMD):
        self.profile = profile
        self.timeout = timeout
        self.workdir = workdir
        self.hermes_cmd = hermes_cmd

    def ask(self, prompt: str) -> str:
        cmd = [
            self.hermes_cmd, "-p", self.profile, "chat",
            "-q", prompt,
            "--source", "swarm",
            "-Q",           # quiet: suppress banner/spinner
        ]
        env = os.environ.copy()
        # Force working directory so agents see consistent paths.
        env["HERMES_CWD"] = self.workdir
        try:
            result = subprocess.run(
                cmd,
                cwd=self.workdir,
                env=env,
                capture_output=True,
                text=True,
                timeout=self.timeout,
            )
            output = result.stdout + result.stderr
            # Strip ANSI escape codes.
            output = re.sub(r"\x1b\[[0-9;]*m", "", output)
            return self._extract_answer(output)
        except subprocess.TimeoutExpired:
            return f"[timeout: agent {self.profile} did not respond within {self.timeout}s]"
        except Exception as e:
            return f"[error running agent {self.profile}: {e}]"

    @staticmethod
    def _extract_answer(output: str) -> str:
        lines = output.splitlines()
        drop_prefixes = (
            "Thinking", "Using model", "▔", "─", "┌", "└", "│",
            "You:", "Hermes:", "◆", "⚕", "Running", "Tool result",
        )
        cleaned = []
        for line in lines:
            stripped = line.strip()
            if not stripped:
                continue
            stripped = re.sub(r"\x1b\[[0-9;]*m", "", stripped)
            if any(stripped.startswith(p) for p in drop_prefixes):
                continue
            if "session_id:" in stripped or "⚠ tirith" in stripped:
                continue
            cleaned.append(line)
        while cleaned and not cleaned[-1].strip():
            cleaned.pop()
        text = "\n".join(cleaned).strip()
        if not text:
            text = re.sub(r"\x1b\[[0-9;]*m", "", output).strip()
        return text


class SwarmChat:
    def __init__(self, room: RedisRoom, agents: List[AgentDef], topic: str, rounds: int = 3):
        self.room = room
        self.agents = agents
        self.topic = topic
        self.rounds = rounds
        self.runner = HermesAgentRunner(profile="")  # profile set per turn

    def _build_prompt(self, agent: AgentDef, round_no: int, history: List[SwarmMessage]) -> str:
        transcript = []
        for m in history:
            transcript.append(f"[Round {m.round} | {m.agent} ({m.role})]\n{m.content}\n")
        transcript_text = "\n".join(transcript) if transcript else "(no messages yet)"

        return (
            f"You are participating in a multi-agent engineering discussion.\n\n"
            f"YOUR ROLE: {agent.role}\n"
            f"{agent.prompt}\n\n"
            f"DISCUSSION TOPIC:\n{self.topic}\n\n"
            f"CURRENT ROUND: {round_no} of {self.rounds}\n\n"
            f"PREVIOUS MESSAGES IN THIS ROOM:\n{transcript_text}\n\n"
            f"TASK:\n"
            f"Read the previous messages, then contribute your perspective as the {agent.role}. "
            f"Be concise (2-5 sentences). Do not repeat what others already said. "
            f"If you disagree, explain why and propose a better approach. "
            f"If the discussion has already converged, say so and add any final detail."
        )

    def run(self) -> List[SwarmMessage]:
        seed = SwarmMessage(
            id=f"seed-{uuid.uuid4().hex[:8]}",
            round=0,
            agent="dispatcher",
            role="moderator",
            content=f"Discussion topic: {self.topic}. Agents: {[a.role for a in self.agents]}. Maximum rounds: {self.rounds}.",
        )
        self.room.post(seed)

        for r in range(1, self.rounds + 1):
            if self.room.should_stop():
                print(f"[dispatcher] stop signal detected before round {r}")
                break
            print(f"\n=== Round {r}/{self.rounds} ===")
            for agent in self.agents:
                if self.room.should_stop():
                    break
                history = self.room.history(since=0)
                prompt = self._build_prompt(agent, r, history)
                self.runner.profile = agent.name
                print(f"[ask {agent.name} ({agent.role})] ...", end="", flush=True)
                answer = self.runner.ask(prompt)
                print(f" ({len(answer)} chars)")
                msg = SwarmMessage(
                    id=f"r{r}-{agent.role}-{uuid.uuid4().hex[:6]}",
                    round=r,
                    agent=agent.name,
                    role=agent.role,
                    content=answer,
                )
                self.room.post(msg)

        self.room.signal_stop("rounds_exhausted")
        return self.room.history(since=0)

    def synthesize(self) -> str:
        history = self.room.history(since=0)
        transcript = []
        for m in history:
            if m.agent == "dispatcher":
                continue
            transcript.append(f"[{m.agent} ({m.role}), round {m.round}]\n{m.content}\n")
        transcript_text = "\n".join(transcript)

        synthesizer = self.agents[0]
        prompt = (
            f"You are the {synthesizer.role}. The multi-agent discussion below has finished.\n\n"
            f"ORIGINAL TOPIC:\n{self.topic}\n\n"
            f"FULL TRANSCRIPT:\n{transcript_text}\n\n"
            f"TASK:\n"
            f"Write a concise final answer (1-3 paragraphs) that incorporates the best ideas, resolves disagreements, "
            f"and gives the user a concrete, actionable result. Do not introduce fictional tools or APIs."
        )
        self.runner.profile = synthesizer.name
        print(f"\n[synthesize via {synthesizer.name}] ...", end="", flush=True)
        answer = self.runner.ask(prompt)
        print(f" ({len(answer)} chars)")
        return answer


def main():
    parser = argparse.ArgumentParser(description="Run a synchronous multi-agent discussion via Redis Pub/Sub.")
    parser.add_argument("--topic", required=True, help="The discussion topic / task.")
    parser.add_argument("--agents", nargs="+", required=True,
                        help="Hermes profile names to use as agents, e.g. swarm-architect swarm-critic swarm-coder")
    parser.add_argument("--roles", nargs="+", default=None,
                        help="Optional role labels matching --agents. If omitted, derived from profile names.")
    parser.add_argument("--rounds", type=int, default=3, help="Number of discussion rounds (default 3).")
    parser.add_argument("--room", default=None, help="Redis room key. Default: hermes:swarm:<random>.")
    parser.add_argument("--no-synth", action="store_true", help="Skip final synthesis.")
    parser.add_argument("--clear", action="store_true", help="Clear room history before starting.")
    parser.add_argument("--timeout", type=int, default=120, help="Timeout per agent call in seconds.")
    parser.add_argument("--output-dir", type=Path, default=None,
                        help="Directory for transcripts/final output. Default: SWARM_OUTPUT_DIR or ./swarm_outputs")
    args = parser.parse_args()

    if len(args.agents) < 2:
        print("Need at least 2 agents for a discussion.", file=sys.stderr)
        sys.exit(1)

    room_name = args.room or f"{DEFAULT_ROOM_PREFIX}:{uuid.uuid4().hex[:8]}"
    room = RedisRoom(room_name)
    if args.clear:
        room.clear()

    roles = args.roles or [a.replace("swarm-", "") for a in args.agents]
    if len(roles) != len(args.agents):
        print("--roles must match --agents length.", file=sys.stderr)
        sys.exit(1)

    role_prompts = {
        "architect": "Look at the big picture, propose structure and design trade-offs.",
        "critic": "Challenge assumptions, find flaws, edge cases, and risks in others' proposals.",
        "coder": "Turn ideas into concrete code or commands. Prefer working examples.",
        "tester": "Identify what needs to be tested, write test cases, and verify edge cases.",
        "reviewer": "Check the proposed solution for clarity, completeness, and maintainability.",
        "planner": "Break the topic into actionable steps and suggest an execution order.",
        "default": "Contribute relevant expertise and help the group converge on a good answer.",
    }

    agents = []
    for profile, role in zip(args.agents, roles):
        agents.append(AgentDef(
            name=profile,
            role=role,
            prompt=role_prompts.get(role, role_prompts["default"]),
        ))

    out_dir = args.output_dir or OUTPUT_DIR
    out_dir.mkdir(parents=True, exist_ok=True)

    print(f"Room: {room_name}")
    print(f"Agents: {[a.name + '(' + a.role + ')' for a in agents]}")
    print(f"Rounds: {args.rounds}")
    print(f"Output: {out_dir}")

    swarm = SwarmChat(room, agents, args.topic, args.rounds)
    swarm.runner.timeout = args.timeout
    history = swarm.run()

    print(f"\nTotal messages in room: {len(history)}")

    if not args.no_synth:
        final = swarm.synthesize()
        print("\n=== FINAL SYNTHESIS ===")
        print(final)
        out_file = out_dir / f"{room_name.replace(':', '_')}_final.md"
        out_file.write_text(f"# {args.topic}\n\n{final}\n", encoding="utf-8")
        print(f"\nSaved synthesis to {out_file}")

    tx_file = out_dir / f"{room_name.replace(':', '_')}_transcript.md"
    lines = [f"# {args.topic}\n\n## Room: {room_name}\n\n"]
    for m in history:
        lines.append(f"### Round {m.round} — {m.agent} ({m.role})\n\n{m.content}\n\n")
    tx_file.write_text("".join(lines), encoding="utf-8")
    print(f"Saved transcript to {tx_file}")


if __name__ == "__main__":
    main()