from datetime import datetime
from typing import (
    Any,
    Dict,
    Generic,
    List,
    Mapping,
    Optional,
    Tuple,
    Type,
    TypeVar,
)

from kumoapi.common import ValidationResponse
from kumoapi.jobs import (
    ArtifactExportRequest,
    ArtifactExportResponse,
    AutoTrainerProgress,
    BaselineJobRequest,
    BaselineJobResource,
    BatchPredictionJobResource,
    BatchPredictionRequest,
    CancelBatchPredictionJobResponse,
    CancelTrainingJobResponse,
    ErrorDetails,
    GeneratePredictionTableJobResource,
    GeneratePredictionTableRequest,
    GenerateTrainTableJobResource,
    GetPredictionsDfUrlResponse,
    JobRequestBase,
    JobResourceBase,
    JobStatus,
    PredictionProgress,
    TrainingJobRequest,
    TrainingJobResource,
    TrainingTableSpec,
)
from kumoapi.json_serde import from_json, to_json_dict
from kumoapi.source_table import LLMRequest, LLMResponse, SourceTableType
from kumoapi.train import TrainingStage
from typing_extensions import override

from kumoai.client import KumoClient
from kumoai.client.utils import (
    Returns,
    parse_patch_response,
    parse_response,
    raise_on_error,
)

TrainingJobID = str
BatchPredictionJobID = str
GenerateTrainTableJobID = str
GeneratePredictionTableJobID = str
LLMJobId = str
BaselineJobID = str

JobRequestType = TypeVar('JobRequestType', bound=JobRequestBase)
JobResourceType = TypeVar('JobResourceType', bound=JobResourceBase)


class CommonJobAPI(Generic[JobRequestType, JobResourceType]):
    def __init__(self, client: KumoClient, base_endpoint: str,
                 res_type: Type[JobResourceType]) -> None:
        self._client = client
        self._base_endpoint = base_endpoint
        self._res_type = res_type

    def create(self, request: JobRequestType) -> str:
        response = self._client._post(self._base_endpoint,
                                      json=to_json_dict(request))
        raise_on_error(response)
        return parse_response(Dict[str, str], response)['id']

    def get(self, id: str) -> JobResourceType:
        response = self._client._get(f'{self._base_endpoint}/{id}')
        raise_on_error(response)
        return parse_response(self._res_type, response)

    def list(
        self,
        *,
        pquery_name: Optional[str] = None,
        pquery_id: Optional[str] = None,
        job_status: Optional[JobStatus] = None,
        limit: Optional[int] = None,
        additional_tags: Mapping[str, str] = {},
    ) -> List[JobResourceType]:
        params: Dict[str, Any] = {
            'pquery_name': pquery_name,
            'pquery_id': pquery_id,
            'job_status': job_status,
            'limit': limit
        }
        params.update(additional_tags)
        response = self._client._get(self._base_endpoint, params=params)
        raise_on_error(response)
        resource_elements = response.json()
        assert isinstance(resource_elements, list)
        return [from_json(e, self._res_type) for e in resource_elements]

    def delete_tags(self, job_id: str, tags: List[str]) -> bool:
        r"""Removes the tags from the job.

        Args:
            job_id (str): The ID of the job.
            tags (List[str]): The tags to remove.
        """
        return self.update_tags(job_id, {t: 'none' for t in tags})

    def update_tags(self, job_id: str,
                    custom_job_tags: Mapping[str, Optional[str]]) -> bool:
        r"""Updates the tags of the job.

        Args:
            job_id (str): The ID of the job.
            custom_job_tags (Mapping[str, Optional[str]]): The tags to update.
                Note that the value 'none' will remove the tag. If the tag is
                not present, it will be added.
        """
        response = self._client._patch(
            f'{self._base_endpoint}/{job_id}/tags',
            data=None,
            params={
                k: str(v)
                for k, v in custom_job_tags.items()
            },
        )
        raise_on_error(response)
        return parse_patch_response(response)


class BaselineJobAPI(CommonJobAPI[BaselineJobRequest, BaselineJobResource]):
    r"""Typed API definition for the baseline job resource."""
    def __init__(self, client: KumoClient) -> None:
        super().__init__(client, '/baseline_jobs', BaselineJobResource)


