from __future__ import annotations

import datetime
import traceback
from typing import TYPE_CHECKING, Any, Self
from uuid import UUID, uuid4

from opentelemetry import trace
from opentelemetry.semconv._incubating.attributes.code_attributes import (
    CODE_NAMESPACE,
)
from opentelemetry.semconv._incubating.attributes.messaging_attributes import (
    MESSAGING_CONSUMER_GROUP_NAME,
    MESSAGING_DESTINATION_NAME,
    MESSAGING_MESSAGE_ID,
    MESSAGING_OPERATION_NAME,
    MESSAGING_OPERATION_TYPE,
    MESSAGING_SYSTEM,
    MessagingOperationTypeValues,
)
from opentelemetry.semconv.attributes.error_attributes import ERROR_TYPE
from opentelemetry.trace import Link, SpanContext, SpanKind

from plain import postgres
from plain.logs import get_framework_logger
from plain.postgres import transaction, types
from plain.postgres.expressions import F
from plain.runtime import settings
from plain.utils import timezone

from .exceptions import DeferError, DeferJob
from .registry import jobs_registry

if TYPE_CHECKING:
    from .jobs import Job

__all__ = ["JobRequest", "JobProcess", "JobResult", "JobResultStatuses"]

logger = get_framework_logger()
tracer = trace.get_tracer("plain.jobs")


@postgres.register_model
class JobRequest(postgres.Model):
    """
    Keep all pending job requests in a single table.
    """

    created_at: datetime.datetime = types.DateTimeField(auto_now_add=True)
    uuid: UUID = types.UUIDField(default=uuid4)

    job_class: str = types.TextField(max_length=255)
    parameters: dict[str, Any] | None = types.JSONField(required=False, allow_null=True)
    priority: int = types.SmallIntegerField(default=0)
    source: str = types.TextField(required=False)
    queue: str = types.TextField(default="default", max_length=255)

    retries: int = types.SmallIntegerField(default=0)
    retry_attempt: int = types.SmallIntegerField(default=0)

    concurrency_key: str = types.TextField(max_length=255, required=False)

    start_at: datetime.datetime | None = types.DateTimeField(
        required=False, allow_null=True
    )

    # OpenTelemetry trace context
    trace_id: str | None = types.TextField(
        max_length=34, required=False, allow_null=True
    )
    span_id: str | None = types.TextField(
        max_length=18, required=False, allow_null=True
    )

    # expires_at = postgres.DateTimeField(required=False, allow_null=True)

    query: postgres.QuerySet[JobRequest] = postgres.QuerySet()

    model_options = postgres.Options(
        ordering=["-priority", "-created_at"],
        indexes=[
            postgres.Index(
                name="plainjobs_jobrequest_priority_idx", fields=["priority"]
            ),
            postgres.Index(
                name="plainjobs_jobrequest_created_at_idx", fields=["created_at"]
            ),
            postgres.Index(name="plainjobs_jobrequest_queue_idx", fields=["queue"]),
            postgres.Index(
                name="plainjobs_jobrequest_start_at_idx", fields=["start_at"]
            ),
            postgres.Index(
                name="plainjobs_jobrequest_concurrency_key_idx",
                fields=["concurrency_key"],
            ),
            postgres.Index(
                name="plainjobs_jobrequest_trace_id_idx", fields=["trace_id"]
            ),
            postgres.Index(name="plainjobs_jobrequest_uuid_idx", fields=["uuid"]),
            # Used for job grouping queries
            postgres.Index(
                name="job_request_concurrency_key",
                fields=["job_class", "concurrency_key"],
            ),
        ],
        constraints=[
            postgres.UniqueConstraint(
                fields=["uuid"], name="plainjobs_jobrequest_unique_uuid"
            ),
        ],
    )

    def __str__(self) -> str:
        return f"{self.job_class} [{self.uuid}]"

    def convert_to_job_process(self) -> JobProcess:
        """
        JobRequests are the pending jobs that are waiting to be executed.
        We immediately convert them to JobProcess when they are picked up.
        """
        with transaction.atomic():
            result = JobProcess.query.create(
                job_request_uuid=self.uuid,
                job_class=self.job_class,
                parameters=self.parameters,
                priority=self.priority,
                source=self.source,
                queue=self.queue,
                retries=self.retries,
                retry_attempt=self.retry_attempt,
                concurrency_key=self.concurrency_key,
                trace_id=self.trace_id,
                span_id=self.span_id,
            )

            # Delete the pending JobRequest now
            self.delete()

        return result


