import asyncio
import json
import logging
from functools import wraps
from typing import Any, Callable, Coroutine, Dict, List, Optional, Union
from .utils import get_free_port
from aiohttp import web
import threading
import time
import aiohttp_cors  # CORS 지원을 위한 패키지 추가

# Configure logging
logging.basicConfig(level=logging.INFO)
log = logging.getLogger("pyloid.rpc")

class RPCError(Exception):
    """
    Custom exception for RPC-related errors.

    Follows the JSON-RPC 2.0 error object structure.

    Attributes
    ----------
    message : str
        A human-readable description of the error.
    code : int, optional
        A number indicating the error type that occurred. Standard JSON-RPC
        codes are used where applicable, with application-specific codes
        also possible. Defaults to -32000 (Server error).
    data : Any, optional
        Additional information about the error, by default None.
    """
    def __init__(self, message: str, code: int = -32000, data: Any = None):
        """
        Initialize the RPCError.

        Parameters
        ----------
        message : str
            The error message.
        code : int, optional
            The error code. Defaults to -32000.
        data : Any, optional
            Additional data associated with the error. Defaults to None.
        """
        self.message = message
        self.code = code
        self.data = data
        super().__init__(self.message)

    def to_dict(self) -> Dict[str, Any]:
        """
        Convert the error details into a dictionary suitable for JSON-RPC responses.

        Returns
        -------
        Dict[str, Any]
            A dictionary representing the JSON-RPC error object.
        """
        error_obj = {"code": self.code, "message": self.message}
        if self.data is not None:
            error_obj["data"] = self.data
        return error_obj

