import asyncio
import logging
from datetime import datetime, timezone

from sqlalchemy import select, update
from sqlalchemy.orm import selectinload

from app.models.base import async_session
from app.models.upload_task import UploadTask
from app.models.generated_resource import GeneratedResource
from app.models.user_oauth import UserOAuth
from app.models.resources_material import ResourcesMaterial
from app.utils.id_gen import generate_id
from app.utils.douyinApi import DouyinApi

import os
import hashlib

logger = logging.getLogger("upload_queue")
douyin_api = DouyinApi()


class UploadQueue:
    def __init__(self):
        self.queue: asyncio.Queue[str] = asyncio.Queue()
        self.running = False

    async def enqueue(self, task_id: str):
        """Add a task to the queue."""
        await self.queue.put(task_id)

    async def recover(self):
        """Recover pending tasks from DB on startup."""
        async with async_session() as db:
            result = await db.execute(
                select(UploadTask).where(
                    UploadTask.status.in_([1, 2]),
                    UploadTask.deleted_at.is_(None),
                )
            )
            tasks = result.scalars().all()
            for task in tasks:
                await self.queue.put(task.id)
                logger.info(f"Recovered upload task: {task.id}")

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

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

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

        logger.info("Upload queue stopped")

    async def _process(self, task_id: str):
        """Process a single upload task."""
        async with async_session() as db:
            result = await db.execute(
                select(UploadTask).where(
                    UploadTask.id == task_id,
                    UploadTask.deleted_at.is_(None),
                ).with_for_update()
            )
            task = result.scalar_one_or_none()

            if not task:
                logger.warning(f"Upload task {task_id} not found or deleted")
                return

            if task.status == 2:
                logger.warning(f"Upload task {task_id} is already running")
                return

            task.status = 2
            await db.commit()

        try:
            result = await _upload_single_material(
                task.id,
                task.user_id,
                task.oauth_id,
                task.advertiser_id,
                task.resource_id,
                db=None
            )

            async with async_session() as db:
                if result.get("success"):
                    await db.execute(
                        update(UploadTask).where(UploadTask.id == task_id).values(
                            status=3,
                            note=result.get("message", "上传成功"),
                        )
                    )
                else:
                    await db.execute(
                        update(UploadTask).where(UploadTask.id == task_id).values(
                            status=4,
                            note=result.get("error", "上传失败"),
                        )
                    )
                await db.commit()

            logger.info(f"Upload task {task_id} completed: {'success' if result.get('success') else 'failed'}")

        except Exception as e:
            async with async_session() as db:
                await db.execute(
                    update(UploadTask).where(UploadTask.id == task_id).values(
                        status=4,
                        note=f"上传失败: {str(e)}",
                    )
                )
                await db.commit()

            logger.error(f"Upload task {task_id} failed: {e}")

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


async def _upload_single_material(
    task_id: str,
    user_id: str,
    oauth_id: str,
    advertiser_id: str,
    resource_id: str,
    db=None
) -> dict:
    if db is None:
        from app.models.base import async_session
        db = async_session()
    
    try:
        oauth = await db.execute(
            select(UserOAuth).where(
                UserOAuth.id == oauth_id,
                UserOAuth.deleted_at.is_(None),
                UserOAuth.user_id == user_id,
            )
        )
        oauth = oauth.scalar_one_or_none()
        if not oauth:
            return {
                "success": False,
                "error": "授权记录不存在"
            }

        resource = await db.execute(
            select(GeneratedResource).where(
                GeneratedResource.id == resource_id,
                GeneratedResource.user_id == user_id,
                GeneratedResource.deleted_at.is_(None),
            )
        )
        resource = resource.scalar_one_or_none()
        
        if not resource:
            return {
                "success": False,
                "error": f"资源 {resource_id} 不存在或不属于当前用户"
            }

        resource_type = resource.resource_type
        storage_path = resource.storage_path

        if resource_type not in ["image", "video"]:
            return {
                "success": False,
                "error": f"不支持的资源类型: {resource_type}，仅支持image和video",
            }

        if not storage_path:
            return {
                "success": False,
                "error": "资源本地存储路径为空",
            }

        return await _upload_to_juliang(
            oauth_id, storage_path, resource_type, advertiser_id, resource, db, user_id
        )
    finally:
        if db is not None:
            await db.close()