class JobQuerySet(postgres.QuerySet["JobProcess"]):
    def running(self) -> Self:
        return self.filter(started_at__isnull=False)

    def waiting(self) -> Self:
        return self.filter(started_at__isnull=True)

    def mark_lost_jobs(self) -> None:
        # Lost jobs are jobs that have been pending for too long,
        # and probably never going to get picked up by a worker process.
        # In theory we could save a timeout per-job and mark them timed-out more quickly,
        # but if they're still running, we can't actually send a signal to cancel it...
        now = timezone.now()
        cutoff = now - datetime.timedelta(seconds=settings.JOBS_TIMEOUT)
        lost_jobs = self.filter(
            created_at__lt=cutoff
        )  # Doesn't matter whether it started or not -- it shouldn't take this long.

        # Note that this will save it in the results,
        # but lost jobs are only retried if they have a retry!
        for job in lost_jobs:
            job.convert_to_result(status=JobResultStatuses.LOST)


@postgres.register_model
class JobProcess(postgres.Model):
    """
    All active jobs are stored in this table.
    """

    uuid: UUID = types.UUIDField(default=uuid4)
    created_at: datetime.datetime = types.DateTimeField(auto_now_add=True)
    started_at: datetime.datetime | None = types.DateTimeField(
        required=False, allow_null=True
    )

    # From the JobRequest
    job_request_uuid: UUID = types.UUIDField()
    job_class: str = types.TextField(max_length=255)
    parameters: dict[str, Any] | None = types.JSONField(required=False, allow_null=True)
    priority: int = types.SmallIntegerField(default=0)
    source: str = types.TextField(required=False)
    queue: str = types.TextField(default="default", max_length=255)
    retries: int = types.SmallIntegerField(default=0)
    retry_attempt: int = types.SmallIntegerField(default=0)
    concurrency_key: str = types.TextField(max_length=255, required=False)

    # OpenTelemetry trace context
    trace_id: str | None = types.TextField(
        max_length=34, required=False, allow_null=True
    )
    span_id: str | None = types.TextField(
        max_length=18, required=False, allow_null=True
    )

    query: JobQuerySet = JobQuerySet()

    model_options = postgres.Options(
        ordering=["-created_at"],
        indexes=[
            postgres.Index(
                name="plainjobs_jobprocess_created_at_idx", fields=["created_at"]
            ),
            postgres.Index(name="plainjobs_jobprocess_queue_idx", fields=["queue"]),
            postgres.Index(
                name="plainjobs_jobprocess_concurrency_key_idx",
                fields=["concurrency_key"],
            ),
            postgres.Index(
                name="plainjobs_jobprocess_started_at_idx", fields=["started_at"]
            ),
            postgres.Index(
                name="plainjobs_jobprocess_job_request_uuid_idx",
                fields=["job_request_uuid"],
            ),
            postgres.Index(
                name="plainjobs_jobprocess_trace_id_idx", fields=["trace_id"]
            ),
            postgres.Index(name="plainjobs_jobprocess_uuid_idx", fields=["uuid"]),
            # Used for job grouping queries
            postgres.Index(
                name="job_concurrency_key",
                fields=["job_class", "concurrency_key"],
            ),
        ],
        constraints=[
            postgres.UniqueConstraint(
                fields=["uuid"], name="plainjobs_job_unique_uuid"
            ),
        ],
    )

    def revert_to_job_request(self) -> JobRequest:
        """Undo convert_to_job_process — put the job back in the request queue."""
        with transaction.atomic():
            job_request = JobRequest.query.create(
                uuid=self.job_request_uuid,
                job_class=self.job_class,
                parameters=self.parameters,
                priority=self.priority,
                source=self.source,
                queue=self.queue,
                retries=self.retries,
                retry_attempt=self.retry_attempt,
                concurrency_key=self.concurrency_key,
                trace_id=self.trace_id,
                span_id=self.span_id,
            )
            self.delete()
        return job_request

    def run(self) -> JobResult:
        links = []
        if self.trace_id and self.span_id:
            try:
                links.append(
                    Link(
                        SpanContext(
                            trace_id=int(self.trace_id, 16),
                            span_id=int(self.span_id, 16),
                            is_remote=True,
                        )
                    )
                )
            except (ValueError, TypeError):
                logger.warning(
                    "Invalid trace context for job",
                    extra={"job_uuid": self.uuid},
                )

        with (
            tracer.start_as_current_span(
                f"run {self.job_class}",
                kind=SpanKind.CONSUMER,
                attributes={
                    MESSAGING_SYSTEM: "plain.jobs",
                    MESSAGING_OPERATION_TYPE: MessagingOperationTypeValues.PROCESS.value,
                    MESSAGING_OPERATION_NAME: "run",
                    MESSAGING_MESSAGE_ID: str(self.uuid),
                    MESSAGING_DESTINATION_NAME: self.queue,
                    MESSAGING_CONSUMER_GROUP_NAME: self.queue,  # Workers consume from specific queues
                    CODE_NAMESPACE: self.job_class,
                },
                links=links,
            ) as span
        ):
            # This is how we know it has been picked up
            self.started_at = timezone.now()
            self.save(update_fields=["started_at"])

            try:
                job = jobs_registry.load_job(self.job_class, self.parameters or {})
                job.job_process = self

                try:
                    job.run()
                except DeferJob as e:
                    # Job deferred - not an error, log at INFO level
                    logger.info(
                        "Job deferred",
                        extra={
                            "delay": e.delay,
                            "increment_retries": e.increment_retries,
                            "job_class": self.job_class,
                            "job_process_uuid": self.uuid,
                        },
                    )
                    span.set_attribute(ERROR_TYPE, "DeferJob")
                    span.set_status(trace.StatusCode.OK)  # Not an error
                    return self.defer(job=job, defer_exception=e)

                # Success case (only reached if no DeferJob was raised)
                span.set_status(trace.StatusCode.OK)
                return self.convert_to_result(status=JobResultStatuses.SUCCESSFUL)

            except DeferError as e:
                # Defer failed (e.g., concurrency limit reached during re-enqueue)
                # The transaction was rolled back, so the JobProcess still exists in DB.
                # The pk was restored in defer() before raising, so we can proceed normally.
                logger.warning(
                    "Defer failed",
                    extra={"job_class": self.job_class, "error": str(e)},
                )
                span.record_exception(e)
                span.set_status(trace.Status(trace.StatusCode.ERROR, str(e)))
                span.set_attribute(ERROR_TYPE, type(e).__name__)
                return self.convert_to_result(
                    status=JobResultStatuses.ERRORED,
                    error=str(e),
                )

            except Exception as e:
                logger.exception(e)
                span.record_exception(e)
                span.set_status(trace.Status(trace.StatusCode.ERROR, str(e)))
                span.set_attribute(ERROR_TYPE, type(e).__name__)
                return self.convert_to_result(
                    status=JobResultStatuses.ERRORED,
                    error="".join(traceback.format_tb(e.__traceback__)),
                )

    def defer(self, *, job: Job, defer_exception: DeferJob) -> JobResult:
        """Defer this job by re-enqueueing it for later execution.

        Atomically deletes the JobProcess, re-enqueues the job, and creates
        a JobResult linking to the new request. This ensures the concurrency
        slot is released before attempting to re-enqueue.

        Raises:
            DeferError: If the job cannot be re-enqueued (e.g., due to concurrency limits).
                       The transaction will be rolled back and the JobProcess will remain.
        """
        # Calculate new retry_attempt based on increment_retries
        retry_attempt = (
            self.retry_attempt + 1
            if defer_exception.increment_retries
            else self.retry_attempt
        )

        with transaction.atomic():
            # 1. Save JobProcess state and delete (releases concurrency slot)
            saved_id = self.id
            job_process_uuid = self.uuid
            job_request_uuid = self.job_request_uuid
            started_at = self.started_at
            self.delete()

            # 2. Re-enqueue job (concurrency check can now pass)
            new_job_request = job.run_in_worker(
                queue=self.queue,
                delay=defer_exception.delay,
                priority=self.priority,
                retries=self.retries,
                retry_attempt=retry_attempt,
                concurrency_key=self.concurrency_key,
            )

            # Check if re-enqueue failed
            if new_job_request is None:
                # Restore id since transaction will roll back and object still exists
                self.id = saved_id
                raise DeferError(
                    f"Failed to re-enqueue deferred job {self.job_class}: "
                    f"concurrency limit reached for key '{self.concurrency_key}'"
                )

            # 3. Create JobResult linking to new request
            result = JobResult.query.create(
                ended_at=timezone.now(),
                error=f"Deferred for {defer_exception.delay} seconds",
                status=JobResultStatuses.DEFERRED,
                retry_job_request_uuid=new_job_request.uuid,
                # From the JobProcess
                job_process_uuid=job_process_uuid,
                started_at=started_at,
                # From the JobRequest
                job_request_uuid=job_request_uuid,
                job_class=self.job_class,
                parameters=self.parameters,
                priority=self.priority,
                source=self.source,
                queue=self.queue,
                retries=self.retries,
                retry_attempt=self.retry_attempt,
                concurrency_key=self.concurrency_key,
                trace_id=self.trace_id,
                span_id=self.span_id,
            )

            return result

    def convert_to_result(self, *, status: str, error: str = "") -> JobResult:
        """
        Convert this JobProcess to a JobResult.
        """
        with transaction.atomic():
            result = JobResult.query.create(
                ended_at=timezone.now(),
                error=error,
                status=status,
                # From the JobProcess
                job_process_uuid=self.uuid,
                started_at=self.started_at,
                # From the JobRequest
                job_request_uuid=self.job_request_uuid,
                job_class=self.job_class,
                parameters=self.parameters,
                priority=self.priority,
                source=self.source,
                queue=self.queue,
                retries=self.retries,
                retry_attempt=self.retry_attempt,
                concurrency_key=self.concurrency_key,
                trace_id=self.trace_id,
                span_id=self.span_id,
            )

            # Delete the JobProcess now
            self.delete()

        return result

    def as_json(self) -> dict[str, str | int | dict | None]:
        """A JSON-compatible representation to make it easier to reference in Sentry or logging"""
        return {
            "uuid": str(self.uuid),
            "created_at": self.created_at.isoformat(),
            "started_at": self.started_at.isoformat() if self.started_at else None,
            "job_request_uuid": str(self.job_request_uuid),
            "job_class": self.job_class,
            "parameters": self.parameters,
            "priority": self.priority,
            "source": self.source,
            "queue": self.queue,
            "retries": self.retries,
            "retry_attempt": self.retry_attempt,
            "concurrency_key": self.concurrency_key,
            "trace_id": self.trace_id,
            "span_id": self.span_id,
        }


