from __future__ import annotations

import json
from dataclasses import dataclass
from enum import Enum
from typing import Any, List, Optional, Union

import requests
from pydantic import BaseModel, ConfigDict

from fused._formatter.formatter_eval_result import fused_eval_result_repr
from fused.models.schema import Schema
from fused.models.udf import AnyBaseUdf


class CacheSources(str, Enum):
    MEMORY = "memory"
    OBJECT_STORAGE = "object_storage"


@dataclass
class CacheStatus:
    is_cached: bool
    cache_source: CacheSources


def is_response_cached(r: requests.Response, fused_metadata: dict | None = None):
    if not fused_metadata:
        _fused_metadata = r.headers.get("x-fused-metadata")
        fused_metadata = json.loads(_fused_metadata) if _fused_metadata else {}

    from_cached_storage = fused_metadata.get("cache") is True
    from_cached_proxy = r.headers.get("x-cache", "").lower() == "hit"
    is_cached = from_cached_storage or from_cached_proxy

    if from_cached_proxy:
        cache_source = CacheSources.MEMORY
    elif from_cached_storage:
        cache_source = CacheSources.OBJECT_STORAGE
    else:
        cache_source = None

    return CacheStatus(is_cached=is_cached, cache_source=cache_source)


class UdfEvaluationResult(BaseModel):
    data: Any = None
    sidecar: Optional[bytes] = None

    udf: Optional[AnyBaseUdf] = None
    table_schema: Optional[Schema] = None

    time_taken_seconds: float

    stdout: Optional[str] = None
    stderr: Optional[str] = None
    has_exception: bool = False
    exception_class: Optional[str] = None
    error_message: Optional[str] = None
    error_lineno: Optional[int] = None
    is_cached: bool = False
    cache_source: Optional[CacheSources] = None

    def _repr_html_(self) -> str:
        return fused_eval_result_repr(self)

    model_config = ConfigDict(arbitrary_types_allowed=True)


class MultiUdfEvaluationResult(BaseModel):
    udf_results: List[Union[UdfEvaluationResult, Any]]

    def _repr_html_(self) -> str:
        # Aggregate reprs
        result_reprs = [
            udf_result._repr_html_()
            if hasattr(udf_result, "_repr_html_")
            else repr(udf_result)
            for udf_result in self.udf_results
        ]
        return "<br><br>".join(result_reprs)

    model_config = ConfigDict(arbitrary_types_allowed=True)
