from typing import Optional, Tuple, Any
import httpx
import psutil
import sentry_sdk
from loguru import logger
import functools

from BowongModalFunctions.models.web_model import WebhookNotify, WebhookMethodEnum, BaseFFMPEGTaskStatusResponse, \
    TaskStatus, ErrorCode


class SentryUtils:
    @staticmethod
    def capture_hardware_info() -> Tuple[Any, int, Any]:
        cpu_freq = psutil.cpu_freq()
        cpu_count = psutil.cpu_count()
        mem = psutil.virtual_memory()

        return cpu_freq, cpu_count, mem

    @staticmethod
    def sentry_tracker(sentry_trace_id: Optional[str], sentry_baggage: Optional[str],
                       op: Optional[str], name: Optional[str], fn_id: str):
        def decorator(func):
            @functools.wraps(func)
            def wrapper(*args, **kwargs):
                if sentry_trace_id and sentry_baggage:
                    transaction = sentry_sdk.continue_trace(environ_or_headers={"sentry-trace": sentry_trace_id,
                                                                                "baggage": sentry_baggage, })
                else:
                    transaction = sentry_sdk.start_transaction(op='modal.function', name='Modal Function直接调用')
                with transaction:
                    with transaction.start_child(op=op, name=name) as span:
                        cpu_freq, cpu_count, mem = SentryUtils.capture_hardware_info()
                        total_mb = mem.total >> 20
                        # available_mb = mem.available >> 20
                        sentry_sdk.set_tag('fn.id', fn_id)
                        span.set_data('cpu.count', cpu_count)
                        span.set_data('cpu.frequency', f"{cpu_freq.current / 1000:.2f} GHz")
                        span.set_data('memory.available', f"{total_mb} Mb")
                        result = func(*args, **kwargs)
                return result

            return wrapper

        return decorator

    @staticmethod
    def webhook_handler(webhook: WebhookNotify, func_id: str):
        def decorator(func):
            @functools.wraps(func)
            async def async_wrapper(*args, **kwargs):
                try:
                    result = await func(*args, **kwargs)
                    status = TaskStatus.success
                    error = None
                    code = ErrorCode.SUCCESS.value
                except Exception as e:
                    logger.exception(e)
                    result = None
                    status = TaskStatus.failed
                    error = e.message if hasattr(e, 'message') else str(e)
                    code = ErrorCode.SYSTEM_ERROR.value
                logger.info(f"webhook = {webhook}")
                if webhook:
                    with httpx.Client() as client:
                        match webhook.method:
                            case WebhookMethodEnum.POST:
                                response = client.post(url=webhook.endpoint.__str__(),
                                                       json=BaseFFMPEGTaskStatusResponse(
                                                           taskId=func_id, status=status, error=error,
                                                           code=code,
                                                           results=[result] if isinstance(result, str) else result,
                                                       ).model_dump())
                                logger.info(f"webhook response = {response}")
                            case WebhookMethodEnum.GET:
                                response = client.post(url=webhook.endpoint.__str__())
                        response.raise_for_status()
                return result

            return async_wrapper

        return decorator