import time
from datetime import timedelta

from django.core.management.base import BaseCommand
from django.db import connection, transaction
from django.db.models import Q
from django.utils import timezone

from finance.exporting import process_formula_export
from finance.importing import process_dataset
from finance.models import Dataset, FormulaExport

# 无咨询锁的数据库（如 SQLite）上等待该时限后才重新认领 processing 任务，
# 避免与仍在运行的 worker 抢任务；PostgreSQL 由咨询锁保证单 worker，可立即恢复。
STALE_PROCESSING_SECONDS = 15 * 60


class Command(BaseCommand):
    help = "Process durable import and formula-export jobs"

    def add_arguments(self, parser):
        parser.add_argument("--once", action="store_true")

    def handle(self, *args, **options):
        postgres = connection.vendor == "postgresql"
        if postgres:
            with connection.cursor() as cursor:
                cursor.execute("SELECT pg_try_advisory_lock(782146291)")
                if not cursor.fetchone()[0]:
                    self.stdout.write("Another worker owns the database lock")
                    return
        try:
            self.run_jobs(options, postgres)
        finally:
            if postgres:
                with connection.cursor() as cursor:
                    cursor.execute("SELECT pg_advisory_unlock(782146291)")

    def run_jobs(self, options, recover_immediately):
        while True:
            processed = self._next_import(recover_immediately)
            if not processed:
                processed = self._next_export(recover_immediately)
            if options["once"]:
                return
            if not processed:
                time.sleep(2)

    def _claim(self, model, immediate_recovery):
        stale_before = timezone.now() - timedelta(seconds=STALE_PROCESSING_SECONDS)
        claimable = Q(status="queued") | Q(status="processing", created_at__lt=stale_before)
        if immediate_recovery:
            claimable = Q(status="queued") | Q(status="processing")
        with transaction.atomic():
            qs = model.objects.filter(claimable).order_by("created_at")
            if connection.features.has_select_for_update_skip_locked:
                qs = qs.select_for_update(skip_locked=True)
            item = qs.first()
            if not item:
                return None
            # CAS 认领 queued 任务互斥；processing 仅在超时或咨询锁保护下重认领，
            # 重复发布对同一目标路径写入相同内容，不会产生重复文件。
            claimed = model.objects.filter(pk=item.pk, status=item.status).update(status="processing")
            return item if claimed else None

    def _next_import(self, recover_immediately):
        item = self._claim(Dataset, recover_immediately)
        if not item:
            return False
        Dataset.objects.filter(pk=item.pk).update(started_at=timezone.now())
        result = process_dataset(item.pk)
        self.stdout.write(f"{result.pk}: {result.status}, {result.row_count} rows")
        return True

    def _next_export(self, recover_immediately):
        item = self._claim(FormulaExport, recover_immediately)
        if not item:
            return False
        try:
            stats = process_formula_export(item.pk)
            self.stdout.write(f"{item.pk}: formula export ready, {stats['rows']} rows")
        except Exception as exc:
            self.stderr.write(f"Formula export {item.pk} failed: {exc}")
        return True