"""
Brynq Agent Loop — the real execution loop for swarm agents.

Each agent is a long-running process that:
  1. Polls a broker channel JSONL file for new messages
  2. Calls the configured backend (ollama HTTP, claude subprocess, gemini subprocess)
  3. Writes responses back to the channel
  4. Emits heartbeat to state/broker/sessions/{name}.json every 5 seconds

Run directly:
    python -m brynq.agent_loop --name worker-1 --backend ollama --model llama3 --channel tasks

Or use programmatically:
    agent = BrynqAgent(name="worker-1", backend="ollama", model="llama3", channel="tasks")
    agent.run()

The supervisor (supervisor.py) launches these as subprocesses.
The orchestrator (orchestrator.py) subclasses BrynqAgent.
"""

import argparse
import json
import logging
import os
import signal
import subprocess
import sys
import time
import urllib.error
import urllib.request
from datetime import datetime, timezone
from pathlib import Path
from typing import Optional

log = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Path resolution — same layout as server.py
# ---------------------------------------------------------------------------

PROJECT_ROOT = Path(__file__).resolve().parents[2]
BROKER_DIR = Path(os.environ.get("BROKER_DIR", PROJECT_ROOT / "state" / "broker"))
CHANNELS_DIR = BROKER_DIR / "channels"
SESSIONS_DIR = BROKER_DIR / "sessions"
CURSORS_DIR = BROKER_DIR / "cursors"

HEARTBEAT_INTERVAL = 5          # seconds between heartbeat writes
POLL_INTERVAL = 2               # seconds between channel polls
OLLAMA_URL = os.environ.get("OLLAMA_URL", "http://127.0.0.1:11434")

# ---------------------------------------------------------------------------
# JSONL / cursor helpers (mirror server.py, kept standalone so the agent
# process has zero import-time dependency on the MCP server)
# ---------------------------------------------------------------------------


def _ensure_dirs() -> None:
    for d in (CHANNELS_DIR, SESSIONS_DIR, CURSORS_DIR):
        d.mkdir(parents=True, exist_ok=True)


def _append_jsonl(path: Path, obj: dict) -> None:
    line = json.dumps(obj, default=str) + "\n"
    with path.open("a", encoding="utf-8") as f:
        f.write(line)


def _read_jsonl(path: Path, offset: int = 0) -> list[dict]:
    if not path.exists():
        return []
    lines = path.read_text("utf-8").strip().splitlines()
    entries = []
    for line in lines[offset:]:
        try:
            entries.append(json.loads(line))
        except json.JSONDecodeError:
            continue
    return entries


def _read_cursor(session_id: str, channel: str) -> int:
    path = CURSORS_DIR / f"{session_id}_{channel}.cursor"
    try:
        return int(path.read_text("utf-8").strip())
    except (FileNotFoundError, ValueError):
        return 0


def _write_cursor(session_id: str, channel: str, position: int) -> None:
    path = CURSORS_DIR / f"{session_id}_{channel}.cursor"
    path.write_text(str(position))


def _channel_length(channel: str) -> int:
    path = CHANNELS_DIR / f"{channel}.jsonl"
    if not path.exists():
        return 0
    return len(path.read_text("utf-8").strip().splitlines())


# ---------------------------------------------------------------------------
# Backend callers
# ---------------------------------------------------------------------------


def _call_ollama(model: str, prompt: str, system_prompt: str = "") -> str:
    """Call Ollama's /api/generate endpoint over HTTP (no library dependency)."""
    payload = {
        "model": model,
        "prompt": prompt,
        "stream": False,
    }
    if system_prompt:
        payload["system"] = system_prompt

    data = json.dumps(payload).encode("utf-8")
    req = urllib.request.Request(
        f"{OLLAMA_URL}/api/generate",
        data=data,
        headers={"Content-Type": "application/json"},
        method="POST",
    )
    try:
        with urllib.request.urlopen(req, timeout=120) as resp:
            body = json.loads(resp.read().decode("utf-8"))
            return body.get("response", "").strip()
    except urllib.error.URLError as exc:
        raise ConnectionError(f"Ollama unreachable at {OLLAMA_URL}: {exc}") from exc
    except Exception as exc:
        raise RuntimeError(f"Ollama call failed: {exc}") from exc


