"""
OpenAI Responses API — WebSocket Agentic Loop (Production-Grade)
Requires: openai>=1.78.0, websockets>=12.0, tenacity>=8.0
Install:  pip install "openai>=1.78.0" "tenacity>=8.0"

Improvements over the original:
  - Variables properly initialized before use (no NameError on turn > 0)
  - All tool calls per turn are executed (not just the first one)
  - JSON argument parsing wrapped in try/except
  - Network errors retried with exponential backoff (tenacity)
  - Per-receive timeout to detect stalled connections
  - Explicit MaxTurnsExceededError instead of silent empty return
  - Fixed list-comprehension variable shadowing
  - Structured logging instead of raw print()
  - Type hints throughout
"""

from __future__ import annotations

import asyncio
import json
import logging
import os
from dataclasses import dataclass, field
from typing import Any

from openai import AsyncOpenAI
from tenacity import (
    retry,
    retry_if_exception_type,
    stop_after_attempt,
    wait_exponential,
)

# ── Logging ────────────────────────────────────────────────────────────────────

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(name)s — %(message)s",
)
logger = logging.getLogger("ws_agent")


# ── Custom exceptions ──────────────────────────────────────────────────────────

class MaxTurnsExceededError(RuntimeError):
    """Raised when the agent loop hits the turn cap without a final answer."""

class ToolExecutionError(RuntimeError):
    """Raised when a tool call cannot be dispatched."""


# ── Data structures ────────────────────────────────────────────────────────────

@dataclass
class ToolCall:
    call_id: str
    name: str
    arguments: dict[str, Any]


@dataclass
class TurnResult:
    """Outcome of a single response.done event."""
    response_id: str
    tool_calls: list[ToolCall] = field(default_factory=list)
    final_text: str = ""

    @property
    def has_tool_calls(self) -> bool:
        return bool(self.tool_calls)


# ── Tool registry ──────────────────────────────────────────────────────────────

TOOLS: list[dict[str, Any]] = [
    {
        "type": "function",
        "name": "read_file",
        "description": "Read the contents of a file by path",
        "parameters": {
            "type": "object",
            "properties": {"path": {"type": "string"}},
            "required": ["path"],
        },
    },
    {
        "type": "function",
        "name": "write_file",
        "description": "Write content to a file at the given path",
        "parameters": {
            "type": "object",
            "properties": {
                "path": {"type": "string"},
                "content": {"type": "string"},
            },
            "required": ["path", "content"],
        },
    },
]


def execute_tool(name: str, args: dict[str, Any]) -> str:
    """
    Dispatch a tool call and return its output as a string.
    Raises ToolExecutionError for unknown tools.
    All I/O errors are caught and returned as error strings
    so the model can self-correct rather than crashing.
    """
    if name == "read_file":
        try:
            with open(args["path"]) as fh:
                content = fh.read()
            logger.info("read_file OK — %s (%d bytes)", args["path"], len(content))
            return content
        except FileNotFoundError:
            return f"Error: file '{args['path']}' not found"
        except OSError as exc:
            return f"Error reading '{args['path']}': {exc}"

    if name == "write_file":
        try:
            with open(args["path"], "w") as fh:
                fh.write(args["content"])
            msg = f"Written {len(args['content'])} bytes to {args['path']}"
            logger.info("write_file OK — %s", msg)
            return msg
        except OSError as exc:
            return f"Error writing '{args['path']}': {exc}"

    raise ToolExecutionError(f"Unknown tool: '{name}'")


# ── Event parsing ──────────────────────────────────────────────────────────────

def _parse_tool_calls(output_items: list[dict[str, Any]]) -> list[ToolCall]:
    """
    Extract and parse every function_call item from a response output.
    Malformed JSON arguments are caught per-call so one bad call
    does not discard the others.
    """
    calls: list[ToolCall] = []
    for item in output_items:
        if item.get("type") != "function_call":
            continue
        raw_args = item.get("arguments", "{}")
        try:
            parsed_args = json.loads(raw_args)
        except json.JSONDecodeError as exc:
            logger.warning(
                "Could not parse arguments for tool '%s': %s — raw: %r",
                item.get("name"), exc, raw_args,
            )
            parsed_args = {}
        calls.append(ToolCall(
            call_id=item["call_id"],
            name=item["name"],
            arguments=parsed_args,
        ))
    return calls


def _parse_final_text(output_items: list[dict[str, Any]]) -> str:
    """
    Collect all output_text content from message items.
    Uses distinct variable names to avoid the shadowing bug in the original.
    """
    parts: list[str] = []
    for msg_item in output_items:
        if msg_item.get("type") != "message":
            continue
        for content_block in msg_item.get("content", []):
            if content_block.get("type") == "output_text":
                parts.append(content_block["text"])
    return "\n".join(parts)


