"""Keyless stub agent for workflow testing.

Connects to the orchestrator over WebSocket, registers, and returns canned results WITHOUT
calling any LLM (no OPENAI_API_KEY needed). To make the master review loop deterministic, the
"testing" task returns a conflicting framework (unittest) on its FIRST execution and an aligned
one (pytest) on retry — so the heuristic critic rejects exactly once and then accepts.

Run standalone against a manually started orchestrator:
    set "ORCHESTRATOR_URL=ws://localhost:8000"
    set "AGENT_ID=stub-1"
    set "AGENT_CAPABILITIES=python,code_generation,testing,pytest,technical-writing,general"
    python scripts/stub_agent.py
"""
import asyncio
import json
import os
import time

import websockets


class StubAgent:
    def __init__(self, orchestrator_url: str, agent_id: str, capabilities: list[str]):
        self.orchestrator_url = orchestrator_url.rstrip("/")
        self.agent_id = agent_id
        self.capabilities = capabilities
        self.ws = None
        self.runs: dict[str, int] = {}   # task_id -> execution count
        self.running = True

    def _result_for(self, task_id: str, description: str) -> dict:
        count = self.runs.get(task_id, 0) + 1
        self.runs[task_id] = count
        tid = task_id.lower()
        if "implementation" in tid:
            summary = "implemented add(a, b); tests should use pytest"
        elif "testing" in tid:
            # First attempt conflicts (unittest); retry aligns with implementation (pytest).
            summary = "wrote tests with unittest" if count == 1 else "wrote tests with pytest"
        elif "documentation" in tid:
            summary = "documented usage of add(a, b)"
        else:
            summary = f"completed: {description[:40]}"
        return {
            "success": True,
            "task_id": task_id,
            "subtasks": [{"status": "completed", "summary": summary, "changes": summary, "files": []}],
            "awaiting_handoff": False,
            "agent_id": self.agent_id,
            "usage": {
                "model_id": "stub", "model_tokens": 0, "prompt_tokens": 0,
                "completion_tokens": 0, "model_cost_usd": 0.0, "runtime_seconds": 0.0,
            },
        }

    async def send(self, payload: dict):
        await self.ws.send(json.dumps(payload))

    async def handle(self, msg: dict):
        msg_type = msg.get("type")
        if msg_type == "task_assignment":
            task_id = msg["task_id"]
            await self.send({"type": "task_accepted", "task_id": task_id, "available_slots": 4})
            await self.send({"type": "task_start", "agent_id": self.agent_id, "task_id": task_id, "timestamp": time.time()})
            await asyncio.sleep(0.2)  # simulate work
            result = self._result_for(task_id, msg.get("description", ""))
            await self.send({
                "type": "task_complete", "agent_id": self.agent_id,
                "task_id": task_id, "result": result, "timestamp": time.time(),
            })
        elif msg_type == "peer_message" and not msg.get("is_reply"):
            # Answer an inbound peer query so collaboration round-trips complete.
            await self.send({
                "type": "peer_message", "agent_id": self.agent_id,
                "target_agent_id": msg.get("from_agent_id"), "task_id": msg.get("task_id"),
                "content": f"stub guidance from {self.agent_id}",
                "correlation_id": msg.get("correlation_id"), "is_reply": True, "timestamp": time.time(),
            })

    async def heartbeat(self):
        while self.running:
            try:
                await self.send({
                    "type": "heartbeat", "agent_id": self.agent_id,
                    "timestamp": time.time(), "active_tasks": 0, "available_slots": 4,
                })
            except Exception:
                return
            await asyncio.sleep(5)

    async def run(self):
        async with websockets.connect(f"{self.orchestrator_url}/ws/{self.agent_id}") as ws:
            self.ws = ws
            await self.send({
                "type": "register", "agent_id": self.agent_id,
                "capabilities": self.capabilities, "available_slots": 4, "active_task_ids": [],
            })
            hb = asyncio.create_task(self.heartbeat())
            try:
                async for raw in ws:
                    await self.handle(json.loads(raw))
            finally:
                self.running = False
                hb.cancel()


async def _main():
    agent = StubAgent(
        os.getenv("ORCHESTRATOR_URL", "ws://localhost:8000"),
        os.getenv("AGENT_ID", "stub-1"),
        os.getenv("AGENT_CAPABILITIES", "python,code_generation,testing,pytest,technical-writing,general").split(","),
    )
    await agent.run()


if __name__ == "__main__":
    asyncio.run(_main())