from aiohttp import web, WSMsgType
from pathlib import Path
from typing import Set, Optional
import asyncio
import json
from dataclasses import dataclass
import mimetypes
from functools import partial

from .web_agent import WebAgent, WebAgentState
from .web_io_buffer import WebIOBuffer

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

@dataclass
class ClientMessage:
    """Messages that can be received from clients."""
    APPROVE_CONTEXT = "APPROVE_CONTEXT"
    APPROVE_RESPONSE = "APPROVE_RESPONSE" 
    MODIFY_RESPONSE = "MODIFY_RESPONSE"
    SEND_INPUT = "SEND_INPUT"

@dataclass
class ServerMessage:
    """Messages that can be sent to clients."""
    STATE_CHANGE = "STATE_CHANGE"
    CONTEXT_UPDATE = "CONTEXT_UPDATE"
    RESPONSE_UPDATE = "RESPONSE_UPDATE"
    OUTPUT_UPDATE = "OUTPUT_UPDATE"
    VALIDATION_ERROR = "VALIDATION_ERROR"

class WebServer:
    """
    AIOHTTP-based web server that manages WebSocket connections and 
    integrates with the WebAgent. Serves static files from a web directory.
    """
    def __init__(self, agent: WebAgent, io_buffer: WebIOBuffer, 
                 host: str = "0.0.0.0", port: int = 8080,
                 static_dir: Optional[Path] = None):
        """Initialize the web server."""
        self.agent = agent
        self.io_buffer = io_buffer
        self.host = host
        self.port = port
        self.app = web.Application()
        self.static_dir = static_dir or Path(__file__).parent.parent / "static"
        self._init_routes()
        self._clients: Set[web.WebSocketResponse] = set()
        self._lock = asyncio.Lock()
        self._response_tasks: List[asyncio.Task] = []  # Initialize tasks list
        self.runner: Optional[web.AppRunner] = None
        
        # Register wrapped agent handlers
        self.agent.add_state_change_handler(self._wrap_async_handler(self._handle_state_change))
        self.agent.add_response_change_handler(self._wrap_async_handler(self._handle_response_change))

    def _wrap_async_handler(self, coro_func):
        """Wrap an async handler function to be called synchronously."""
        def wrapper(*args, **kwargs):
            loop = asyncio.get_event_loop()
            task = loop.create_task(coro_func(*args, **kwargs))
            self._response_tasks.append(task)
            self._response_tasks = [t for t in self._response_tasks if not t.done()]
            return task
        return wrapper

    def _init_routes(self):
        """Initialize application routes."""
        self.app.middlewares.append(self._cors_middleware)
        self.app.router.add_get("/ws", self._handle_websocket)

        if self.static_dir.exists():
            self.app.router.add_get("/", self._serve_index)
            self.app.router.add_static("/static/", self.static_dir, show_index=False)
            self.app.router.add_static("/assets/", self.static_dir / "assets", show_index=False)
            self.app.router.add_get("/{path:.*}", self._serve_index)

    async def _handle_client_message(self, data: dict):
        """Handle incoming client message."""
        message_type = data.get("type")
        
        if message_type == ClientMessage.APPROVE_CONTEXT:
            if self.agent.current_state == WebAgentState.CONTEXT_APPROVAL:
                self.agent.approve_context()
                
        elif message_type == ClientMessage.APPROVE_RESPONSE:
            if self.agent.current_state == WebAgentState.RESPONSE_APPROVAL:
                self.agent.approve_response()
                
        elif message_type == ClientMessage.MODIFY_RESPONSE:
            if self.agent.current_state == WebAgentState.RESPONSE_APPROVAL:
                new_response = data.get("data")
                if new_response:
                    self.agent._set_response(new_response)
                    
        elif message_type == ClientMessage.SEND_INPUT:
            input_text = data.get("data")
            if input_text:
                if not input_text.endswith('\n'):
                    input_text += '\n'
                
                self.io_buffer.append_stdin(input_text)
                
                current_length = self.io_buffer.buffer_length()
                return
        else:
            print(f"Agent in wrong state: {self.agent.current_state}")

    async def _serve_index(self, request: web.Request) -> web.Response:
        """Serve the React application HTML for any unmatched routes."""
        index_path = self.static_dir / "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

    @property
    def clients(self) -> Set[web.WebSocketResponse]:
        """Get the set of connected clients."""
        return self._clients

    async def start(self):
        """Start the web server."""
        # Verify static directory exists
        if not self.static_dir.exists():
            raise FileNotFoundError(f"Static directory not found: {self.static_dir}")

        self.runner = web.AppRunner(self.app)
        await self.runner.setup()
        site = web.TCPSite(self.runner, self.host, self.port)
        await site.start()
        print(f"Web server started at http://{self.host}:{self.port}")
        print(f"Serving static files from: {self.static_dir}")
        print(f"WebSocket endpoint at ws://{self.host}:{self.port}/ws")

    async def stop(self):
        """Stop the web server and clean up resources."""
        if self.runner:
            await self.runner.cleanup()

    async def _send_to_client(self, ws: web.WebSocketResponse, message: dict) -> bool:
        """Send message to client, return True if successful."""
        if not ws.closed:
            try:
                await ws.send_json(message)
                return True
            except Exception as e:
                print(f"Error sending to client: {e}")
                return False
        return False
            
    async def broadcast(self, message_type: str, data: str, **kwargs):
        """Broadcast a message to all connected clients."""
        async with self._lock:
            message = {
                "type": message_type,
                "data": data,
                **kwargs
            }
            
            tasks = []
            disconnected = set()
            
            for ws in self._clients:
                if not ws.closed:
                    task = asyncio.create_task(self._send_to_client(ws, message))
                    tasks.append((ws, task))
                else:
                    disconnected.add(ws)
                    
            if tasks:
                for ws, task in tasks:
                    try:
                        success = await task
                        if not success:
                            disconnected.add(ws)
                    except Exception:
                        disconnected.add(ws)
                    
            self._clients -= disconnected

    async def _handle_websocket(self, request: web.Request) -> web.WebSocketResponse:
        """Handle new WebSocket connections and messages."""
        ws = web.WebSocketResponse(heartbeat=30)
        await ws.prepare(request)
        
        async with self._lock:
            self._clients.add(ws)
        
        try:
            # Send initial state
            await self._send_to_client(ws, {
                "type": ServerMessage.STATE_CHANGE,
                "data": self.agent.current_state.name if self.agent.current_state else "INITIAL"
            })
            
            # Send current context if available
            if self.agent.context:
                await self._send_to_client(ws, {
                    "type": ServerMessage.CONTEXT_UPDATE,
                    "data": self.agent.context
                })
                
            # Send current response if available
            if self.agent.response:
                await self._send_to_client(ws, {
                    "type": ServerMessage.RESPONSE_UPDATE,
                    "data": self.agent.response,
                    "validation_error": self.agent._validation_error
                })
            
            # Handle incoming messages
            async for msg in ws:
                if msg.type == WSMsgType.TEXT:
                    try:
                        data = json.loads(msg.data)
                        await self._handle_client_message(data)
                                
                    except json.JSONDecodeError:
                        print(f"Invalid JSON message: {msg.data}")
                    except Exception as e:
                        print(f"Error handling message: {e}")
                        
                elif msg.type == WSMsgType.ERROR:
                    print(f"WebSocket connection closed with error: {ws.exception()}")
                    break
                    
        finally:
            async with self._lock:
                self._clients.remove(ws)
                
        return ws

    async def _handle_client_message(self, data: dict):
        message_type = data.get("type")
    
        if message_type == ClientMessage.APPROVE_CONTEXT:
            if self.agent.current_state == WebAgentState.CONTEXT_APPROVAL:
                await self.agent.approve_context()                
        elif message_type == ClientMessage.APPROVE_RESPONSE:
            if self.agent.current_state == WebAgentState.RESPONSE_APPROVAL:
                self.agent.approve_response()  # Now synchronous
                
        elif message_type == ClientMessage.MODIFY_RESPONSE:
            if self.agent.current_state == WebAgentState.RESPONSE_APPROVAL:
                new_response = data.get("data")
                if new_response:
                    self.agent._set_response(new_response)
                    
        elif message_type == ClientMessage.SEND_INPUT:
            input_text = data.get("data")
            if input_text:
                self.agent.set_input(input_text)
                self.io_buffer.append_stdin(input_text + "\n")
        else:
            print(f"Agent in wrong state: {self.agent.current_state}")
        
    async def _handle_state_change(self, new_state: WebAgentState):
        """Handle state changes from the WebAgent."""
        await self.broadcast(ServerMessage.STATE_CHANGE, new_state.name)
        
    async def _handle_response_change(self, response: str):
        """Handle response changes from the WebAgent."""
        await self.broadcast(
            ServerMessage.RESPONSE_UPDATE,
            response,
            validation_error=self.agent._validation_error
        )

    async def broadcast(self, message_type: str, data: str, **kwargs):
        """Broadcast a message to all connected clients."""
        async with self._lock:
            message = {
                "type": message_type,
                "data": data,
                **kwargs
            }
            
            tasks = []
            disconnected = set()
            
            for ws in self._clients:
                if not ws.closed:
                    try:
                        await self._send_to_client(ws, message)
                    except Exception as e:
                        print(f"Error sending to client: {e}")
                        disconnected.add(ws)
                else:
                    disconnected.add(ws)
                    
            self._clients -= disconnected