#
# Copyright 2023-2025 DataRobot, Inc. and its affiliates.
#
# All rights reserved.
#
# DataRobot, Inc.
#
# This is proprietary source code of DataRobot, Inc. and its
# affiliates.
#
# Released under the terms of DataRobot Tool and Utility Agreement.
"""

  AbstractJob
  ^        ^
  |        |
  +        +
Job       AbstractSpecificJob
            ^             ^
            |             |
            +             +
        ModelJob   AbstractBatchJob
                    ^           ^
                    |           |
                    +           +
        BatchPredictionJob    BatchMonitoringJob


"""

from __future__ import annotations

import csv
import io
import os
import threading
import time
import typing
from typing import Any, Dict, Iterable, Iterator, List, Optional, Tuple, Union

import pandas as pd
import requests
import trafaret as t

from datarobot._compat import Int, String, TypedDict
from datarobot.utils import recognize_sourcedata as orig_recognize_sourcedata

from .. import errors
from ..enums import AVAILABLE_STATEMENT_TYPES, IntakeAdapters, OutputAdapters, TrainingDataSubsets
from .dataset import Dataset
from .job import AbstractSpecificJob


class IntakeSettings(TypedDict, total=False):
    """Intake settings typed dict"""

    type: IntakeAdapters
    file: Optional[Union[str, pd.DataFrame, io.IOBase]]
    url: Optional[str]
    credential_id: Optional[str]
    data_store_id: Optional[str]
    query: Optional[str]
    table: Optional[str]
    schema: Optional[str]
    catalog: Optional[str]
    fetch_size: Optional[int]
    format: Optional[str]
    endpoint_url: Optional[str]
    dataset: Optional[Dataset]
    dataset_id: Optional[str]
    dataset_version_id: Optional[str]


class OutputSettings(TypedDict, total=False):
    """Output settings typed dict"""

    type: OutputAdapters
    path: Optional[str]
    url: Optional[str]
    credential_id: Optional[str]
    data_store_id: Optional[str]
    table: Optional[str]
    schema: Optional[str]
    catalog: Optional[str]
    statement_type: Optional[str]
    update_columns: Optional[List[str]]
    where_columns: Optional[List[str]]
    create_table_if_not_exists: Optional[bool]


class CsvSettings(TypedDict):
    delimiter: Optional[str]
    quotechar: Optional[str]
    encoding: Optional[str]


class Schedule(TypedDict):
    day_of_week: List[Union[int, str]]
    month: List[Union[int, str]]
    hour: List[Union[int, str]]
    minute: List[Union[int, str]]
    day_of_month: List[Union[int, str]]


def recognize_sourcedata(sourcedata: Any, default_fname: str) -> Dict[str, Any]:
    """Override the default recognize_sourcedata with one that converts
    DataFrame to io.BytesIO"""
    if not isinstance(sourcedata, pd.DataFrame):
        return orig_recognize_sourcedata(sourcedata, default_fname)
    # The original recognize_sourcedata will encode dataframes into StringIO
    # To pass this to requests.put, we need BytesIO - otherwise, requests will
    # try to encode it using latin-1 encoding (the default HTTP encoding),
    # which will crash when a charater not contained in latin-1 is used
    #
    # https://github.com/python/cpython/blob/b9797417315cc2d1700cb2d427685016d3380711/Lib/http/client.py#L1046
    buf = io.BytesIO()
    sourcedata.to_csv(buf, encoding="utf-8", index=False, quoting=csv.QUOTE_ALL)
    buf.seek(0)
    return {"filelike": buf, "fname": default_fname}


