import gzip
from io import BytesIO
from typing import Optional, Union, List, Any

from pydantic import BaseModel, AnyUrl

from apipod.compatibility.upload import is_param_media_toolkit_file
from apipod.core.job.base_job import JOB_STATUS, BaseJob
from apipod.settings import DEFAULT_DATE_TIME_FORMAT
from media_toolkit import IMediaContainer
from media_toolkit.utils.data_type_utils import is_file_model_dict


class FileModel(BaseModel):
    file_name: str
    content_type: str
    content: Union[str, AnyUrl]  # base64 encoded or url
    max_size_mb: Optional[float] = 4000

    class Config:
        json_schema_extra = {
            "x-media-type": "MediaFile",
            "example": {
                "file_name": "example.csv",
                "content_type": "text/csv",
                "content": "https://example.com/example.csv"   # bytes, base64 encoded or url
            }
        }


class ImageFileModel(FileModel):
    class Config:
        json_schema_extra = {
            "x-media-type": "ImageFile",
            "example": {
                "file_name": "example.png",
                "content_type": "image/png",
                "content": "base64 encoded image data"
            }
        }


class AudioFileModel(FileModel):
    class Config:
        json_schema_extra = {
            "x-media-type": "AudioFile",
            "example": {
                "file_name": "example.mp3",
                "content_type": "audio/mpeg",
                "content": "base64 encoded audio data"
            }
        }


class VideoFileModel(FileModel):
    class Config:
        json_schema_extra = {
            "x-media-type": "VideoFile",
            "example": {
                "file_name": "example.mp4",
                "content_type": "video/mp4",
                "content": "base64 encoded video data"
            }
        }


class JobProgress(BaseModel):
    progress: float = 0.0
    message: Optional[str] = None


class JobResult(BaseModel):
    """
    When the user (client) sends a request to an Endpoint, a ClientJob is created.
    This job contains the information about the request and the response.
    """
    id: str
    status: Optional[str] = None
    progress: Optional[JobProgress] = None
    error: Optional[str] = None
    result: Union[FileModel, List[FileModel], List, str, Any, None] = None
    refresh_job_url: Optional[str] = None
    cancel_job_url: Optional[str] = None

    created_at: Optional[str] = None
    queued_at: Optional[str] = None
    execution_started_at: Optional[str] = None
    execution_finished_at: Optional[str] = None

    endpoint_protocol: Optional[str] = "socaity"


class JobResultFactory:

    @staticmethod
    def _serialize_result(data: Any) -> Union[FileModel, List[FileModel], List, str, None]:
        # Handle single MediaList or MediaFile
        if isinstance(data, IMediaContainer):
            return data.to_json()

        if is_param_media_toolkit_file(data):
            return FileModel(**data.to_json())

        if isinstance(data, FileModel):
            return data

        if is_file_model_dict(data):
            return FileModel(**data)

        # Handle list of MediaLists/MediaFiles/other items
        if isinstance(data, list):
            return [JobResultFactory._serialize_result(item) for item in data]

        if isinstance(data, dict):
            return {key: JobResultFactory._serialize_result(value) for key, value in data.items()}

        return data

    @staticmethod
    def from_base_job(ij: BaseJob) -> JobResult:
        def format_date(date):
            return date.strftime(DEFAULT_DATE_TIME_FORMAT) if date else None

        created_at = format_date(ij.created_at)
        queued_at = format_date(ij.queued_at)
        execution_started_at = format_date(ij.execution_started_at)
        execution_finished_at = format_date(ij.execution_finished_at)

        # if the internal job returned a media-toolkit file, convert it to a json serializable FileModel
        result = ij.result
        result = JobResultFactory._serialize_result(result)

        # Job_status is an Enum, convert it to a string to return it as json
        status = ij.status
        if isinstance(status, JOB_STATUS):
            status = status.value

        try:
            jp = JobProgress(progress=ij.job_progress._progress, message=ij.job_progress._message)
        except Exception:
            jp = JobProgress(progress=0.0, message='')

        return JobResult(
            id=ij.id,
            status=status,
            progress=jp,
            error=ij.error,
            result=result,
            created_at=created_at,
            queued_at=queued_at,
            execution_started_at=execution_started_at,
            execution_finished_at=execution_finished_at
        )

    @staticmethod
    def gzip_job_result(job_result: JobResult) -> bytes:
        job_result_bytes = job_result.json().encode('utf-8')
        # Compress the serialized bytes with gzip
        gzip_buffer = BytesIO()
        with gzip.GzipFile(fileobj=gzip_buffer, mode='wb') as gzip_file:
            gzip_file.write(job_result_bytes)

        # Retrieve the gzipped data
        return gzip_buffer.getvalue()

    @staticmethod
    def job_not_found(job_id: str) -> JobResult:
        return JobResult(
            id=job_id,
            status=JOB_STATUS.FAILED,
            error="Job not found."
        )