class PyloidRPC:
    """
    A simple JSON-RPC server wrapper based on aiohttp.

    Allows registering asynchronous functions as RPC methods using the `@rpc`
    decorator and handles JSON-RPC 2.0 request parsing, validation,
    method dispatching, and response formatting.

    Attributes
    ----------
    _host : str
        The hostname or IP address to bind the server to.
    _port : int
        The port number to listen on.
    _rpc_path : str
        The URL path for handling RPC requests.
    _functions : Dict[str, Callable[..., Coroutine[Any, Any, Any]]]
        A dictionary mapping registered RPC method names to their
        corresponding asynchronous functions.
    _app : web.Application
        The underlying aiohttp web application instance.
    """
    def __init__(self):
        """
        Initialize the PyloidRPC server instance.
        """
        self._host = "127.0.0.1"
        self._port = get_free_port()
        self._rpc_path = "/rpc"
        
        self.url = f"http://{self._host}:{self._port}{self._rpc_path}"
        
        self._functions: Dict[str, Callable[..., Coroutine[Any, Any, Any]]] = {}
        self._app = web.Application()
        
        # CORS 설정 추가
        cors = aiohttp_cors.setup(self._app, defaults={
            "*": aiohttp_cors.ResourceOptions(
                allow_credentials=True,
                expose_headers="*",
                allow_headers="*",
                allow_methods=["POST"]
            )
        })
        
        # CORS 적용된 라우트 추가
        resource = cors.add(self._app.router.add_resource(self._rpc_path))
        cors.add(resource.add_route("POST", self._handle_rpc))
        
        log.info(f"RPC server initialized.")
        self._runner: Optional[web.AppRunner] = None
        self._site: Optional[web.TCPSite] = None

    def rpc(self, name: Optional[str] = None) -> Callable:
        """
        Decorator to register an async function as an RPC method.

        Parameters
        ----------
        name : Optional[str], optional
            The name to register the RPC method under. If None, the
            function's name is used. Defaults to None.

        Returns
        -------
        Callable
            The decorator function.

        Raises
        ------
        TypeError
            If the decorated function is not an async function (`coroutinefunction`).
        ValueError
            If an RPC function with the specified name is already registered.
        """
        def decorator(func: Callable[..., Coroutine[Any, Any, Any]]):
            rpc_name = name or func.__name__
            if not asyncio.iscoroutinefunction(func):
                raise TypeError(f"RPC function '{rpc_name}' must be an async function.")
            if rpc_name in self._functions:
                raise ValueError(f"RPC function name '{rpc_name}' is already registered.")

            self._functions[rpc_name] = func
            log.info(f"RPC function registered: {rpc_name}")

            @wraps(func)
            async def wrapper(*args, **kwargs):
                 # This wrapper exists to follow the decorator pattern.
                 # The actual call uses the original function stored in _functions.
                return await func(*args, **kwargs)
            return wrapper
        return decorator

    def _validate_jsonrpc_request(self, data: Any) -> Optional[Dict[str, Any]]:
        """
        Validate the structure of a potential JSON-RPC request object.

        Checks for required fields ('jsonrpc', 'method') and validates the
        types of fields like 'params' and 'id' according to the JSON-RPC 2.0 spec.

        Parameters
        ----------
        data : Any
            The parsed JSON data from the request body.

        Returns
        -------
        Optional[Dict[str, Any]]
            None if the request is valid according to the basic structure,
            otherwise a dictionary representing the JSON-RPC error object
            to be returned to the client.
        """
        # Attempt to extract the ID if possible, even for invalid requests
        request_id = data.get("id") if isinstance(data, dict) else None

        if not isinstance(data, dict):
            return {"code": -32600, "message": "Invalid Request: Request must be a JSON object."}
        if data.get("jsonrpc") != "2.0":
            return {"code": -32600, "message": "Invalid Request: 'jsonrpc' version must be '2.0'."}
        if "method" not in data or not isinstance(data["method"], str):
            return {"code": -32600, "message": "Invalid Request: 'method' must be a string."}
        if "params" in data and not isinstance(data["params"], (list, dict)):
            # JSON-RPC 2.0: "params" must be array or object if present
            return {"code": -32602, "message": "Invalid params: 'params' must be an array or object."}
        # JSON-RPC 2.0: "id" is optional, but if present, must be string, number, or null.
        # This validation is simplified here. A more robust check could be added.
        # if "id" in data and not isinstance(data.get("id"), (str, int, float, type(None))):
        #     return {"code": -32600, "message": "Invalid Request: 'id', if present, must be a string, number, or null."}
        return None # Request structure is valid

    async def _handle_rpc(self, request: web.Request) -> web.Response:
        """
        Handles incoming JSON-RPC requests.

        Parses the request, validates it, dispatches to the appropriate
        registered RPC method, executes the method, and returns the
        JSON-RPC response or error object.

        Parameters
        ----------
        request : web.Request
            The incoming aiohttp request object.

        Returns
        -------
        web.Response
            An aiohttp JSON response object containing the JSON-RPC response or error.
        """
        request_id: Optional[Union[str, int, None]] = None
        data: Any = None # Define data outside try block for broader scope if needed

        try:
            # 1. Check Content-Type
            if request.content_type != 'application/json':
                 # Cannot determine ID if content type is wrong, respond with null ID
                 error_resp = {"jsonrpc": "2.0", "error": {"code": -32700, "message": "Parse error: Content-Type must be application/json."}, "id": None}
                 return web.json_response(error_resp, status=415) # Unsupported Media Type

            # 2. Parse JSON Body
            try:
                raw_data = await request.read()
                data = json.loads(raw_data)
                # Extract ID early for inclusion in potential error responses
                if isinstance(data, dict):
                    request_id = data.get("id") # Can be str, int, null, or absent
            except json.JSONDecodeError:
                # Invalid JSON, ID might be unknown, respond with null ID
                error_resp = {"jsonrpc": "2.0", "error": {"code": -32700, "message": "Parse error: Invalid JSON format."}, "id": None}
                return web.json_response(error_resp, status=400) # Bad Request

            # 3. Validate JSON-RPC Structure
            validation_error = self._validate_jsonrpc_request(data)
            if validation_error:
                 # Use extracted ID if available, otherwise it remains None
                 error_resp = {"jsonrpc": "2.0", "error": validation_error, "id": request_id}
                 return web.json_response(error_resp, status=400) # Bad Request

            # Assuming validation passed, data is a dict with 'method'
            method_name: str = data["method"]
            # Use empty list/dict if 'params' is omitted, as per spec flexibility
            params: Union[List, Dict] = data.get("params", [])

            # 4. Find and Call Method
            func = self._functions.get(method_name)
            if func is None:
                error_resp = {"jsonrpc": "2.0", "error": {"code": -32601, "message": "Method not found."}, "id": request_id}
                return web.json_response(error_resp, status=404) # Not Found

            try:
                log.debug(f"Executing RPC method: {method_name}(params={params})")
                # Call the function with positional or keyword arguments
                if isinstance(params, list):
                    result = await func(*params)
                else: # isinstance(params, dict)
                    result = await func(**params)

                # 5. Format Success Response (only for non-notification requests)
                if request_id is not None: # Notifications (id=null or absent) don't get responses
                    response_data = {"jsonrpc": "2.0", "result": result, "id": request_id}
                    return web.json_response(response_data)
                else:
                    # No response for notifications, return 204 No Content might be appropriate
                    # or just an empty response. aiohttp handles this implicitly if nothing is returned.
                    # For clarity/standard compliance, maybe return 204?
                    return web.Response(status=204)


            except RPCError as e:
                 # Application-specific error during method execution
                 log.warning(f"RPC execution error in method '{method_name}': {e}", exc_info=False)
                 if request_id is not None:
                     error_resp = {"jsonrpc": "2.0", "error": e.to_dict(), "id": request_id}
                     # Use 500 or a more specific 4xx/5xx if applicable based on error code?
                     # Sticking to 500 for server-side execution errors.
                     return web.json_response(error_resp, status=500)
                 else:
                     return web.Response(status=204) # No response for notification errors
            except Exception as e:
                # Unexpected error during method execution
                log.exception(f"Unexpected error during execution of RPC method '{method_name}':") # Log full traceback
                if request_id is not None:
                    # Minimize internal details exposed to the client
                    error_resp = {"jsonrpc": "2.0", "error": {"code": -32000, "message": f"Server error: {type(e).__name__}"}, "id": request_id}
                    return web.json_response(error_resp, status=500) # Internal Server Error
                else:
                    return web.Response(status=204) # No response for notification errors

        except Exception as e:
            # Catch-all for fatal errors during request handling itself (before/after method call)
            log.exception("Fatal error in RPC handler:")
            # ID might be uncertain at this stage, include if available
            error_resp = {"jsonrpc": "2.0", "error": {"code": -32603, "message": "Internal error"}, "id": request_id}
            return web.json_response(error_resp, status=500)

    async def start_async(self, **kwargs):
        """Starts the server asynchronously without blocking."""
        self._runner = web.AppRunner(self._app, access_log=None, **kwargs)
        await self._runner.setup()
        self._site = web.TCPSite(self._runner, self._host, self._port)
        await self._site.start()
        log.info(f"RPC server started asynchronously on {self.url}")
        # 서버가 백그라운드에서 실행되도록 여기서 블로킹하지 않습니다.
        # 이 코루틴은 서버 시작 후 즉시 반환됩니다.

    async def stop_async(self):
        """Stops the server asynchronously."""
        if self._runner:
            await self._runner.cleanup()
            log.info("RPC server stopped.")
        self._site = None
        self._runner = None

    def start(self, **kwargs):
        """
        Start the aiohttp web server to listen for RPC requests (blocking).

        This method wraps `aiohttp.web.run_app` and blocks until the server stops.
        Prefer `start_async` for non-blocking operation within an asyncio event loop.

        Parameters
        ----------
        **kwargs
            Additional keyword arguments to pass directly to `aiohttp.web.run_app`.
            For example, `ssl_context` for HTTPS. By default, suppresses the
            default `aiohttp` startup message using `print=None`.
        """
        log.info(f"Starting RPC server")
        # Default to print=None to avoid duplicate startup messages, can be overridden via kwargs
        run_app_kwargs = {'print': None, 'access_log': None}
        run_app_kwargs.update(kwargs)
        try:
            web.run_app(self._app, host=self._host, port=self._port, **run_app_kwargs)
        except Exception as e:
            log.exception(f"Failed to start or run the server: {e}")
            raise