class AbstractBatchJob(AbstractSpecificJob):
    """
    Generic representation of batch job, that covers common functions/checks
    for predictions, and monitoring jobs.

    """

    _s3_settings = t.Dict({
        t.Key("url"): String(),
        t.Key("credential_id", optional=True): String(),
        t.Key("endpoint_url", optional=True): String(),
        t.Key("format", optional=True): String(),
    })

    _gcp_settings = t.Dict({
        t.Key("url"): String(),
        t.Key("credential_id", optional=True): String(),
        t.Key("format", optional=True): String(),
    })

    _azure_settings = t.Dict({
        t.Key("url"): String(),
        t.Key("credential_id", optional=True): String(),
        t.Key("format", optional=True): String(),
    })

    _dataset_intake_settings = t.Dict({
        t.Key("dataset"): t.Type(Dataset),
        t.Key("dataset_version_id", optional=True): String(),
    })

    _dss_intake_settings = t.Dict({
        t.Key("project_id"): String(),
        t.Key("dataset_id", optional=True): String(),
        t.Key("partition", optional=True): t.Enum(
            TrainingDataSubsets.HOLDOUT,
            TrainingDataSubsets.VALIDATION,
            TrainingDataSubsets.ALL_BACKTESTS,
        ),
    })

    _jdbc_intake_settings = t.Dict({
        t.Key("data_store_id"): String(),
        t.Key("query", optional=True): String(),
        t.Key("table", optional=True): String(),
        t.Key("schema", optional=True): String(),
        t.Key("catalog", optional=True): String(),
        t.Key("fetch_size", optional=True): Int(),
        t.Key("credential_id", optional=True): String(),
    })

    _jdbc_output_settings = t.Dict({
        t.Key("data_store_id"): String(),
        t.Key("table"): String(),
        t.Key("schema", optional=True): String(),
        t.Key("catalog", optional=True): String(),
        t.Key("statement_type"): t.Enum(
            AVAILABLE_STATEMENT_TYPES.INSERT,
            AVAILABLE_STATEMENT_TYPES.UPDATE,
            AVAILABLE_STATEMENT_TYPES.INSERT_UPDATE,
            AVAILABLE_STATEMENT_TYPES.CREATE_TABLE,
        ),
        t.Key("update_columns", optional=True): t.List(String),
        t.Key("where_columns", optional=True): t.List(String),
        t.Key("credential_id", optional=True): String(),
        t.Key("create_table_if_not_exists", optional=True): t.Bool(),
    })

    _csv_settings = t.Dict({
        t.Key("delimiter", optional=True): t.Atom("tab") | String(min_length=1, max_length=1),
        t.Key("quotechar", optional=True): String(),
        t.Key("encoding", optional=True): String(),
    })

    _datasphere_settings = t.Dict({
        t.Key("data_store_id"): String(),
        t.Key("table", optional=True): String(),
        t.Key("schema", optional=True): String(),
        t.Key("credential_id", optional=True): String(),
    })

    @classmethod
    def validate_intake_settings(cls, input_intake_settings: Optional[IntakeSettings] = None) -> IntakeSettings:
        """Validates intake settings based on type run specific trafaret check
        Creates a copy in order to avoid mutating input data

        :param input_intake_settings:
        :return: A validated copy of intake settings
        """
        if input_intake_settings is None:
            intake_settings = IntakeSettings({"type": IntakeAdapters.LOCAL_FILE})
        else:
            # avoid mutating the input argument
            intake_settings = input_intake_settings.copy()

        # Validate the intake settings
        if intake_settings.get("type") not in IntakeAdapters:
            raise ValueError("Unsupported type parameter for intake_settings: {}".format(intake_settings.get("type")))

        elif intake_settings["type"] == IntakeAdapters.LOCAL_FILE:
            # This intake option requires us to upload the source
            # data ourselves

            if intake_settings.get("file") is None:
                raise ValueError(
                    "Missing source data. Either supply the `file` "
                    "parameter or switch to an intake option that does not "
                    "require it."
                )

        elif intake_settings["type"] == IntakeAdapters.S3:
            del intake_settings["type"]
            intake_settings = cls._s3_settings.check(intake_settings)
            intake_settings["type"] = IntakeAdapters.S3

        elif intake_settings["type"] == IntakeAdapters.GCP:
            del intake_settings["type"]
            intake_settings = cls._gcp_settings.check(intake_settings)
            intake_settings["type"] = IntakeAdapters.GCP

        elif intake_settings["type"] == IntakeAdapters.AZURE:
            del intake_settings["type"]
            intake_settings = cls._azure_settings.check(intake_settings)
            intake_settings["type"] = IntakeAdapters.AZURE

        elif intake_settings["type"] == IntakeAdapters.JDBC:
            del intake_settings["type"]
            intake_settings = cls._jdbc_intake_settings.check(intake_settings)
            intake_settings["type"] = IntakeAdapters.JDBC

        elif intake_settings["type"] == IntakeAdapters.DATASET:
            del intake_settings["type"]
            intake_settings = cls._dataset_intake_settings.check(intake_settings)
            intake_settings["type"] = IntakeAdapters.DATASET

            dataset = intake_settings["dataset"]
            intake_settings["dataset_id"] = dataset.id

            if "dataset_version_id" not in intake_settings:
                intake_settings["dataset_version_id"] = dataset.version_id

            del intake_settings["dataset"]

        elif intake_settings["type"] == IntakeAdapters.DSS:
            del intake_settings["type"]
            intake_settings = cls._dss_intake_settings.check(intake_settings)
            intake_settings["type"] = IntakeAdapters.DSS

        elif intake_settings["type"] == IntakeAdapters.DATASPHERE:
            del intake_settings["type"]
            intake_settings = cls._datasphere_settings.check(intake_settings)
            intake_settings["type"] = IntakeAdapters.DATASPHERE

        return intake_settings

    @classmethod
    def validate_output_settings(cls, input_output_settings: Optional[OutputSettings] = None) -> OutputSettings:
        """Validates output settings based on type run specific trafaret check
        Creates a copy in order to avoid mutating input data

        :param input_output_settings:
        :return: A validated copy of output settings
        """
        if input_output_settings is None:
            output_settings = OutputSettings({"type": OutputAdapters.LOCAL_FILE})
        else:
            output_settings = input_output_settings.copy()

        # Validate the output settings

        if output_settings.get("type") not in OutputAdapters:
            raise ValueError("Unsupported type parameter for output_settings: {}".format(output_settings.get("type")))
        elif output_settings["type"] == OutputAdapters.LOCAL_FILE:
            output_settings["path"] = output_settings.get("path")

        elif output_settings["type"] == OutputAdapters.S3:
            del output_settings["type"]
            output_settings = cls._s3_settings.check(output_settings)
            output_settings["type"] = OutputAdapters.S3

        elif output_settings["type"] == OutputAdapters.GCP:
            del output_settings["type"]
            output_settings = cls._gcp_settings.check(output_settings)
            output_settings["type"] = OutputAdapters.GCP

        elif output_settings["type"] == OutputAdapters.AZURE:
            del output_settings["type"]
            output_settings = cls._azure_settings.check(output_settings)
            output_settings["type"] = OutputAdapters.AZURE

        elif output_settings["type"] == OutputAdapters.JDBC:
            del output_settings["type"]
            output_settings = cls._jdbc_output_settings.check(output_settings)
            output_settings["type"] = OutputAdapters.JDBC

        elif output_settings["type"] == OutputAdapters.DATASPHERE:
            del output_settings["type"]
            output_settings = cls._datasphere_settings.check(output_settings)
            output_settings["type"] = OutputAdapters.DATASPHERE

        return output_settings

    @classmethod
    def _upload_intake_file(
        cls,
        intake_file: Dict[str, Any],
        intake_settings: IntakeSettings,
        job_response: Dict[str, Any],
        upload_read_timeout: int,
        output_file: Optional[typing.IO[Any]] = None,
    ) -> Optional[threading.Thread]:
        """Uploads intake file to running job, when intake type is localFile"""

        if intake_settings["type"] != IntakeAdapters.LOCAL_FILE:
            return None

        # There is source data to upload, so spin up a thread to handle
        # the upload concurrently and for thread safety issues, make
        # a copy of the REST client object
        _upload_client = cls._client.copy()
        job_csv_settings = job_response.get("jobSpec", {}).get("csvSettings", {})
        upload_thread = None

        def _get_file_size(fileobj: typing.BinaryIO) -> int:
            # To cover both files and filelike obj utilize .tell
            cur = fileobj.tell()
            fileobj.seek(0, os.SEEK_END)
            file_size = fileobj.tell()
            fileobj.seek(cur)

            return file_size

        # pylint: disable-next=unused-argument
        def _create_csv_chunk(
            header: Iterable[Any],
            reader: Iterator[Any],
            max_size: int,
            delimiter: str,
            encoding: str,
            quotechar: str,
        ) -> Tuple[io.StringIO, int]:
            chunk = io.StringIO()
            bytes_written = 0
            writer = csv.writer(chunk, delimiter=delimiter, quotechar=quotechar)
            writer.writerow(header)
            while bytes_written < max_size:
                try:
                    csv_chunk_content = next(reader)
                    written = writer.writerow(csv_chunk_content)
                    bytes_written += written
                except StopIteration:
                    break

            return chunk, bytes_written

        def _fileobj_to_csv_stream(fileobj: typing.BinaryIO, encoding: str) -> Iterator[Any]:
            stream = io.TextIOWrapper(fileobj, encoding=encoding)

            yield from csv.reader(stream)

            stream.close()

        def _upload_multipart(fileobj: typing.BinaryIO, base_upload_url: str) -> None:
            is_async = intake_settings.get("async", True)
            MAX_RETRY = 1 if is_async else 3
            MB_PER_CHUNK = 5
            CHUNK_MAX_SIZE = MB_PER_CHUNK * 1024 * 1024

            delimiter = job_csv_settings.get("delimiter", ",")
            encoding = job_csv_settings.get("encoding", "utf-8")
            quotechar = job_csv_settings.get("quotechar", '"')

            file_size = _get_file_size(fileobj)
            csv_stream = _fileobj_to_csv_stream(fileobj, encoding)

            # grab the header so it can be added on all parts
            header = next(csv_stream)

            part_number = 0
            bytes_written = 0
            while bytes_written <= file_size:
                part_upload_url = f"{base_upload_url}part/{part_number}/"

                # Read the inputfile in chunks of CHUNK_MAX_SIZE
                # Then call put multiple times increasing the part_number each time
                chunk, chunk_bytes = _create_csv_chunk(
                    header, csv_stream, CHUNK_MAX_SIZE, delimiter, encoding, quotechar
                )
                bytes_written += chunk_bytes
                if chunk_bytes == 0:
                    break

                for attempts in range(MAX_RETRY):
                    try:
                        chunk.seek(0)
                        response = _upload_client.put(
                            url=part_upload_url,
                            data=chunk,
                            headers={"content-type": "text/csv"},
                            timeout=(_upload_client.connect_timeout, upload_read_timeout),
                        )

                        # Success! don't retry
                        if response.status_code == 202:
                            chunk.close()
                            break
                    except (requests.exceptions.ConnectionError, requests.exceptions.Timeout):
                        attempts += 1
                        if attempts == MAX_RETRY:
                            raise

                part_number += 1

            # finalize the upload to indicate no more data is arriving
            _upload_client.post(url=f"{base_upload_url}finalizeMultipart/")

        def _uploader() -> None:
            upload_url = job_response["links"]["csvUpload"]

            if "file_path" in intake_file:
                fileobj = open(  # pylint: disable=consider-using-with
                    intake_file["file_path"], "rb"
                )
            else:
                fileobj = intake_file["filelike"]
                fileobj.seek(0)

            try:
                if intake_settings.get("multipart"):
                    _upload_multipart(fileobj, upload_url)

                else:
                    _upload_client.put(
                        url=upload_url,
                        data=fileobj,
                        headers={"content-type": "text/csv"},
                        timeout=(_upload_client.connect_timeout, upload_read_timeout),
                    )
            finally:
                if hasattr(fileobj, "close"):
                    fileobj.close()

        if output_file is not None:
            # If output_file is specified, we upload and download
            # concurrently

            upload_thread = threading.Thread(target=_uploader)
            upload_thread.daemon = True
            upload_thread.start()

        else:
            # Otherwise, upload is synchronous

            _uploader()
        return upload_thread

    @classmethod
    def _job_path(cls, project_id: str, job_id: str) -> str:
        raise NotImplementedError

    def _delete(self, ignore_404_errors: bool = False) -> None:
        """
        Cancel this job. If this job has not finished running, it will be
        removed and canceled.
        """
        status = self._get_status()

        prediction_job_id = status["links"]["self"].split("/")[-2]
        try:
            self._client.delete(self._job_path(project_id="", job_id=prediction_job_id))
        except errors.ClientError as exc:
            if exc.status_code == 404 and ignore_404_errors:
                return
            raise

    def _get_status(self) -> Any:
        """Get status of batch job"""
        batch_job = super().get(project_id="", job_id=self.id)
        batch_job.id = self.id

        return batch_job._safe_data

    def _download(self, fileobj: typing.IO[Any], timeout: int = 120, read_timeout: int = 660) -> None:
        """Downloads the CSV result of a batch job"""

        status = self._wait_for_download(timeout=timeout)
        download_iter = self._client.get(
            status["links"]["download"],
            stream=True,
            timeout=read_timeout,
        ).iter_content(chunk_size=8192)

        for chunk in download_iter:
            if chunk:
                fileobj.write(chunk)

        # Check if job was aborted during download (and the download is incomplete)
        status = self._get_status()
        if status["status"] in ("ABORTED", "FAILED"):
            raise RuntimeError("Job {} was aborted: {}".format(self.id, status["status_details"]))

    def _wait_for_download(self, timeout: int = 120) -> Any:
        """Waits for download to become available"""
        start = time.time()
        status = None
        while True:
            status = self._get_status()

            output_adapter_type = status["job_spec"].get("output_settings", {}).get("type")
            if output_adapter_type and not output_adapter_type == OutputAdapters.LOCAL_FILE:
                raise RuntimeError(
                    (
                        "You cannot download predictions from jobs that did not use local_file as "
                        "the output adapter. Job with ID {} had the output adapter defined as {}."
                    ).format(self.id, output_adapter_type)
                )

            if status["status"] in ("ABORTED", "FAILED"):
                raise RuntimeError("Job {} was aborted: {}".format(self.id, status["status_details"]))

            if "download" in status["links"]:
                break

            if timeout >= 0 and time.time() - start > timeout:  # pylint: disable=chained-comparison
                break

            time.sleep(1)

        if "download" not in status["links"]:
            # Ignore 404 errors here if the job never started - then we can't abort it
            self._delete(ignore_404_errors=True)
            raise RuntimeError(
                (
                    "Timed out waiting for download to become available for job ID {}. "
                    "Other jobs may be occupying the queue. Consider raising the timeout."
                ).format(self.id)
            )

        return status