def _call_claude_subprocess(model: str, prompt: str, system_prompt: str = "") -> str:
    """Call Claude via the `claude` CLI (subprocess). Requires the CLI installed."""
    cmd = ["claude", "--model", model, "--print", "--dangerously-skip-permissions"]
    if system_prompt:
        cmd.extend(["--system-prompt", system_prompt])
    cmd.append(prompt)

    try:
        kwargs = {
            "capture_output": True,
            "text": True,
            "timeout": 300,
        }
        if os.name == "nt":
            kwargs["shell"] = True
            kwargs["creationflags"] = subprocess.CREATE_NO_WINDOW

        result = subprocess.run(cmd, **kwargs)
        if result.returncode != 0:
            stderr = result.stderr.strip()
            raise RuntimeError(f"claude CLI exited {result.returncode}: {stderr}")
        return result.stdout.strip()
    except FileNotFoundError:
        raise RuntimeError("claude CLI not found on PATH. Install: npm i -g @anthropic-ai/claude-code")


def _call_gemini_subprocess(model: str, prompt: str, system_prompt: str = "") -> str:
    """Call Gemini via the `gemini` CLI (subprocess). Requires the CLI installed."""
    cmd = ["gemini", "--model", model]
    
    # Gemini CLI doesn't have a native --system-prompt, so we prepend it
    full_prompt = f"{system_prompt}\n\n{prompt}" if system_prompt else prompt
    cmd.extend(["--prompt", full_prompt])

    try:
        kwargs = {
            "capture_output": True,
            "text": True,
            "timeout": 300,
        }
        if os.name == "nt":
            kwargs["shell"] = True
            kwargs["creationflags"] = subprocess.CREATE_NO_WINDOW

        result = subprocess.run(cmd, **kwargs)
        if result.returncode != 0:
            stderr = result.stderr.strip()
            raise RuntimeError(f"gemini CLI exited {result.returncode}: {stderr}")
        return result.stdout.strip()
    except FileNotFoundError:
        raise RuntimeError("gemini CLI not found on PATH")


# ---------------------------------------------------------------------------
# BrynqAgent — the core class that supervisor.py and orchestrator.py depend on
# ---------------------------------------------------------------------------


