import asyncio
import json
import logging
import os
from datetime import datetime

from sqlalchemy import select

from app.models.base import async_session
from app.models.generation_record import GenerationRecord
from app.services.video_gen import get_active_engine, poll_task_status, download_video, _log_video_response
from app.services.image_gen import get_active_image_engine, download_image
from app.services.media_token_usage_snapshot_service import sync_generation_record_media_token_snapshot
from app.services.resource_accounting_service import (
    record_generation_record_generated_resource,
    safe_file_size,
)
from app.services.video_cover_service import create_video_cover_for_local_video
from app.config import settings
from app.services.generation_refund_service import mark_generation_record_failed_and_refund_once

logger = logging.getLogger("videogen")

POLL_INTERVAL = 30  # seconds between polls
MAX_POLLS = 60      # max 30 minutes total


class TaskQueue:
    def __init__(self):
        self.queue: asyncio.Queue[str] = asyncio.Queue()
        self.running = False
        self._active: dict[str, int] = {}  # record_id -> poll count

    async def enqueue(self, record_id: str):
        """Add a record to the polling queue."""
        await self.queue.put(record_id)

    async def recover(self):
        """Recover in-progress tasks from DB on startup."""
        async with async_session() as db:
            result = await db.execute(
                select(GenerationRecord).where(
                    GenerationRecord.status == "generating",
                    GenerationRecord.seedance_task_id.isnot(None),
                    GenerationRecord.deleted_at.is_(None),
                )
            )
            records = result.scalars().all()
            for record in records:
                await self.queue.put(record.id)
                logger.info(f"Recovered task: {record.id} (seedance: {record.seedance_task_id})")

    async def run(self):
        """Main polling loop."""
        self.running = True
        logger.info("Video queue started")

        while self.running:
            try:
                record_id = await asyncio.wait_for(self.queue.get(), timeout=5.0)
            except asyncio.TimeoutError:
                continue

            try:
                await self._process(record_id)
            except Exception as e:
                logger.error(f"Error processing {record_id}: {e}")
            finally:
                self.queue.task_done()

        logger.info("Video queue stopped")

    async def _process(self, record_id: str):
        """Process a single record: poll status and update DB."""
        async with async_session() as db:
            result = await db.execute(
                select(GenerationRecord).where(
                    GenerationRecord.id == record_id,
                    GenerationRecord.deleted_at.is_(None),
                )
                .with_for_update()
                .limit(1)
            )
            record = result.scalar_one_or_none()
            if not record or record.status != "generating":
                return

            if record.gen_type == "video":
                await self._process_video(db, record)
            else:
                await self._process_image(db, record)

    async def _process_video(self, db, record):
        """Process video generation task."""
        record_id = record.id
        
        if not record.seedance_task_id:
            await mark_generation_record_failed_and_refund_once(
                db,
                record=record,
                error_message="缺少外部任务ID",
            )
            await db.commit()
            return

        try:
            engine = await get_active_engine(db)
            poll_result = await poll_task_status(engine, record.seedance_task_id)
        except Exception as e:
            logger.error(f"Poll error for {record_id}: {e}")
            count = self._active.get(record_id, 0) + 1
            self._active[record_id] = count
            if count >= MAX_POLLS:
                await mark_generation_record_failed_and_refund_once(
                    db,
                    record=record,
                    error_message=f"轮询超时: {e}",
                )
                await db.commit()
                del self._active[record_id]
            else:
                await asyncio.sleep(POLL_INTERVAL)
                await self.queue.put(record_id)
            return

        status = poll_result["status"]
        try:
            resp_data = json.loads(poll_result.get("response_data", "{}"))
        except (json.JSONDecodeError, TypeError):
            resp_data = {}

        _log_video_response(record_id, resp_data, poll_result.get("error"))

        if status == "succeeded":
            file_url = poll_result.get("video_url", "")
            storage_path = None
            file_size_bytes = 0
            if settings.STORAGE_TYPE == "local" and file_url:
                try:
                    date_dir = datetime.now().strftime("%Y/%m/%d")
                    dest_dir = os.path.join(settings.STORAGE_LOCAL_PATH, date_dir)
                    os.makedirs(dest_dir, exist_ok=True)
                    dest = os.path.join(dest_dir, f"{record_id}.mp4")
                    await download_video(file_url, dest)
                    record.video_url = f"/generate/videos/{date_dir}/{record_id}.mp4"
                    cover_url, _cover_storage_path = create_video_cover_for_local_video(
                        record_id=record_id,
                        video_path=dest,
                        date_dir=date_dir,
                        log_prefix=f"GenerationRecord视频封面生成 record_id={record_id}",
                    )
                    record.video_cover_url = cover_url
                    storage_path = dest
                    file_size_bytes = safe_file_size(dest)
                except Exception as e:
                    logger.warning(f"Download failed, using remote URL: {e}")
                    record.video_url = file_url
            else:
                record.video_url = file_url
            record.video_tokens_used = poll_result.get("video_tokens", 0)
            await sync_generation_record_media_token_snapshot(db, record, provider_response=resp_data)
            record.status = "completed"
            record.generated_at = datetime.now()
            if record.video_url:
                await record_generation_record_generated_resource(
                    db,
                    record,
                    resource_url=record.video_url,
                    storage_path=storage_path,
                    file_size_bytes=file_size_bytes,
                    remote_url=file_url,
                    generated_at=record.generated_at,
                )
            self._active.pop(record_id, None)
            await db.commit()
            logger.info(f"Video task completed: {record_id}")

        elif status == "failed":
            await mark_generation_record_failed_and_refund_once(
                db,
                record=record,
                error_message=poll_result.get("error", "视频生成失败"),
            )
            self._active.pop(record_id, None)
            await db.commit()
            logger.info(f"Video task failed: {record_id}")

        else:
            count = self._active.get(record_id, 0) + 1
            self._active[record_id] = count
            if count >= MAX_POLLS:
                await mark_generation_record_failed_and_refund_once(
                    db,
                    record=record,
                    error_message="视频生成超时",
                )
                self._active.pop(record_id, None)
                await db.commit()
                logger.info(f"Video task timed out: {record_id}")
            else:
                await db.commit()
                await asyncio.sleep(POLL_INTERVAL)
                await self.queue.put(record_id)

    async def _process_image(self, db, record):
        """Process image generation task - calls API directly."""
        record_id = record.id
        from app.services.image_gen import submit_image_task, _log_image_response

        try:
            engine = await get_active_image_engine(db)
            poll_result = await asyncio.to_thread(
                submit_image_task,
                db,
                engine,
                record,
                include_media_references=False,
            )

            if poll_result["error"] == "":
                remote_url = poll_result.get("image_url")
                storage_path = None
                file_size_bytes = 0
                if settings.STORAGE_TYPE == "local" and remote_url:
                    try:
                        date_dir = datetime.now().strftime("%Y/%m/%d")
                        dest_dir = os.path.join(settings.STORAGE_IMAGE_LOCAL_PATH, date_dir)
                        os.makedirs(dest_dir, exist_ok=True)
                        dest = os.path.join(dest_dir, f"{record_id}.png")
                        await download_image(remote_url, dest)
                        record.image_url = f"/generate/images/{date_dir}/{record_id}.png"
                        storage_path = dest
                        file_size_bytes = safe_file_size(dest)
                    except Exception as e:
                        logger.warning(f"Download failed, using remote URL: {e}")
                        record.image_url = remote_url
                else:
                    record.image_url = remote_url
                record.image_tokens_used = poll_result.get("image_tokens", 0)
                await sync_generation_record_media_token_snapshot(db, record, provider_response=poll_result)
                record.status = "completed"
                record.generated_at = datetime.now()
                if record.image_url:
                    await record_generation_record_generated_resource(
                        db,
                        record,
                        resource_url=record.image_url,
                        storage_path=storage_path,
                        file_size_bytes=file_size_bytes,
                        remote_url=remote_url,
                        generated_at=record.generated_at,
                    )
                await db.commit()
                logger.info(f"Image task completed: {record_id}")
            else:
                await mark_generation_record_failed_and_refund_once(
                    db,
                    record=record,
                    error_message=poll_result.get("error", "图片生成失败"),
                )
                await db.commit()
                logger.info(f"Image task failed: {record_id}")
            _log_image_response(record_id, poll_result)

        except Exception as e:
            await mark_generation_record_failed_and_refund_once(
                db,
                record=record,
                error_message=str(e),
            )
            _log_image_response(record_id, {}, str(e))
            await db.commit()
            logger.error(f"Image task failed: {record_id}, error: {e}")

    def stop(self):
        """Signal the queue to stop."""
        self.running = False


task_queue = TaskQueue()
