from aiohttp import web
import json
import asyncio

from ..auto_approver import AutoApprover
from ..web_agent import WebAgent
from ..web_io_buffer import WebIOBuffer

class Api:
    def __init__(self, app: web.Application, agent: WebAgent, io_buffer: WebIOBuffer, auto_approver: AutoApprover):
        self._app = app
        self._agent = agent
        self._io_buffer = io_buffer
        self._auto_approver = auto_approver
        
        self._init_routes()
        
    def _init_routes(self):
        """Initialize REST API and WebSocket routes."""
        self._app.router.add_post("/api/inference/{llm}", self._run_inference)
        self._app.router.add_post("/api/approve/{llm}", self._approve_response)
        self._app.router.add_post("/api/context", self._modify_context)
        self._app.router.add_post("/api/input", self._send_input)
        self._app.router.add_post("/api/clear", self._clear_output)
        self._app.router.add_get("/api/output/{llm}", self._get_output)
        self._app.router.add_get("/api/llms", self._get_llms)
        self._app.router.add_get("/api/auto_approver/config", self._get_auto_approver_config)
        self._app.router.add_post("/api/auto_approver/config", self._set_auto_approver_config)
        self._app.router.add_post("/api/auto_approver/context_enabled", self._set_context_enabled)
        self._app.router.add_post("/api/auto_approver/response_enabled", self._set_response_enabled) 
        self._app.router.add_post("/api/auto_approver/context_timeout", self._set_context_timeout)
        self._app.router.add_post("/api/auto_approver/response_timeout", self._set_response_timeout)
        self._app.router.add_post("/api/auto_approver/llm", self._set_llm_name)

    async def _run_inference(self, request: web.Request) -> web.Response:
        """Start inference on specified LLM."""
        llm_name = request.match_info["llm"]
        try:
            await asyncio.get_event_loop().run_in_executor(None, self._agent.run_inference, llm_name)
            return web.Response(status=200)
        except (ValueError, RuntimeError) as e:
            return web.Response(status=400, text=str(e))

    async def _approve_response(self, request: web.Request) -> web.Response:
        """Approve response from specified LLM."""
        llm_name = request.match_info["llm"]
        data = await request.json()
        response = data.get("response")
        if not response:
            return web.Response(status=400, text="Missing response in request body")
        try:
            self._agent.approve_response(llm_name, response)
            return web.Response(status=200)
        except ValueError as e:
            return web.Response(status=400, text=str(e))

    async def _modify_context(self, request: web.Request) -> web.Response:
        """Modify the current context."""
        data = await request.json()
        context = data.get("context")
        if not context:
            return web.Response(status=400, text="Missing context in request body")
        self._agent.modify_context(context)
        return web.Response(status=200)

    async def _send_input(self, request: web.Request) -> web.Response:
        """Send input to the IO buffer."""
        data = await request.json()
        input_text = data.get("input")
        if not input_text:
            return web.Response(status=400, text="Missing input in request body")
        self._io_buffer.append_stdin(input_text)
        return web.Response(status=200)

    async def _clear_output(self, request: web.Request) -> web.Response:
        """Clear the stdout buffer."""
        self._io_buffer.clear_stdout()
        return web.Response(status=200)

    async def _get_output(self, request: web.Request) -> web.Response:
        """Get complete output for specified LLM."""
        llm_name = request.match_info["llm"]
        try:
            output = self._agent.get_output(llm_name)
            return web.Response(
                text=json.dumps({"output": output}),
                content_type="application/json"
            )
        except ValueError as e:
            return web.Response(status=400, text=str(e))

    async def _get_llms(self, request: web.Request) -> web.Response:
        """Get all LLMs and their current states."""
        states = self._agent.llms
        return web.Response(
            text=json.dumps({
                name: state.name 
                for name, state in states.items()
            }),
            content_type="application/json"
        )

    async def _get_auto_approver_config(self, request: web.Request) -> web.Response:
        """Get current auto approver configuration."""
        return web.Response(
            text=json.dumps(self._auto_approver.config),
            content_type="application/json"
        )

    async def _set_auto_approver_config(self, request: web.Request) -> web.Response:
        """Update auto approver configuration."""
        data = await request.json()
        try:
            self._auto_approver.set_config(data)
            return web.Response(status=200)
        except (ValueError, KeyError) as e:
            return web.Response(status=400, text=str(e))

    async def _set_context_enabled(self, request: web.Request) -> web.Response:
        """Set context auto-approval enabled state."""
        data = await request.json()
        enabled = data.get("enabled")
        if enabled is None:
            return web.Response(status=400, text="Missing enabled parameter")
        try:
            self._auto_approver.context_enabled = enabled
            return web.Response(status=200)
        except ValueError as e:
            return web.Response(status=400, text=str(e))

    async def _set_response_enabled(self, request: web.Request) -> web.Response:
        """Set response auto-approval enabled state."""
        data = await request.json()
        enabled = data.get("enabled")
        if enabled is None:
            return web.Response(status=400, text="Missing enabled parameter")
        try:
            self._auto_approver.response_enabled = enabled
            return web.Response(status=200)
        except ValueError as e:
            return web.Response(status=400, text=str(e))

    async def _set_context_timeout(self, request: web.Request) -> web.Response:
        """Set context auto-approval timeout."""
        data = await request.json()
        timeout = data.get("timeout")
        if timeout is None:
            return web.Response(status=400, text="Missing timeout parameter")
        try:
            self._auto_approver.context_timeout = float(timeout)
            return web.Response(status=200)
        except ValueError as e:
            return web.Response(status=400, text=str(e))

    async def _set_response_timeout(self, request: web.Request) -> web.Response:
        """Set response auto-approval timeout."""
        data = await request.json()
        timeout = data.get("timeout")
        if timeout is None:
            return web.Response(status=400, text="Missing timeout parameter")
        try:
            self._auto_approver.response_timeout = float(timeout)
            return web.Response(status=200)
        except ValueError as e:
            return web.Response(status=400, text=str(e))

    async def _set_llm_name(self, request: web.Request) -> web.Response:
        """Set LLM name for auto-approval."""
        data = await request.json()
        name = data.get("name")
        if name is None:
            return web.Response(status=400, text="Missing name parameter")
        try:
            self._auto_approver.llm_name = name
            return web.Response(status=200)
        except ValueError as e:
            return web.Response(status=400, text=str(e))