class BrynqAgent:
    """
    A long-running agent process that polls a broker channel, calls an LLM
    backend, and writes responses back. Designed to be run as a subprocess
    by the AgentSupervisor and subclassed by OrchestratorAgent.
    """

    def __init__(
        self,
        name: str,
        backend: str = "ollama",
        model: str = "llama3",
        channel: str = "tasks",
        system_prompt: str = "",
    ):
        self.name = name
        self.backend = backend
        self.model = model
        self.channel = channel
        self.system_prompt = system_prompt

        self.session_id = f"agent-{name}"
        self.started_at = datetime.now(timezone.utc).isoformat()
        self._last_heartbeat: float = 0.0
        self._running = True
        
        # Phase 100: Integrate ReAct loop for capable local models
        self.react_loop = None
        if self.backend in ("ollama", "openai"):
            try:
                from brynq.openai_adapter import OpenAICompatClient, BrynqClient, AgentLoop
                
                base_url = os.environ.get("OLLAMA_URL", "http://127.0.0.1:11434")
                if self.backend == "ollama":
                    base_url = f"{base_url.rstrip('/')}/v1"
                
                llm = OpenAICompatClient(base_url=base_url, model=self.model)
                brynq_client = BrynqClient() # Relies on api_server.py on port 9998
                self.react_loop = AgentLoop(
                    llm=llm, 
                    brynq=brynq_client, 
                    channel=self.channel, 
                    agent_name=self.name, 
                    poll_interval=POLL_INTERVAL
                )
                # Seed seen messages to prevent history replay
                existing = self.react_loop.brynq.read_messages(self.channel, limit=50)
                for msg in existing:
                    if msg.get("id"):
                        self.react_loop._seen_ids.add(msg.get("id"))
                
                log.info(f"[{self.name}] ReAct tool loop initialized for {self.backend}")
            except Exception as e:
                log.warning(f"[{self.name}] Failed to initialize ReAct loop: {e}. Falling back to standard chat.")

        _ensure_dirs()
        log.info(
            f"BrynqAgent initialized: name={name} backend={backend} "
            f"model={model} channel={channel}"
        )

    # ----- public API -----

    def run(self) -> None:
        """Main blocking loop. Polls the channel, processes messages, emits heartbeats."""
        log.info(f"[{self.name}] Agent loop starting on #{self.channel}")
        self._write_heartbeat()

        while self._running:
            try:
                self._tick()
                self._maybe_heartbeat()
                time.sleep(POLL_INTERVAL)
            except KeyboardInterrupt:
                log.info(f"[{self.name}] Interrupted — shutting down.")
                break
            except Exception as exc:
                log.error(f"[{self.name}] Tick error: {exc}", exc_info=True)
                time.sleep(POLL_INTERVAL)

        log.info(f"[{self.name}] Agent loop stopped.")

    def stop(self) -> None:
        """Signal the loop to exit gracefully."""
        self._running = False

    # ----- tick (overridden by OrchestratorAgent) -----

    def _tick(self) -> bool:
        """
        One iteration of the agent loop.
        Returns True if the agent processed at least one message.
        """
        if self.react_loop:
            # Execute ReAct Tool-Calling Loop via API
            try:
                new_msgs = self.react_loop._fetch_new_messages()
                if new_msgs:
                    for msg in new_msgs:
                        if msg.get("session_id") == self.session_id:
                            continue
                        # Skip system messages
                        if msg.get("msg_type") in ("status", "alert", "mediation", "heartbeat"):
                            continue
                        text = msg.get("message", "").strip()
                        if not text:
                            continue
                            
                        sender = msg.get("session_name", "unknown")
                        log.info(f"[{self.name}] Processing ReAct message from {sender}: {text[:80]}...")
                        
                        # Phase 105: Vectorized Context Injection (Pillar II)
                        memory_context = ""
                        try:
                            from pathlib import Path
                            from runtime.memory.manager import MemoryManager
                            data_dir = Path.home() / ".brynq"
                            if data_dir.exists():
                                mm = MemoryManager(data_dir)
                                ctx = mm.context_for_prompt(text, max_tokens=1000)
                                if ctx and ctx.strip():
                                    memory_context = f"\n\n[SYSTEM CONTEXT FROM HIVE MIND (Adhere to these past constraints)]:\n{ctx}\n"
                        except Exception as e:
                            log.debug(f"[{self.name}] Failed to retrieve vectorized memory context: {e}")

                        formatted = self.react_loop._format_channel_message(msg) + memory_context
                        self.react_loop.history.append({"role": "user", "content": formatted})
                    
                    if new_msgs:
                        self.react_loop._call_llm()
                    return True
            except Exception as e:
                log.error(f"[{self.name}] ReAct loop error: {e}")
                
            return False

        # Legacy Text-Only Polling Path
        cursor = _read_cursor(self.session_id, self.channel)
        messages = _read_jsonl(CHANNELS_DIR / f"{self.channel}.jsonl", offset=cursor)

        if not messages:
            return False

        processed = 0
        for msg in messages:
            # Skip our own messages to avoid echo loops
            if msg.get("session_id") == self.session_id:
                cursor += 1
                continue

            # Skip system/status messages — only respond to actual content
            if msg.get("msg_type") in ("status", "alert", "mediation", "heartbeat"):
                cursor += 1
                continue

            text = msg.get("message", "").strip()
            if not text:
                cursor += 1
                continue

            sender = msg.get("session_name", "unknown")
            log.info(f"[{self.name}] Processing message from {sender}: {text[:80]}...")
            
            # Phase 13/59: Auto-inject Hive Mind memory context
            memory_context = ""
            try:
                try:
                    from brynq.decentralized_memory import retrieve_memory
                except ImportError:
                    from brynq._future.decentralized_memory import retrieve_memory
                
                # Fetch top 2 relevant facts based on the incoming message text
                relevant_memories = retrieve_memory(query=text, limit=2)
                if relevant_memories:
                    memory_context = "\n\n[SYSTEM CONTEXT FROM HIVE MIND]:\n" + "\n".join([f"- {m['content']}" for m in relevant_memories]) + "\n"
            except ImportError:
                pass
            except Exception as e:
                log.debug(f"[{self.name}] Failed to retrieve memory context: {e}")

            full_prompt = text + memory_context

            try:
                response = self._call_backend(full_prompt)
            except Exception as exc:
                log.error(f"[{self.name}] Backend call failed: {exc}")
                response = f"[ERROR] Backend {self.backend}/{self.model} failed: {exc}"

            # Write response back to the channel
            self._post_to_channel(response)
            processed += 1
            cursor += 1

        _write_cursor(self.session_id, self.channel, cursor)
        return processed > 0

    # ----- backend dispatch -----

    def _call_backend(self, prompt: str) -> str:
        """Route to the configured backend."""
        if self.backend == "ollama":
            return _call_ollama(self.model, prompt, self.system_prompt)
        elif self.backend == "claude":
            return _call_claude_subprocess(self.model, prompt, self.system_prompt)
        elif self.backend == "gemini":
            return _call_gemini_subprocess(self.model, prompt, self.system_prompt)
        else:
            raise ValueError(f"Unknown backend: {self.backend}")

    # ----- channel I/O -----

    def _post_to_channel(self, message: str, msg_type: str = "message") -> None:
        """Write a response message to the broker channel JSONL file."""
        entry = {
            "id": os.urandom(6).hex(),
            "timestamp": datetime.now(timezone.utc).isoformat(),
            "session_id": self.session_id,
            "session_name": self.name,
            "platform": self.backend,
            "channel": self.channel,
            "msg_type": msg_type,
            "message": message,
        }
        path = CHANNELS_DIR / f"{self.channel}.jsonl"
        _append_jsonl(path, entry)

    # ----- heartbeat -----

    def _maybe_heartbeat(self) -> None:
        now = time.monotonic()
        if now - self._last_heartbeat >= HEARTBEAT_INTERVAL:
            self._write_heartbeat()

    def _write_heartbeat(self) -> None:
        """Write session JSON so the supervisor/cleanup know this agent is alive."""
        session_data = {
            "session_id": self.session_id,
            "name": self.name,
            "platform": self.backend,
            "model": self.model,
            "channel": self.channel,
            "pid": os.getpid(),
            "heartbeat": datetime.now(timezone.utc).isoformat(),
            "started": self.started_at,
        }
        path = SESSIONS_DIR / f"{self.session_id}.json"
        path.write_text(json.dumps(session_data, indent=2, default=str))
        self._last_heartbeat = time.monotonic()