class JobResultQuerySet(postgres.QuerySet["JobResult"]):
    def successful(self) -> Self:
        return self.filter(status=JobResultStatuses.SUCCESSFUL)

    def cancelled(self) -> Self:
        return self.filter(status=JobResultStatuses.CANCELLED)

    def lost(self) -> Self:
        return self.filter(status=JobResultStatuses.LOST)

    def errored(self) -> Self:
        return self.filter(status=JobResultStatuses.ERRORED)

    def retried(self) -> Self:
        return self.filter(
            postgres.Q(retry_job_request_uuid__isnull=False)
            | postgres.Q(retry_attempt__gt=0)
        )

    def failed(self) -> Self:
        return self.filter(
            status__in=[
                JobResultStatuses.ERRORED,
                JobResultStatuses.LOST,
                JobResultStatuses.CANCELLED,
            ]
        )

    def retryable(self) -> Self:
        return self.failed().filter(
            retry_job_request_uuid__isnull=True,
            retries__gt=0,
            retry_attempt__lt=F("retries"),
        )

    def retry_failed_jobs(self) -> None:
        for result in self.retryable():
            try:
                result.retry_job()
            except Exception:
                # If something went wrong (like a job class being deleted)
                # then we immediately increment the retry_attempt on the existing obj
                # so it won't retry forever.
                logger.exception(
                    "Failed to retry job, incrementing retry_attempt",
                    extra={"result": str(result)},
                )
                result.retry_attempt += 1
                result.save(update_fields=["retry_attempt"])