class TrainingJobAPI(CommonJobAPI[TrainingJobRequest, TrainingJobResource]):
    r"""Typed API definition for the training job resource."""
    def __init__(self, client: KumoClient) -> None:
        super().__init__(client, '/training_jobs', TrainingJobResource)

    def get_progress(self, id: TrainingJobID) -> AutoTrainerProgress:
        response = self._client._get(f'{self._base_endpoint}/{id}/progress')
        raise_on_error(response)
        return parse_response(AutoTrainerProgress, response)

    def holdout_data_url(self, id: TrainingJobID,
                         presigned: bool = True) -> str:
        response = self._client._get(f'{self._base_endpoint}/{id}/holdout',
                                     params={'presigned': presigned})
        raise_on_error(response)
        return response.text

    def cancel(self, id: str) -> CancelTrainingJobResponse:
        response = self._client._post(f'{self._base_endpoint}/{id}/cancel')
        raise_on_error(response)
        return parse_response(CancelTrainingJobResponse, response)


class BatchPredictionJobAPI(CommonJobAPI[BatchPredictionRequest,
                                         BatchPredictionJobResource]):
    r"""Typed API definition for the prediction job resource."""
    def __init__(self, client: KumoClient) -> None:
        super().__init__(client, '/prediction_jobs',
                         BatchPredictionJobResource)

    @override
    def create(self, request: BatchPredictionRequest) -> str:
        # TODO(manan): eventually, all `create` methods should
        # return a validation response:
        raise NotImplementedError

    def maybe_create(
        self, request: BatchPredictionRequest
    ) -> Tuple[Optional[str], ValidationResponse]:
        response = self._client._post(self._base_endpoint,
                                      json=to_json_dict(request))
        raise_on_error(response)
        return parse_response(
            Returns[Tuple[Optional[str], ValidationResponse]], response)

    def list(
        self,
        *,
        model_id: Optional[TrainingJobID] = None,
        pquery_name: Optional[str] = None,
        pquery_id: Optional[str] = None,
        job_status: Optional[JobStatus] = None,
        limit: Optional[int] = None,
        additional_tags: Mapping[str, str] = {},
    ) -> List[BatchPredictionJobResource]:
        if model_id:
            additional_tags = {**additional_tags, 'model_id': model_id}
        return super().list(pquery_name=pquery_name, pquery_id=pquery_id,
                            job_status=job_status, limit=limit,
                            additional_tags=additional_tags)

    def get_progress(self, id: str) -> PredictionProgress:
        response = self._client._get(f'{self._base_endpoint}/{id}/progress')
        raise_on_error(response)
        return parse_response(PredictionProgress, response)

    def cancel(self, id: str) -> CancelBatchPredictionJobResponse:
        response = self._client._post(f'{self._base_endpoint}/{id}/cancel')
        raise_on_error(response)
        return parse_response(CancelBatchPredictionJobResponse, response)

    def get_batch_predictions_url(self, id: str) -> List[str]:
        """Returns presigned url pointing to the location where the
        predictions are made. Depending on the environment where
        this is run, it could be a s3 path, a path in Snowflake stage
        or a Databricks UC volume path.

        Arguments:
        id (str): ID of the batch predictions job for which predictions
                are requested
        """
        response = self._client._get(
            f'{self._base_endpoint}/{id}/get_prediction_df_urls')
        raise_on_error(response)
        return parse_response(
            GetPredictionsDfUrlResponse,
            response,
        ).prediction_partitions


class GenerateTrainTableJobAPI(CommonJobAPI[GeneratePredictionTableRequest,
                                            GenerateTrainTableJobResource]):
    r"""Typed API definition for training table generation job resource."""
    def __init__(self, client: KumoClient) -> None:
        super().__init__(client, '/gentraintable_jobs',
                         GenerateTrainTableJobResource)

    def get_table_data(self, id: GenerateTrainTableJobID,
                       presigned: bool = True) -> List[str]:
        """Return a list of URLs to access train table parquet data.
        There might be multiple URLs if the table data is partitioned into
        multiple files.
        """
        return self._get_table_data(id, presigned)

    def _get_table_data(self, id: GenerateTrainTableJobID,
                        presigned: bool = True,
                        raw_path: bool = False) -> List[str]:
        """Helper function to get train table data."""
        # Raw path to get local file path instead of SPCS stage path
        params: Dict[str, Any] = {'presigned': presigned, 'raw_path': raw_path}
        resp = self._client._get(f'{self._base_endpoint}/{id}/table_data',
                                 params=params)
        raise_on_error(resp)
        return parse_response(List[str], resp)

    def get_split_masks(
            self, id: GenerateTrainTableJobID) -> Dict[TrainingStage, str]:
        """Return a dictionary of presigned URLs keyed by training stage.
        Each URL points to a torch-serialized (default pickle protocol) file of
        the mask tensor for that training stage.

        Example:
            >>> # code to load a mask tensor:
            >>> import io
            >>> import torch
            >>> import requests
            >>> masks = get_split_masks('some-gen-traintable-job-id')
            >>> data_bytes = requests.get(masks[TrainingStage.TEST]).content
            >>> test_mask_tensor = torch.load(io.BytesIO(data))
        """
        resp = self._client._get(f'{self._base_endpoint}/{id}/split_masks')
        raise_on_error(resp)
        return parse_response(Dict[TrainingStage, str], resp)

    def get_progress(self, id: str) -> Dict[str, int]:
        response = self._client._get(f'{self._base_endpoint}/{id}/progress')
        raise_on_error(response)
        return parse_response(Dict[str, int], response)

    def cancel(self, id: str) -> None:
        response = self._client._post(f'{self._base_endpoint}/{id}/cancel')
        raise_on_error(response)

    def validate_custom_train_table(
        self,
        id: str,
        source_table_type: SourceTableType,
        train_table_mod: TrainingTableSpec,
    ) -> ValidationResponse:
        response = self._client._post(
            f'{self._base_endpoint}/{id}/validate_custom_train_table',
            json=to_json_dict({
                'custom_table': source_table_type,
                'train_table_mod': train_table_mod,
            }),
        )
        return parse_response(ValidationResponse, response)

    def get_job_error(self, id: str) -> ErrorDetails:
        """Thin API wrapper for fetching errors from the jobs.

        Arguments:
        id (str): Id of the job whose related errors are expected to be
            queried.
        """
        response = self._client._get(f'{self._base_endpoint}/{id}/get_errors')
        raise_on_error(response)
        return parse_response(ErrorDetails, response)


