#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Enhanced JSON-RPC Communication Module
增强版 JSON-RPC 通信模块

Uses json-rpc library for robust JSON-RPC 2.0 communication.
"""

import json
import sys
import time
import uuid
import asyncio
from typing import Any, Dict, Optional, Union, Callable, List
from dataclasses import dataclass, asdict
from enum import Enum

# json-rpc library imports
from jsonrpc import JSONRPCResponseManager, dispatcher
from jsonrpc.exceptions import JSONRPCError, JSONRPCInvalidRequest, JSONRPCMethodNotFound
from jsonrpc.jsonrpc2 import JSONRPC20Request, JSONRPC20Response

class ProgressLevel(str, Enum):
    """进度级别"""
    INFO = "info"
    SUCCESS = "success"
    WARNING = "warning"
    ERROR = "error"
    DEBUG = "debug"


@dataclass
class ProgressUpdate:
    """进度更新数据结构"""
    step: str
    type: str
    progress: int
    message: str
    timestamp: float
    level: ProgressLevel = ProgressLevel.INFO
    data: Optional[Dict[str, Any]] = None


class EnhancedJSONRPCResponse:
    """增强版 JSON-RPC 2.0 Response handler"""
    
    def __init__(self, request_id: Optional[Union[str, int]] = None):
        # 如果没有提供request_id，生成一个UUID
        if request_id is None:
            self.request_id = str(uuid.uuid4())
        else:
            self.request_id = request_id
    
    def success(self, result: Any) -> None:
        """发送成功响应"""
        response = JSONRPC20Response(result=result, _id=self.request_id)
        self._send_response(response.data)
    
    def error(self, code: int, message: str, data: Any = None) -> None:
        """发送错误响应"""
        error = JSONRPCError(code=code, message=message, data=data)
        response = JSONRPC20Response(error=error, _id=self.request_id)
        self._send_response(response.data)
    
    def progress(self, step: str, progress: int, message: str, 
                level: ProgressLevel = ProgressLevel.INFO, data: Optional[Dict[str, Any]] = None) -> None:
        """发送进度通知"""
        progress_data = ProgressUpdate(
            step=step,
            type="progress",
            progress=progress,
            message=message,
            timestamp=time.time(),
            level=level,
            data=data
        )
        
        self.notification(asdict(progress_data))
    
    def notification(self, params: Any = None) -> None:
        """发送通知（无需响应）"""
        response = JSONRPC20Response(result=params, _id=self.request_id)
        self._send_response(response.data)
    
    def _send_response(self, response: Dict[str, Any]) -> None:
        """发送响应到标准输出"""
        json_str = json.dumps(response, ensure_ascii=False, separators=(',', ':'))
        print(f"JSONRPC:{json_str}", file=sys.stdout, flush=True)

class EnhancedProgressReporter:
    """增强版进度报告器"""
    
    def __init__(self, total: int = 0):
        self.response = EnhancedJSONRPCResponse()
        self.step = 0
        self.total = total
    
    def update(self, message: str, step: Optional[int] = None, data: Optional[Dict[str, Any]] = None) -> None:
        """更新进度"""
        if step is not None:
            self.step = step
        else:
            self.step += 1
        
        progress = int((self.step / self.total * 100)) if self.total > 0 else -1
        self.response.progress("update", progress, message, ProgressLevel.INFO, data)
    
    def info(self, message: str, data: Optional[Dict[str, Any]] = None) -> None:
        """信息消息"""
        self.response.progress("info", -1, message, ProgressLevel.INFO, data)
    
    def success(self, message: str, data: Optional[Dict[str, Any]] = None) -> None:
        """成功消息"""
        self.response.progress("success", -1, message, ProgressLevel.SUCCESS, data)
    
    def warning(self, message: str, data: Optional[Dict[str, Any]] = None) -> None:
        """警告消息"""
        self.response.progress("warning", -1, message, ProgressLevel.WARNING, data)
    
    def error(self, message: str, data: Optional[Dict[str, Any]] = None) -> None:
        """错误消息"""
        self.response.progress("error", -1, message, ProgressLevel.ERROR, data)
    
    def debug(self, message: str, data: Optional[Dict[str, Any]] = None) -> None:
        """调试消息"""
        self.response.progress("debug", -1, message, ProgressLevel.DEBUG, data)


class JSONRPCMethodRegistry:
    """JSON-RPC 方法注册器"""
    
    def __init__(self):
        self.methods: Dict[str, Callable] = {}
        self.dispatcher = dispatcher
    
    def register(self, name: Optional[str] = None):
        """注册方法装饰器"""
        def decorator(func: Callable):
            method_name = name or func.__name__
            self.methods[method_name] = func
            self.dispatcher.add_method(func, method_name)   
            return func
        return decorator
    
    def register_function(self, func: Callable, name: Optional[str] = None) -> None:
        """直接注册函数"""
        method_name = name or func.__name__
        self.methods[method_name] = func
        self.dispatcher.add_method(func, method_name)
            
    
    def handle_request(self, request_data: str) -> Optional[str]:
        """处理JSON-RPC请求"""
        try:
            # 解析请求
            request_json = json.loads(request_data)
            method_name = request_json.get("method")
            params = request_json.get("params", {})
            request_id = request_json.get("id")

            # 检查方法是否存在
            if method_name not in self.methods:
                error_response = {
                    "jsonrpc": "2.0",
                    "id": request_id,
                    "error": {
                        "code": -32601,
                        "message": f"Method '{method_name}' not found"
                    }
                }
                return json.dumps(error_response, ensure_ascii=False, separators=(',', ':'))

            try:
                # 调用方法
                if isinstance(params, dict):
                    result = self.methods[method_name](**params)
                elif isinstance(params, list):
                    result = self.methods[method_name](*params)
                else:
                    result = self.methods[method_name](params)

                # 如果方法返回None，表示响应已经异步发送，不需要额外响应
                if result is None:
                    return None

                # 正常响应
                success_response = {
                    "jsonrpc": "2.0",
                    "id": request_id,
                    "result": result
                }
                return json.dumps(success_response, ensure_ascii=False, separators=(',', ':'))

            except Exception as e:
                error_response = {
                    "jsonrpc": "2.0",
                    "id": request_id,
                    "error": {
                        "code": -32603,
                        "message": f"Internal error: {str(e)}"
                    }
                }
                return json.dumps(error_response, ensure_ascii=False, separators=(',', ':'))

        except Exception as e:
            error_response = {
                "jsonrpc": "2.0",
                "id": None,
                "error": {
                    "code": -32700,
                    "message": f"Parse error: {str(e)}"
                }
            }
            return json.dumps(error_response, ensure_ascii=False, separators=(',', ':'))

# 全局实例
enhanced_progress_reporter = EnhancedProgressReporter()
method_registry = JSONRPCMethodRegistry()


# 便捷函数
def create_response_handler(request_id: Optional[Union[str, int]] = None) -> EnhancedJSONRPCResponse:
    """创建响应处理器"""
    return EnhancedJSONRPCResponse(request_id)


def create_progress_reporter(total: int = 0) -> EnhancedProgressReporter:
    """创建进度报告器"""
    return EnhancedProgressReporter(total)


def register_method(name: Optional[str] = None):
    """注册JSON-RPC方法"""
    return method_registry.register(name)


def handle_jsonrpc_request(request_data: str) -> str:
    """处理JSON-RPC请求"""
    return method_registry.handle_request(request_data)


# 向后兼容的别名
JSONRPCResponse = EnhancedJSONRPCResponse
ProgressReporter = EnhancedProgressReporter