class JobResultStatuses(postgres.TextChoices):
    SUCCESSFUL = "SUCCESSFUL", "Successful"
    ERRORED = "ERRORED", "Errored"  # Threw an error
    CANCELLED = "CANCELLED", "Cancelled"  # Interrupted by shutdown/deploy
    DEFERRED = "DEFERRED", "Deferred"  # Intentionally rescheduled (will run again)
    LOST = (
        "LOST",
        "Lost",
    )  # Either process lost, lost in transit, or otherwise never finished


@postgres.register_model
class JobResult(postgres.Model):
    """
    All in-process and completed jobs are stored in this table.
    """

    uuid: UUID = types.UUIDField(default=uuid4)
    created_at: datetime.datetime = types.DateTimeField(auto_now_add=True)

    # From the Job
    job_process_uuid: UUID = types.UUIDField()
    started_at: datetime.datetime | None = types.DateTimeField(
        required=False, allow_null=True
    )
    ended_at: datetime.datetime | None = types.DateTimeField(
        required=False, allow_null=True
    )
    error: str = types.TextField(required=False)
    status: str = types.TextField(
        max_length=20,
        choices=JobResultStatuses.choices,
    )

    # From the JobRequest
    job_request_uuid: UUID = types.UUIDField()
    job_class: str = types.TextField(max_length=255)
    parameters: dict[str, Any] | None = types.JSONField(required=False, allow_null=True)
    priority: int = types.SmallIntegerField(default=0)
    source: str = types.TextField(required=False)
    queue: str = types.TextField(default="default", max_length=255)
    retries: int = types.SmallIntegerField(default=0)
    retry_attempt: int = types.SmallIntegerField(default=0)
    concurrency_key: str = types.TextField(max_length=255, required=False)

    # Retries
    retry_job_request_uuid: UUID | None = types.UUIDField(
        required=False, allow_null=True
    )

    # OpenTelemetry trace context
    trace_id: str | None = types.TextField(
        max_length=34, required=False, allow_null=True
    )
    span_id: str | None = types.TextField(
        max_length=18, required=False, allow_null=True
    )

    query: JobResultQuerySet = JobResultQuerySet()

    model_options = postgres.Options(
        ordering=["-created_at"],
        indexes=[
            postgres.Index(
                name="plainjobs_jobresult_created_at_idx", fields=["created_at"]
            ),
            postgres.Index(
                name="plainjobs_jobresult_job_process_uuid_idx",
                fields=["job_process_uuid"],
            ),
            postgres.Index(
                name="plainjobs_jobresult_started_at_idx", fields=["started_at"]
            ),
            postgres.Index(
                name="plainjobs_jobresult_ended_at_idx", fields=["ended_at"]
            ),
            postgres.Index(name="plainjobs_jobresult_status_idx", fields=["status"]),
            postgres.Index(
                name="plainjobs_jobresult_job_request_uuid_idx",
                fields=["job_request_uuid"],
            ),
            postgres.Index(
                name="plainjobs_jobresult_job_class_idx", fields=["job_class"]
            ),
            postgres.Index(name="plainjobs_jobresult_queue_idx", fields=["queue"]),
            postgres.Index(
                name="plainjobs_jobresult_trace_id_idx", fields=["trace_id"]
            ),
            postgres.Index(name="plainjobs_jobresult_uuid_idx", fields=["uuid"]),
        ],
        constraints=[
            postgres.UniqueConstraint(
                fields=["uuid"], name="plainjobs_jobresult_unique_uuid"
            ),
        ],
    )

    def retry_job(self, delay: int | None = None) -> JobRequest | None:
        retry_attempt = self.retry_attempt + 1
        job = jobs_registry.load_job(self.job_class, self.parameters or {})

        if delay is None:
            retry_delay = job.calculate_retry_delay(retry_attempt)
        else:
            retry_delay = delay

        with transaction.atomic():
            result = job.run_in_worker(
                # Pass most of what we know through so it stays consistent
                queue=self.queue,
                delay=retry_delay,
                priority=self.priority,
                retries=self.retries,
                retry_attempt=retry_attempt,
                concurrency_key=self.concurrency_key,
            )
            if result:
                self.retry_job_request_uuid = result.uuid
                self.save(update_fields=["retry_job_request_uuid"])
                return result

        return None
