from concurrent.futures import ThreadPoolExecutor
from aiohttp import web
from pathlib import Path
import asyncio
import mimetypes
import time


from .config import Config
from .hf_llm_engine import HfLlmEngine
from .llm_engine import LlmEngine
from .local_llm_engine import LocalLlmEngine
from .response_parser import ResponseParser
from .system_metrics import SystemMetrics
from .web_agent import WebAgent
from .web_agent import WebAgent
from .web_io_buffer import WebIOBuffer
from .web_io_buffer import WebIOBuffer
from .web_socket_manager import WebSocketManager
from .working_memory import WorkingMemory
from .xml_validator import XMLValidator

mimetypes.add_type("application/javascript", ".js")
mimetypes.add_type("application/javascript", ".jsx")
mimetypes.add_type("text/javascript", ".js")
mimetypes.add_type("text/javascript", ".jsx")

class TestLLM:
    def infer(self, prompt: str, context: str):
        yield "<reasoning>"
        time.sleep(2)
        yield "test reasoning"
        time.sleep(2)
        yield "</reasoning>"

class Main:
    def __init__(self):
        self._config = Config()

        self._system_prompt = self._config.system_prompt.read_text()
        self._action_schema = self._config.action_schema.read_text()

        match self._config.llm_engine:
            case "local":
                self._llm = LocalLlmEngine(self._config.model)
            case "hf":
                self._llm = HfLlmEngine(
                    model_id=self._config.model,
                    api_token=self._config.hf_api_token
                )
            case "test":
                self._llm = TestLLM()
            case _:
                raise ValueError(f"Invalid LLM engine: {self._config.llm_engine}")
        self._io_buffer = WebIOBuffer()
        self._agent = WebAgent(
            system_prompt=self._system_prompt,
            action_schema=self._action_schema,
            working_memory=WorkingMemory(),
            system_metrics=SystemMetrics(),
            llm=self._llm,
            validator=XMLValidator(self._action_schema),
            parser=ResponseParser(self._io_buffer)
        )
        self._ws_manager = WebSocketManager(
            self._agent,
            self._io_buffer
        )

        self._app = web.Application()
        self._init_routes()

    @property
    def app(self):
        return self._app

    def _init_routes(self):
        """Initialize application routes."""
        self._app.middlewares.append(self._cors_middleware)
        self._app.router.add_get("/ws", self._ws_manager.handle_websocket)
        self._app.router.add_get("/", self._serve_index)
        self._app.router.add_static("/static/", self._config.static_files, show_index=False)
        self._app.router.add_static("/assets/", self._config.static_files / "assets", show_index=False)
        self._app.router.add_get("/{path:.*}", self._serve_index)

    async def _serve_index(self, request: web.Request) -> web.Response:
        """Serve the React application HTML for any unmatched routes."""
        index_path = self._config.static_files / "index.html"
        if not index_path.exists():
            raise web.HTTPNotFound()

        with open(index_path, "r") as f:
            html_content = f.read()

        return web.Response(
            text=html_content,
            content_type="text/html"
        )

    @web.middleware
    async def _cors_middleware(self, request: web.Request, handler):
        """Handle CORS headers."""
        if request.method == "OPTIONS":
            response = web.Response()
        else:
            response = await handler(request)

        response.headers.update({
            'Access-Control-Allow-Origin': '*',
            'Access-Control-Allow-Methods': 'GET, POST, OPTIONS',
            'Access-Control-Allow-Headers': 'Content-Type',
            'Access-Control-Max-Age': '86400',
        })
        return response

if __name__ == "__main__":
    main = Main()
    print("Web server started at http://localhost:8080")
    web.run_app(main.app, port=8080)