class GeneratePredictionTableJobAPI(
        CommonJobAPI[GeneratePredictionTableRequest,
                     GeneratePredictionTableJobResource]):
    r"""Typed API definition for prediction table generation job resource."""
    def __init__(self, client: KumoClient) -> None:
        super().__init__(client, '/genpredtable_jobs',
                         GeneratePredictionTableJobResource)

    def get_anchor_time(self, id: BatchPredictionJobID) -> Optional[datetime]:
        response = self._client._get(
            f'{self._base_endpoint}/{id}/get_anchor_time')
        raise_on_error(response)
        return parse_response(Returns[Optional[datetime]], response)

    def get_table_data(self, id: GeneratePredictionTableJobID,
                       presigned: bool = True) -> List[str]:
        """Return a list of URLs to access prediction table parquet data.
        There might be multiple URLs if the table data is partitioned into
        multiple files.
        """
        params: Dict[str, Any] = {'presigned': presigned}
        resp = self._client._get(f'{self._base_endpoint}/{id}/table_data',
                                 params=params)
        raise_on_error(resp)
        return parse_response(List[str], resp)

    def cancel(self, id: str) -> None:
        response = self._client._post(f'{self._base_endpoint}/{id}/cancel')
        raise_on_error(response)

    def get_job_error(self, id: str) -> ErrorDetails:
        """Thin API wrapper for fetching errors from the jobs.

        Arguments:
        id (str): Id of the job whose related errors are expected to be
            queried.
        """
        response = self._client._get(f'{self._base_endpoint}/{id}/get_errors')
        raise_on_error(response)
        return parse_response(ErrorDetails, response)


class LLMJobAPI:
    r"""Typed API definition for LLM job resource."""
    def __init__(self, client: KumoClient) -> None:
        self._client = client
        self._base_endpoint = '/llm_embedding_job'

    def create(self, request: LLMRequest) -> LLMJobId:
        response = self._client._post(
            self._base_endpoint,
            json=to_json_dict(request),
        )
        raise_on_error(response)
        return parse_response(LLMResponse, response).job_id

    def get(self, id: LLMJobId) -> JobStatus:
        response = self._client._get(f'{self._base_endpoint}/status/{id}')
        raise_on_error(response)
        return parse_response(JobStatus, response)

    def cancel(self, id: LLMJobId) -> JobStatus:
        response = self._client._delete(f'{self._base_endpoint}/cancel/{id}')
        return response


class ArtifactExportJobAPI:
    r"""Typed API definition for artifact export job resource."""
    def __init__(self, client: KumoClient) -> None:
        self._client = client
        self._base_endpoint = '/artifact'

    def create(self, request: ArtifactExportRequest) -> str:
        response = self._client._post(
            self._base_endpoint,
            json=to_json_dict(request),
        )
        raise_on_error(response)
        return parse_response(ArtifactExportResponse, response).job_id

    # TODO Add an API in artifact export to get
    # JobStatusReport and not just JobStatus
    def get(self, id: str) -> JobStatus:
        response = self._client._get(f'{self._base_endpoint}/{id}')
        raise_on_error(response)
        return parse_response(JobStatus, response)

    def cancel(self, id: str) -> JobStatus:
        response = self._client._post(f'{self._base_endpoint}/{id}/cancel')
        raise_on_error(response)
        return parse_response(JobStatus, response)
