import os
from typing import Optional, Tuple, Any, List

import backoff
import httpx
import psutil
import sentry_sdk
from loguru import logger
import functools

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


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):
                logger.info(f"sentry-trace={sentry_trace_id}, baggage={sentry_baggage}")
                if sentry_trace_id and sentry_baggage:
                    transaction = sentry_sdk.continue_trace(environ_or_headers={"sentry-trace": sentry_trace_id,
                                                                                "baggage": sentry_baggage, },
                                                            op=op,
                                                            name=name)
                else:
                    transaction = sentry_sdk.start_transaction(op=op, name=name)
                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_tags({
                            'fn.id': fn_id,
                            'x_trace.id': sentry_trace_id,
                            'x_trace.baggage': sentry_baggage,
                            'cpu.count': cpu_count,
                            'cpu.frequency': f"{cpu_freq.current / 1000:.2f}",
                            'cpu.frequency.format': f"{cpu_freq.current / 1000:.2f} GHz",
                            'memory.available': total_mb,
                            'memory.available.format': f"{total_mb} Mb",
                            'modal.cloud.region': os.environ.get('MODAL_REGION', 'unknown'),
                            'modal.cloud.provider': os.environ.get('MODAL_CLOUD_PROVIDER', 'unknown'),
                            'modal.task.id': os.environ.get('MODAL_TASK_ID', 'unknown'),
                            'modal.identity.token': os.environ.get('MODAL_IDENTITY_TOKEN', 'unknown'),
                            'modal.image.id': os.environ.get('MODAL_IMAGE_ID', 'unknown'),
                        })
                        result = func(*args, **kwargs)
                        span.set_status('ok')
                return result

            return wrapper

        return decorator

    @staticmethod
    def webhook_handler(webhook: WebhookNotify, func_id: str, raise_on_giveup: bool = False):
        def decorator(func):
            @backoff.on_exception(exception=Exception, wait_gen=backoff.constant,
                                  max_time=15, max_tries=5, raise_on_giveup=raise_on_giveup)
            def webhook_with_retry(webhook: WebhookNotify, body: BaseFFMPEGTaskStatusResponse):
                if webhook.method == WebhookMethodEnum.POST:
                    response = httpx.post(url=webhook.endpoint.__str__(),
                                          headers=webhook.headers,
                                          json=body.model_dump(),
                                          timeout=3)
                else:
                    response = httpx.get(url=webhook.endpoint.__str__(),
                                         headers=webhook.headers,
                                         params=body.model_dump(),
                                         timeout=3)
                response.raise_for_status()
                logger.info(f"[{webhook.method}] response = HTTP {response.status_code} {response.text}")

            @functools.wraps(func)
            async def async_wrapper(*args, **kwargs):
                try:
                    result = await func(*args, **kwargs)
                    status = TaskStatus.success
                    error = None
                    code = ErrorCode.SUCCESS.value
                    if not result:
                        status = TaskStatus.failed
                        error = "Internal Error"
                        code = ErrorCode.BUSINESS_ERROR.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:
                    logger.info(f"result = {result}")
                    if isinstance(result, str):
                        results = [result]
                    elif isinstance(result, FFMPEGResult):
                        results = [result]
                    elif isinstance(result, List):
                        results = result
                    else:
                        results = [result]
                    webhook_with_retry(webhook=webhook,
                                       body=BaseFFMPEGTaskStatusResponse(
                                           taskId=func_id, status=status,
                                           error=error, code=code,
                                           results=results,
                                       ))
                return result

            return async_wrapper

        return decorator