async def _upload_to_juliang(
    oauth_id: str,
    storage_path: str,
    resource_type: str,
    advertiser_id: str,
    resource: GeneratedResource,
    db,
    current_user_id: str,
) -> dict:
    filename = os.path.basename(storage_path)

    if resource_type == "image":
        if resource.file_size_bytes > 5 * 1024 * 1024:
            return {
                "resource_id": resource.id,
                "advertiser_id": advertiser_id,
                "filename": filename,
                "success": False,
                "error": "图片大小超过5M，不支持上传",
            }

        with open(storage_path, "rb") as f:
            file_content = f.read()
            image_signature = hashlib.md5(file_content).hexdigest()

        data = {
            "advertiser_id": advertiser_id,
            "upload_type": "UPLOAD_BY_FILE",
            "image_signature": image_signature,
            "filename": filename,
        }

        files = {
            "image_file": (filename, file_content, "image/png"),
        }

        response = await douyin_api.upload_image_material(oauth_id, data, files)

        if response["code"] != 0:
            return {
                "resource_id": resource.id,
                "advertiser_id": advertiser_id,
                "filename": filename,
                "success": False,
                "error": response.get("message", "上传失败"),
            }

        image_id = response["data"]["id"]
        material_id = response["data"]["material_id"]

        existing = await db.execute(
            select(ResourcesMaterial).where(
                ResourcesMaterial.oauth_id == oauth_id,
                ResourcesMaterial.advertiser_id == advertiser_id,
                ResourcesMaterial.material_id == str(material_id),
                ResourcesMaterial.upload_id == str(image_id),
                ResourcesMaterial.deleted_at.is_(None),
            )
        )
        existing = existing.scalar_one_or_none()

        if existing:
            return {
                "resource_id": resource.id,
                "advertiser_id": advertiser_id,
                "filename": filename,
                "success": True,
                "material_id": str(material_id),
                "upload_id": str(image_id),
                "message": "素材已存在，跳过添加",
            }

        db.add(ResourcesMaterial(
            id=generate_id(),
            oauth_id=oauth_id,
            advertiser_id=advertiser_id,
            target_table="generated_resources",
            target_id=resource.id,
            material_id=str(material_id),
            upload_id=str(image_id),
            resource_type=resource_type,
            user_id=current_user_id,
        ))
        await db.commit()

        return {
            "resource_id": resource.id,
            "advertiser_id": advertiser_id,
            "filename": filename,
            "success": True,
            "material_id": str(material_id),
            "upload_id": str(image_id),
        }

    elif resource_type == "video":
        if resource.file_size_bytes > 500 * 1024 * 1024:
            return {
                "resource_id": resource.id,
                "advertiser_id": advertiser_id,
                "filename": filename,
                "success": False,
                "error": "视频大小超过500M，不支持上传",
            }

        with open(storage_path, "rb") as f:
            file_content = f.read()
            video_signature = hashlib.md5(file_content).hexdigest()

        data = {
            "advertiser_id": advertiser_id,
            "upload_type": "UPLOAD_BY_FILE",
            "video_signature": video_signature,
            "filename": filename,
        }

        files = {
            "video_file": (filename, file_content, "video/mp4"),
        }

        response = await douyin_api.upload_video_material(oauth_id, data, files)

        if response["code"] != 0:
            return {
                "resource_id": resource.id,
                "advertiser_id": advertiser_id,
                "filename": filename,
                "success": False,
                "error": response.get("message", "上传失败"),
            }

        video_id = response["data"]["video_id"]
        material_id = response["data"]["material_id"]

        existing = await db.execute(
            select(ResourcesMaterial).where(
                ResourcesMaterial.oauth_id == oauth_id,
                ResourcesMaterial.advertiser_id == advertiser_id,
                ResourcesMaterial.material_id == str(material_id),
                ResourcesMaterial.upload_id == str(video_id),
                ResourcesMaterial.deleted_at.is_(None),
            )
        )
        existing = existing.scalar_one_or_none()

        if existing:
            return {
                "resource_id": resource.id,
                "advertiser_id": advertiser_id,
                "filename": filename,
                "success": True,
                "material_id": str(material_id),
                "upload_id": str(video_id),
                "message": "素材已存在，跳过添加",
            }

        db.add(ResourcesMaterial(
            id=generate_id(),
            oauth_id=oauth_id,
            advertiser_id=advertiser_id,
            target_table="generated_resources",
            target_id=resource.id,
            material_id=str(material_id),
            upload_id=str(video_id),
            resource_type=resource_type,
            user_id=current_user_id,
        ))
        await db.commit()

        return {
            "resource_id": resource.id,
            "advertiser_id": advertiser_id,
            "filename": filename,
            "success": True,
            "material_id": str(material_id),
            "upload_id": str(video_id),
        }

    else:
        return {
            "resource_id": resource.id,
            "advertiser_id": advertiser_id,
            "filename": filename,
            "success": False,
            "error": f"不支持的资源类型: {resource_type}",
        }


upload_queue = UploadQueue()