import time
from datetime import timedelta

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

from api_v1.mdy_plate_order_staging_qiniu_upload import (
    DEFAULT_DOWNLOAD_CHUNK_SIZE,
    DEFAULT_DOWNLOAD_TIMEOUT,
    DEFAULT_MAX_IMAGE_SIZE,
    DEFAULT_SPOOL_MAX_SIZE,
    DEFAULT_UPLOAD_STATE_NAME,
    build_default_qiniu_rate_limiter,
    get_or_create_audit_for_work_item,
    get_or_create_default_upload_state,
    iter_plate_image_work_items,
    process_qiniu_image_upload_audit,
)
from api_v1.models import (
    MDYPlateOrderStaging,
    MDYPlateOrderStagingQiniuImageUploadAudit,
    MDYPlateOrderStagingQiniuImageUploadState,
)


class Command(BaseCommand):
    help = '将明道云开版暂存 raw.开版图 图片分批上传到七牛，并记录逐图审计'
    lock_key = 'api_v1.upload_mdy_plate_order_staging_images_to_qiniu'

    def add_arguments(self, parser):
        parser.add_argument('--dry-run', action='store_true', help='只输出配置和当前游标，不扫描、不写库、不下载、不上传')
        parser.add_argument('--scan-batch-size', type=int, default=20, help='每轮最多扫描多少条 staging')
        parser.add_argument('--upload-batch-size', type=int, default=5, help='每轮最多上传多少条审计记录')
        parser.add_argument(
            '--max-upload-attempts',
            type=int,
            default=None,
            help='最多尝试处理多少条审计上传记录；测试/灰度用，不传则持续运行到结束',
        )
        parser.add_argument('--qps', type=float, default=None, help='下载/上传限速；默认 settings 或 2 qps；<=0 表示不限速')
        parser.add_argument('--loop-sleep-seconds', type=float, default=0.2, help='每轮结束后的短暂休眠，降低 CPU 压力')
        parser.add_argument('--reset-cursor', action='store_true', help='重置 staging 扫描游标，从头重新发现审计记录')
        parser.add_argument('--retry-failed', action='store_true', help='重试 status=failed 的审计记录')
        parser.add_argument(
            '--retry-stale-processing',
            action='store_true',
            help='重试长时间停留在 processing 的审计记录',
        )
        parser.add_argument(
            '--stale-after-minutes',
            type=int,
            default=120,
            help='processing 超过多少分钟视为可重试的陈旧记录',
        )
        parser.add_argument('--max-attempts', type=int, default=3, help='单张图最多尝试次数；不限制请传 0')
        parser.add_argument(
            '--max-image-size-mb',
            type=int,
            default=DEFAULT_MAX_IMAGE_SIZE // 1024 // 1024,
            help='单张图片最大下载体积，超过则失败落审计',
        )
        parser.add_argument(
            '--spool-max-size-mb',
            type=int,
            default=DEFAULT_SPOOL_MAX_SIZE // 1024 // 1024,
            help='单张图内存缓冲上限，超过后临时落盘',
        )
        parser.add_argument(
            '--download-chunk-size-mb',
            type=int,
            default=DEFAULT_DOWNLOAD_CHUNK_SIZE // 1024 // 1024,
            help='流式下载的单块大小',
        )
        parser.add_argument('--timeout', type=int, default=DEFAULT_DOWNLOAD_TIMEOUT, help='单次下载超时时间（秒）')
        parser.add_argument('--progress-every-loops', type=int, default=1, help='每多少轮输出一次心跳，默认每轮输出')

    def handle(self, *args, **options):
        dry_run = options['dry_run']
        scan_batch_size = options['scan_batch_size']
        upload_batch_size = options['upload_batch_size']
        max_upload_attempts = options['max_upload_attempts']
        loop_sleep_seconds = options['loop_sleep_seconds']
        retry_failed = options['retry_failed']
        retry_stale_processing = options['retry_stale_processing']
        stale_before = timezone.now() - timedelta(minutes=options['stale_after_minutes'])
        max_attempts = options['max_attempts'] or None

        if scan_batch_size <= 0:
            raise CommandError('--scan-batch-size 必须大于 0')
        if upload_batch_size <= 0:
            raise CommandError('--upload-batch-size 必须大于 0')
        if max_upload_attempts is not None and max_upload_attempts <= 0:
            raise CommandError('--max-upload-attempts 必须大于 0')

        if dry_run:
            state = MDYPlateOrderStagingQiniuImageUploadState.objects.filter(
                name=DEFAULT_UPLOAD_STATE_NAME,
            ).first()
            config = {
                'dry_run': True,
                'scan_batch_size': scan_batch_size,
                'upload_batch_size': upload_batch_size,
                'max_upload_attempts': max_upload_attempts,
                'qps': options['qps'],
                'loop_sleep_seconds': loop_sleep_seconds,
                'retry_failed': retry_failed,
                'retry_stale_processing': retry_stale_processing,
                'max_attempts': max_attempts,
                'state_exists': state is not None,
                'state_last_staging_id': state.last_staging_id if state else None,
                'state_staging_finished': state.staging_finished if state else None,
            }
            self.stdout.write(f'start: {config}')
            self.stdout.write(self.style.SUCCESS('dry-run stop: no database writes, no upload'))
            return

        if not self._try_acquire_db_lock():
            raise CommandError('已有 upload_mdy_plate_order_staging_images_to_qiniu 命令正在运行')
        try:
            self._run_upload_loop(
                options=options,
                scan_batch_size=scan_batch_size,
                upload_batch_size=upload_batch_size,
                max_upload_attempts=max_upload_attempts,
                loop_sleep_seconds=loop_sleep_seconds,
                retry_failed=retry_failed,
                retry_stale_processing=retry_stale_processing,
                stale_before=stale_before,
                max_attempts=max_attempts,
            )
        finally:
            self._release_db_lock()

    def _run_upload_loop(
        self,
        *,
        options,
        scan_batch_size: int,
        upload_batch_size: int,
        max_upload_attempts: int | None,
        loop_sleep_seconds: float,
        retry_failed: bool,
        retry_stale_processing: bool,
        stale_before,
        max_attempts: int | None,
    ) -> None:
        state = get_or_create_default_upload_state()
        if options['reset_cursor']:
            state.last_staging_id = 0
            state.staging_finished = False
            state.scanned_staging_rows = 0
            state.discovered_image_items = 0
            state.created_audit_rows = 0
            state.existing_audit_rows = 0
            state.last_heartbeat_at = timezone.now()
            state.save(
                update_fields=[
                    'last_staging_id',
                    'staging_finished',
                    'scanned_staging_rows',
                    'discovered_image_items',
                    'created_audit_rows',
                    'existing_audit_rows',
                    'last_heartbeat_at',
                    'updated_at',
                ]
            )

        config = {
            'dry_run': False,
            'scan_batch_size': scan_batch_size,
            'upload_batch_size': upload_batch_size,
            'max_upload_attempts': max_upload_attempts,
            'qps': options['qps'],
            'loop_sleep_seconds': loop_sleep_seconds,
            'retry_failed': retry_failed,
            'retry_stale_processing': retry_stale_processing,
            'max_attempts': max_attempts,
            'state_last_staging_id': state.last_staging_id,
            'state_staging_finished': state.staging_finished,
        }
        self.stdout.write(f'start: {config}')

        rate_limiter = build_default_qiniu_rate_limiter(options['qps'])
        totals = {
            'loops': 0,
            'scanned_staging_rows': 0,
            'discovered_image_items': 0,
            'created_audit_rows': 0,
            'existing_audit_rows': 0,
            'uploaded': 0,
            'reused': 0,
            'failed': 0,
            'skipped': 0,
        }

        while True:
            totals['loops'] += 1
            effective_upload_batch_size = upload_batch_size
            if max_upload_attempts is not None:
                remaining_attempts = max_upload_attempts - self._upload_attempt_count(totals)
                if remaining_attempts <= 0:
                    self.stdout.write(self.style.SUCCESS(f'max-upload-attempts stop: {totals}'))
                    return
                effective_upload_batch_size = min(upload_batch_size, remaining_attempts)

            discover_stats = self._discover_next_staging_batch(
                scan_batch_size=scan_batch_size,
            )
            upload_stats = self._upload_next_audit_batch(
                upload_batch_size=effective_upload_batch_size,
                retry_failed=retry_failed,
                retry_stale_processing=retry_stale_processing,
                stale_before=stale_before,
                max_attempts=max_attempts,
                rate_limiter=rate_limiter,
                max_image_size=options['max_image_size_mb'] * 1024 * 1024,
                spool_max_size=options['spool_max_size_mb'] * 1024 * 1024,
                chunk_size=options['download_chunk_size_mb'] * 1024 * 1024,
                timeout=options['timeout'],
            )
            self._add_stats(totals, discover_stats)
            self._add_stats(totals, upload_stats)

            state = get_or_create_default_upload_state()
            heartbeat = {
                'loop': totals['loops'],
                'state_last_staging_id': state.last_staging_id,
                'state_staging_finished': state.staging_finished,
                'discover': discover_stats,
                'upload': upload_stats,
                'totals': totals,
            }
            state.last_heartbeat_at = timezone.now()
            state.save(update_fields=['last_heartbeat_at', 'updated_at'])

            if options['progress_every_loops'] > 0 and totals['loops'] % options['progress_every_loops'] == 0:
                self.stdout.write(f'progress: {heartbeat}')

            if state.staging_finished and upload_stats['selected_audit_rows'] == 0:
                self.stdout.write(self.style.SUCCESS(f'finished: {heartbeat}'))
                return

            if max_upload_attempts is not None and self._upload_attempt_count(totals) >= max_upload_attempts:
                self.stdout.write(self.style.SUCCESS(f'max-upload-attempts stop: {heartbeat}'))
                return

            if loop_sleep_seconds > 0:
                time.sleep(loop_sleep_seconds)

    def _try_acquire_db_lock(self) -> bool:
        if connection.vendor != 'postgresql':
            return True
        with connection.cursor() as cursor:
            cursor.execute('select pg_try_advisory_lock(hashtext(%s))', [self.lock_key])
            return bool(cursor.fetchone()[0])

    def _release_db_lock(self) -> None:
        if connection.vendor != 'postgresql':
            return
        with connection.cursor() as cursor:
            cursor.execute('select pg_advisory_unlock(hashtext(%s))', [self.lock_key])

    def _discover_next_staging_batch(self, *, scan_batch_size: int) -> dict[str, int]:
        state = get_or_create_default_upload_state()
        stats = {
            'scanned_staging_rows': 0,
            'discovered_image_items': 0,
            'created_audit_rows': 0,
            'existing_audit_rows': 0,
        }
        if state.staging_finished:
            return stats

        rows = list(
            MDYPlateOrderStaging.objects.filter(id__gt=state.last_staging_id)
            .only('id', 'mdy_rowid', 'raw')
            .order_by('id')[:scan_batch_size]
        )
        if not rows:
            state.staging_finished = True
            state.last_heartbeat_at = timezone.now()
            state.save(update_fields=['staging_finished', 'last_heartbeat_at', 'updated_at'])
            return stats

        for staging in rows:
            stats['scanned_staging_rows'] += 1
            for work_item in iter_plate_image_work_items(staging):
                stats['discovered_image_items'] += 1
                _audit, created = get_or_create_audit_for_work_item(work_item)
                if created:
                    stats['created_audit_rows'] += 1
                else:
                    stats['existing_audit_rows'] += 1

        state.last_staging_id = rows[-1].id
        state.scanned_staging_rows += stats['scanned_staging_rows']
        state.discovered_image_items += stats['discovered_image_items']
        state.created_audit_rows += stats['created_audit_rows']
        state.existing_audit_rows += stats['existing_audit_rows']
        state.last_heartbeat_at = timezone.now()
        state.save(
            update_fields=[
                'last_staging_id',
                'scanned_staging_rows',
                'discovered_image_items',
                'created_audit_rows',
                'existing_audit_rows',
                'last_heartbeat_at',
                'updated_at',
            ]
        )
        return stats

    def _upload_next_audit_batch(self, **kwargs) -> dict[str, int]:
        upload_batch_size = kwargs.pop('upload_batch_size')
        retry_failed = kwargs['retry_failed']
        retry_stale_processing = kwargs['retry_stale_processing']
        stale_before = kwargs['stale_before']

        statuses = MDYPlateOrderStagingQiniuImageUploadAudit.Status
        status_filter = Q(status=statuses.PENDING)
        if retry_failed:
            status_filter |= Q(status=statuses.FAILED)
        if retry_stale_processing:
            status_filter |= Q(status=statuses.PROCESSING, last_attempt_at__lte=stale_before)

        audits = list(
            MDYPlateOrderStagingQiniuImageUploadAudit.objects.filter(status_filter)
            .select_related('staging')
            .order_by('id')[:upload_batch_size]
        )
        stats = {
            'selected_audit_rows': len(audits),
            'uploaded': 0,
            'reused': 0,
            'failed': 0,
            'skipped': 0,
        }

        process_kwargs = dict(kwargs)
        process_kwargs.pop('retry_failed', None)
        process_kwargs.pop('retry_stale_processing', None)
        process_kwargs.pop('stale_before', None)
        for audit in audits:
            result = process_qiniu_image_upload_audit(
                audit,
                retry_failed=retry_failed,
                retry_stale_processing=retry_stale_processing,
                stale_before=stale_before,
                **process_kwargs,
            )
            if result.action == 'uploaded':
                stats['uploaded'] += 1
            elif result.action == 'reused':
                stats['reused'] += 1
            elif result.action == 'failed':
                stats['failed'] += 1
            else:
                stats['skipped'] += 1

        return stats

    @staticmethod
    def _add_stats(target: dict[str, int], source: dict[str, int]) -> None:
        for key, value in source.items():
            if key in target:
                target[key] += value

    @staticmethod
    def _upload_attempt_count(stats: dict[str, int]) -> int:
        return stats.get('uploaded', 0) + stats.get('reused', 0) + stats.get('failed', 0) + stats.get('skipped', 0)
