#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
JSON-RPC Server Implementation
JSON-RPC 服务器实现

Provides HTTP and WebSocket JSON-RPC server implementations.
"""

import json
import asyncio
from typing import Any, Dict, Optional, Callable, List
from dataclasses import dataclass
from http.server import HTTPServer, BaseHTTPRequestHandler
import threading
import logging

try:
    from jsonrpc import JSONRPCResponseManager, dispatcher
    from jsonrpc.exceptions import JSONRPCError
    JSON_RPC_AVAILABLE = True
except ImportError:
    JSON_RPC_AVAILABLE = False

from .jsonrpc_enhanced import JSONRPCMethodRegistry, EnhancedProgressReporter


@dataclass
class ServerConfig:
    """服务器配置"""
    host: str = "localhost"
    port: int = 8080
    debug: bool = False
    cors_enabled: bool = True
    max_request_size: int = 1024 * 1024  # 1MB


class JSONRPCHTTPHandler(BaseHTTPRequestHandler):
    """JSON-RPC HTTP 请求处理器"""
    
    def __init__(self, method_registry: JSONRPCMethodRegistry, config: ServerConfig, *args, **kwargs):
        self.method_registry = method_registry
        self.config = config
        super().__init__(*args, **kwargs)
    
    def do_POST(self):
        """处理POST请求"""
        try:
            # 检查Content-Type
            content_type = self.headers.get('Content-Type', '')
            if 'application/json' not in content_type:
                self._send_error(400, "Content-Type must be application/json")
                return
            
            # 读取请求体
            content_length = int(self.headers.get('Content-Length', 0))
            if content_length > self.config.max_request_size:
                self._send_error(413, "Request too large")
                return
            
            request_data = self.rfile.read(content_length).decode('utf-8')
            
            # 处理JSON-RPC请求
            response_data = self.method_registry.handle_request(request_data)

            # 如果响应为None，表示响应已经异步发送，不需要额外响应
            if response_data is not None:
                # 发送响应
                self._send_json_response(response_data)
            else:
                # 发送空的200响应，表示请求已处理但无需响应体
                self.send_response(200)
                if self.config.cors_enabled:
                    self._send_cors_headers()
                self.end_headers()
            
        except Exception as e:
            if self.config.debug:
                logging.exception("Error handling request")
            self._send_error(500, f"Internal server error: {str(e)}")
    
    def do_OPTIONS(self):
        """处理OPTIONS请求（CORS预检）"""
        if self.config.cors_enabled:
            self._send_cors_headers()
            self.end_headers()
        else:
            self._send_error(405, "Method not allowed")
    
    def _send_json_response(self, data: str):
        """发送JSON响应"""
        self.send_response(200)
        self.send_header('Content-Type', 'application/json')
        if self.config.cors_enabled:
            self._send_cors_headers()
        self.end_headers()
        self.wfile.write(data.encode('utf-8'))
    
    def _send_cors_headers(self):
        """发送CORS头"""
        self.send_header('Access-Control-Allow-Origin', '*')
        self.send_header('Access-Control-Allow-Methods', 'POST, OPTIONS')
        self.send_header('Access-Control-Allow-Headers', 'Content-Type')
    
    def _send_error(self, code: int, message: str):
        """发送错误响应"""
        self.send_response(code)
        self.send_header('Content-Type', 'application/json')
        if self.config.cors_enabled:
            self._send_cors_headers()
        self.end_headers()
        
        error_response = {
            "jsonrpc": "2.0",
            "id": None,
            "error": {
                "code": code,
                "message": message
            }
        }
        self.wfile.write(json.dumps(error_response).encode('utf-8'))
    
    def log_message(self, format, *args):
        """重写日志方法"""
        if self.config.debug:
            super().log_message(format, *args)


class JSONRPCServer:
    """JSON-RPC HTTP 服务器"""
    
    def __init__(self, config: Optional[ServerConfig] = None):
        self.config = config or ServerConfig()
        self.method_registry = JSONRPCMethodRegistry()
        self.server: Optional[HTTPServer] = None
        self.server_thread: Optional[threading.Thread] = None
        self.running = False
    
    def register_method(self, name: Optional[str] = None):
        """注册方法装饰器"""
        return self.method_registry.register(name)
    
    def register_function(self, func: Callable, name: Optional[str] = None):
        """注册函数"""
        self.method_registry.register_function(func, name)
    
    def start(self, blocking: bool = True):
        """启动服务器"""
        if self.running:
            raise RuntimeError("Server is already running")
        
        # 创建处理器工厂
        def handler_factory(*args, **kwargs):
            return JSONRPCHTTPHandler(self.method_registry, self.config, *args, **kwargs)
        
        # 创建服务器
        self.server = HTTPServer((self.config.host, self.config.port), handler_factory)
        self.running = True
        
        print(f"🚀 JSON-RPC Server started on http://{self.config.host}:{self.config.port}")
        
        if blocking:
            try:
                self.server.serve_forever()
            except KeyboardInterrupt:
                self.stop()
        else:
            self.server_thread = threading.Thread(target=self.server.serve_forever)
            self.server_thread.daemon = True
            self.server_thread.start()
    
    def stop(self):
        """停止服务器"""
        if self.server and self.running:
            print("🛑 Stopping JSON-RPC Server...")
            self.server.shutdown()
            self.server.server_close()
            self.running = False
            
            if self.server_thread:
                self.server_thread.join(timeout=5)
    
    def __enter__(self):
        """上下文管理器入口"""
        return self
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        """上下文管理器出口"""
        self.stop()


# 异步WebSocket服务器（需要额外依赖）
try:
    import websockets
    import asyncio
    
    class JSONRPCWebSocketServer:
        """JSON-RPC WebSocket 服务器"""
        
        def __init__(self, config: Optional[ServerConfig] = None):
            self.config = config or ServerConfig()
            self.method_registry = JSONRPCMethodRegistry()
            self.clients: List[Any] = []
        
        def register_method(self, name: Optional[str] = None):
            """注册方法装饰器"""
            return self.method_registry.register(name)
        
        async def handle_client(self, websocket, path):
            """处理WebSocket客户端"""
            self.clients.append(websocket)
            try:
                async for message in websocket:
                    try:
                        response = self.method_registry.handle_request(message)
                        await websocket.send(response)
                    except Exception as e:
                        error_response = {
                            "jsonrpc": "2.0",
                            "id": None,
                            "error": {
                                "code": -32603,
                                "message": f"Internal error: {str(e)}"
                            }
                        }
                        await websocket.send(json.dumps(error_response))
            except websockets.exceptions.ConnectionClosed:
                pass
            finally:
                if websocket in self.clients:
                    self.clients.remove(websocket)
        
        async def broadcast(self, method: str, params: Any = None):
            """广播通知到所有客户端"""
            if not self.clients:
                return
            
            notification = {
                "jsonrpc": "2.0",
                "method": method,
                "params": params
            }
            message = json.dumps(notification)
            
            # 发送到所有连接的客户端
            disconnected = []
            for client in self.clients:
                try:
                    await client.send(message)
                except websockets.exceptions.ConnectionClosed:
                    disconnected.append(client)
            
            # 清理断开的连接
            for client in disconnected:
                self.clients.remove(client)
        
        async def start(self):
            """启动WebSocket服务器"""
            print(f"🚀 JSON-RPC WebSocket Server started on ws://{self.config.host}:{self.config.port}")
            
            async with websockets.serve(
                self.handle_client,
                self.config.host,
                self.config.port
            ):
                await asyncio.Future()  # 永远运行

except ImportError:
    class JSONRPCWebSocketServer:
        def __init__(self, *args, **kwargs):
            raise ImportError("WebSocket server requires 'websockets' package")


# 便捷函数
def create_server(config: Optional[ServerConfig] = None) -> JSONRPCServer:
    """创建HTTP服务器"""
    return JSONRPCServer(config)


def create_websocket_server(config: Optional[ServerConfig] = None) -> JSONRPCWebSocketServer:
    """创建WebSocket服务器"""
    return JSONRPCWebSocketServer(config)
