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 .openai_llm_engine import OpenAILlmEngine
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 Main:
    @classmethod
    async def create(cls, config: Config):
        self = cls()
        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,
                    self._config.temperature,
                    self._config.token_limit
                )
            case "hf":
                self._llm = HfLlmEngine(
                    self._config.model,
                    self._config.temperature,
                    self._config.api_token,
                )
            case "openai":
                self._llm = OpenAILlmEngine(
                    self._config.model,
                    self._config.temperature,
                    self._config.api_token,
                    self._config.token_limit
                )
            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(),
            metrics=SystemMetrics(),
            llm=self._llm,
            validator=XMLValidator(self._action_schema),
            parser=ResponseParser(self._io_buffer)
        )
        self._ws_manager = await WebSocketManager.create(
            self._agent,
            self._io_buffer
        )

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

    @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__":
    loop = asyncio.new_event_loop()
    config = Config()
    main = loop.run_until_complete(Main.create(config))
    print(f"Web server started at http://{config.host}:{config.port}")
    web.run_app(main.app, loop=loop, host=config.host, port=config.port)