import uuid
from typing import Iterator

import sqlalchemy

import mlflow.store.db.utils
from mlflow.entities._job import Job
from mlflow.entities._job_status import JobStatus
from mlflow.exceptions import MlflowException
from mlflow.protos.databricks_pb2 import RESOURCE_DOES_NOT_EXIST
from mlflow.store.jobs.abstract_store import AbstractJobStore
from mlflow.store.tracking.dbmodels.models import SqlJob
from mlflow.utils.time import get_current_time_millis
from mlflow.utils.uri import extract_db_type_from_uri

sqlalchemy.orm.configure_mappers()

_LIST_JOB_PAGE_SIZE = 100


class SqlAlchemyJobStore(AbstractJobStore):
    """
    SQLAlchemy compliant backend store for storing Job metadata.
    This store interacts with SQL store using SQLAlchemy abstractions defined
    for MLflow Job entities.
    """

    def __init__(self, db_uri):
        """
        Create a database backed store.

        Args:
            db_uri: The SQLAlchemy database URI string to connect to the database.
        """
        super().__init__()
        self.db_uri = db_uri
        self.db_type = extract_db_type_from_uri(db_uri)
        self.engine = mlflow.store.db.utils.create_sqlalchemy_engine_with_retry(db_uri)
        if not mlflow.store.db.utils._all_tables_exist(self.engine):
            mlflow.store.db.utils._initialize_tables(self.engine)

        SessionMaker = sqlalchemy.orm.sessionmaker(bind=self.engine)
        self.ManagedSessionMaker = mlflow.store.db.utils._get_managed_session_maker(
            SessionMaker, self.db_type
        )

    def create_job(self, function_fullname: str, params: str, timeout: float | None = None) -> Job:
        """
        Create a new job with the specified function and parameters.

        Args:
            function_fullname: The full name of the function to execute
            params: The job parameters that are serialized as a JSON string
            timeout: The job execution timeout in seconds

        Returns:
            Job entity instance
        """
        with self.ManagedSessionMaker() as session:
            job_id = str(uuid.uuid4())
            creation_time = get_current_time_millis()

            job = SqlJob(
                id=job_id,
                creation_time=creation_time,
                function_fullname=function_fullname,
                params=params,
                timeout=timeout,
                status=JobStatus.PENDING.to_int(),
                result=None,
            )

            session.add(job)
            session.flush()
            return job.to_mlflow_entity()

    def start_job(self, job_id: str) -> None:
        """
        Start a job by setting its status to RUNNING.

        Args:
            job_id: The ID of the job to start
        """
        with self.ManagedSessionMaker() as session:
            job = self._get_sql_job(session, job_id)

            job.status = JobStatus.RUNNING.to_int()

    def reset_job(self, job_id: str) -> None:
        """
        Reset a job by setting its status to PENDING.

        Args:
            job_id: The ID of the job to re-enqueue.
        """
        with self.ManagedSessionMaker() as session:
            job = self._get_sql_job(session, job_id)

            job.status = JobStatus.PENDING.to_int()

    def finish_job(self, job_id: str, result: str) -> None:
        """
        Finish a job by setting its status to DONE and setting the result.

        Args:
            job_id: The ID of the job to finish
            result: The job result as a string
        """
        with self.ManagedSessionMaker() as session:
            job = self._get_sql_job(session, job_id)

            job.status = JobStatus.SUCCEEDED.to_int()
            job.result = result

    def fail_job(self, job_id: str, error: str) -> None:
        """
        Fail a job by setting its status to FAILED and setting the error message.

        Args:
            job_id: The ID of the job to fail
            error: The error message as a string
        """
        with self.ManagedSessionMaker() as session:
            job = self._get_sql_job(session, job_id)

            job.status = JobStatus.FAILED.to_int()
            job.result = error

    def mark_job_timed_out(self, job_id: str) -> None:
        """
        Set a job status to Timeout.

        Args:
            job_id: The ID of the job
        """
        with self.ManagedSessionMaker() as session:
            job = self._get_sql_job(session, job_id)

            job.status = JobStatus.TIMEOUT.to_int()

    def retry_or_fail_job(self, job_id: str, error: str) -> int | None:
        """
        If the job retry_count is less than maximum allowed retry count,
        increment the retry_count and reset the job to PENDING status,
        otherwise set the job to FAILED status and fill the job's error field.

        Args:
            job_id: The ID of the job to fail
            error: The error message as a string

        Returns:
            If the job is allowed to retry, returns the retry count,
            otherwise returns None.
        """
        from mlflow.environment_variables import MLFLOW_SERVER_JOB_TRANSIENT_ERROR_MAX_RETRIES

        max_retries = MLFLOW_SERVER_JOB_TRANSIENT_ERROR_MAX_RETRIES.get()

        with self.ManagedSessionMaker() as session:
            job = self._get_sql_job(session, job_id)

            if job.retry_count >= max_retries:
                job.status = JobStatus.FAILED.to_int()
                job.result = error
                return None
            job.retry_count += 1
            job.status = JobStatus.PENDING.to_int()
            return job.retry_count

    def list_jobs(
        self,
        function_fullname: str | None = None,
        statuses: list[JobStatus] | None = None,
        begin_timestamp: int | None = None,
        end_timestamp: int | None = None,
    ) -> Iterator[Job]:
        """
        List jobs based on the provided filters.

        Args:
            function_fullname: Filter by function full name (exact match)
            statuses: Filter by a list of job status (PENDING, RUNNING, DONE, FAILED, TIMEOUT)
            begin_timestamp: Filter jobs created after this timestamp (inclusive)
            end_timestamp: Filter jobs created before this timestamp (inclusive)

        Returns:
            Iterator of Job entities that match the filters, ordered by creation time (oldest first)
        """
        offset = 0

        while True:
            with self.ManagedSessionMaker() as session:
                # Select all columns needed for Job entity
                query = session.query(SqlJob)

                # Apply filters
                if function_fullname is not None:
                    query = query.filter(SqlJob.function_fullname == function_fullname)

                if statuses:
                    query = query.filter(
                        SqlJob.status.in_([status.to_int() for status in statuses])
                    )

                if begin_timestamp is not None:
                    query = query.filter(SqlJob.creation_time >= begin_timestamp)

                if end_timestamp is not None:
                    query = query.filter(SqlJob.creation_time <= end_timestamp)

                # Order by creation time (oldest first) and apply pagination
                jobs = (
                    query.order_by(SqlJob.creation_time)
                    .offset(offset)
                    .limit(_LIST_JOB_PAGE_SIZE)
                    .all()
                )

                # If no jobs returned, we've reached the end
                if not jobs:
                    break

                # Yield each job
                for job in jobs:
                    yield job.to_mlflow_entity()

                # If we got fewer jobs than page_size, we've reached the end
                if len(jobs) < _LIST_JOB_PAGE_SIZE:
                    break

                # Move to next page
                offset += _LIST_JOB_PAGE_SIZE

    def _get_sql_job(self, session, job_id) -> SqlJob:
        job = session.query(SqlJob).filter(SqlJob.id == job_id).one_or_none()
        if job is None:
            raise MlflowException(
                f"Job with ID {job_id} not found", error_code=RESOURCE_DOES_NOT_EXIST
            )
        return job

    def get_job(self, job_id: str) -> Job:
        """
        Get a job by its ID.

        Args:
            job_id: The ID of the job to retrieve

        Returns:
            Job entity

        Raises:
            MlflowException: If job with the given ID is not found
        """
        with self.ManagedSessionMaker() as session:
            job = self._get_sql_job(session, job_id)
            if job is None:
                raise MlflowException(
                    f"Job with ID {job_id} not found", error_code=RESOURCE_DOES_NOT_EXIST
                )
            return job.to_mlflow_entity()