def _parse_turn_result(response: dict[str, Any]) -> TurnResult:
    output = response.get("output", [])
    tool_calls = _parse_tool_calls(output)
    final_text = "" if tool_calls else _parse_final_text(output)
    return TurnResult(
        response_id=response["id"],
        tool_calls=tool_calls,
        final_text=final_text,
    )


# ── WebSocket receive loop ─────────────────────────────────────────────────────

async def _receive_until_done(
    ws: Any,
    event_timeout: float,
) -> TurnResult:
    """
    Drain WebSocket events until response.done, with a per-event timeout.
    Raises asyncio.TimeoutError if no event arrives within event_timeout seconds.
    """
    async for raw in ws:
        event = json.loads(raw)
        etype = event.get("type", "")

        if etype == "response.output_text.delta":
            print(event.get("delta", ""), end="", flush=True)

        elif etype == "error":
            raise RuntimeError(
                f"Server error {event.get('code')}: {event.get('message')}"
            )

        elif etype == "response.done":
            print()
            return _parse_turn_result(event["response"])

    raise RuntimeError("WebSocket closed before response.done was received")


# ── Core agent loop ────────────────────────────────────────────────────────────

@retry(
    retry=retry_if_exception_type((OSError, asyncio.TimeoutError)),
    stop=stop_after_attempt(3),
    wait=wait_exponential(multiplier=1, min=2, max=10),
    reraise=True,
)
async def run_websocket_agent(
    user_message: str,
    model: str = "gpt-4o",
    max_turns: int = 30,
    event_timeout: float = 60.0,
) -> str:
    """
    Run a full agentic loop over a single persistent WebSocket connection.

    Args:
        user_message:   The user's initial prompt.
        model:          OpenAI model identifier.
        max_turns:      Hard cap on tool-call rounds before raising.
        event_timeout:  Seconds to wait for any single WebSocket event.

    Returns:
        The model's final text answer.

    Raises:
        MaxTurnsExceededError: If the loop reaches max_turns without finishing.
        RuntimeError:           On unrecoverable server or network errors.
    """
    client = AsyncOpenAI(api_key=os.environ["OPENAI_API_KEY"])

    # Initialised here so they are always defined before the continuation branch
    previous_response_id: str | None = None
    pending_tool_results: list[dict[str, Any]] = []

    async with client.responses.websocket() as ws:
        for turn in range(max_turns):
            logger.info("Turn %d/%d", turn + 1, max_turns)

            # ── Build request payload ──────────────────────────────────────────
            if turn == 0:
                payload: dict[str, Any] = {
                    "model": model,
                    "tools": TOOLS,
                    "store": False,
                    "input": [
                        {"role": "system", "content": "You are a helpful coding assistant."},
                        {"role": "user",   "content": user_message},
                    ],
                }
            else:
                # Send ALL tool results accumulated from the previous turn
                payload = {
                    "model": model,
                    "tools": TOOLS,
                    "store": False,
                    "previous_response_id": previous_response_id,
                    "input": pending_tool_results,
                }

            # ── Send ───────────────────────────────────────────────────────────
            await ws.send(json.dumps({"type": "response.create", "response": payload}))

            # ── Receive with per-event timeout ─────────────────────────────────
            result = await asyncio.wait_for(
                _receive_until_done(ws, event_timeout),
                timeout=event_timeout * 10,
            )

            previous_response_id = result.response_id

            # ── Terminal condition ─────────────────────────────────────────────
            if not result.has_tool_calls:
                logger.info("Final answer reached after %d turn(s)", turn + 1)
                return result.final_text

            # ── Execute ALL tool calls and collect results ─────────────────────
            pending_tool_results = []
            for call in result.tool_calls:
                logger.info("Executing tool '%s' (call_id=%s)", call.name, call.call_id)
                try:
                    output = execute_tool(call.name, call.arguments)
                except ToolExecutionError as exc:
                    output = str(exc)

                pending_tool_results.append({
                    "type":    "function_call_output",
                    "call_id": call.call_id,
                    "output":  output,
                })

    raise MaxTurnsExceededError(
        f"Agent did not produce a final answer within {max_turns} turns."
    )


# ── Entry point ────────────────────────────────────────────────────────────────

if __name__ == "__main__":
    import sys

    try:
        answer = asyncio.run(run_websocket_agent(
            user_message="Read main.py, fix any type errors, and write the fixed version back.",
            model="gpt-4o",
            max_turns=30,
        ))
        print("\n--- FINAL ANSWER ---")
        print(answer)
    except MaxTurnsExceededError as exc:
        logger.error("Loop aborted: %s", exc)
        sys.exit(1)
    except KeyboardInterrupt:
        logger.info("Interrupted by user")
        sys.exit(0)