from __future__ import annotations

import json
from typing import Any
from urllib import request

from config import ServerConfig
from oauth import OAuthError, QwenOAuthManager
from tools import ToolError, ToolRegistry


DEFAULT_SYSTEM_PROMPT = """You are new-qwen-serv, a server-side coding agent.
You help the user through a remote client.
Use tools when they are necessary.
Be concise, precise and action-oriented.
When you modify files, explain what changed.
Do not claim to have executed tools unless a tool result confirms it."""


class QwenAgent:
    def __init__(self, config: ServerConfig, oauth: QwenOAuthManager, tools: ToolRegistry) -> None:
        self.config = config
        self.oauth = oauth
        self.tools = tools

    def _request_completion(self, messages: list[dict[str, Any]]) -> dict[str, Any]:
        creds = self.oauth.get_valid_credentials()
        base_url = self.oauth.get_openai_base_url(creds)
        payload = {
            "model": self.config.model,
            "messages": messages,
            "tools": self.tools.schemas(),
            "tool_choice": "auto",
        }
        data = json.dumps(payload).encode("utf-8")
        req = request.Request(
            f"{base_url}/chat/completions",
            data=data,
            headers={
                "Content-Type": "application/json",
                "Authorization": f"Bearer {creds['access_token']}",
            },
            method="POST",
        )
        with request.urlopen(req, timeout=180) as response:
            return json.loads(response.read().decode("utf-8"))

    def run(self, history: list[dict[str, Any]], user_message: str) -> dict[str, Any]:
        system_prompt = self.config.system_prompt or DEFAULT_SYSTEM_PROMPT
        messages: list[dict[str, Any]] = [{"role": "system", "content": system_prompt}]
        messages.extend(history)
        messages.append({"role": "user", "content": user_message})
        events: list[dict[str, Any]] = []

        for _ in range(self.config.max_tool_rounds):
            response = self._request_completion(messages)
            choice = response["choices"][0]["message"]
            tool_calls = choice.get("tool_calls") or []
            content = choice.get("content")
            if content:
                events.append({"type": "assistant", "content": content})

            if not tool_calls:
                return {
                    "answer": content or "",
                    "events": events,
                    "usage": response.get("usage"),
                    "messages": messages + [{"role": "assistant", "content": content or ""}],
                }

            messages.append(
                {
                    "role": "assistant",
                    "content": content or "",
                    "tool_calls": tool_calls,
                }
            )

            for call in tool_calls:
                tool_name = call["function"]["name"]
                try:
                    arguments = json.loads(call["function"]["arguments"] or "{}")
                except json.JSONDecodeError:
                    arguments = {}

                events.append({"type": "tool_call", "name": tool_name, "arguments": arguments})
                try:
                    result = self.tools.execute(tool_name, arguments)
                except ToolError as exc:
                    result = {"error": str(exc)}
                except Exception as exc:
                    result = {"error": f"Unexpected tool failure: {exc}"}

                events.append({"type": "tool_result", "name": tool_name, "result": result})
                messages.append(
                    {
                        "role": "tool",
                        "tool_call_id": call["id"],
                        "content": self.tools.encode_result(result),
                    }
                )

        raise OAuthError("Max tool rounds exceeded")