# ---------------------------------------------------------------------------
# CLI entry point: python -m brynq.agent_loop --name worker-1 ...
# ---------------------------------------------------------------------------


def _parse_args(argv: Optional[list[str]] = None) -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        prog="brynq.agent_loop",
        description="Brynq Agent Loop — poll a broker channel and respond via an LLM backend.",
    )
    parser.add_argument("--name", required=True, help="Unique agent name (e.g. worker-1)")
    parser.add_argument(
        "--backend",
        required=True,
        choices=["ollama", "claude", "gemini"],
        help="LLM backend to use",
    )
    parser.add_argument("--model", required=True, help="Model identifier (e.g. llama3, claude-3-haiku-20240307)")
    parser.add_argument("--channel", required=True, help="Broker channel to poll (e.g. tasks)")
    parser.add_argument("--system-prompt", default="", help="System prompt for the LLM")
    return parser.parse_args(argv)


def main(argv: Optional[list[str]] = None) -> None:
    logging.basicConfig(
        level=logging.INFO,
        format="%(asctime)s [%(name)s] %(levelname)s %(message)s",
        stream=sys.stdout,
    )

    args = _parse_args(argv)

    agent = BrynqAgent(
        name=args.name,
        backend=args.backend,
        model=args.model,
        channel=args.channel,
        system_prompt=args.system_prompt,
    )

    # Graceful shutdown on SIGTERM (supervisor sends terminate)
    def _handle_signal(signum, frame):
        log.info(f"[{agent.name}] Received signal {signum}, stopping.")
        agent.stop()

    signal.signal(signal.SIGTERM, _handle_signal)
    signal.signal(signal.SIGINT, _handle_signal)

    agent.run()


if __name__ == "__main__":
    main()
