import asyncio
from datetime import datetime, timezone

from sqlalchemy import select, update
from sqlalchemy.orm import selectinload
from sqlalchemy.ext.asyncio import AsyncSession

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.models.user_oauth_account import UserOAuthAccount
from app.models.pre_test_template import PreTestTemplate
from app.utils.id_gen import generate_id
from app.utils.douyinApi import DouyinApi
from app.utils.logger import get_logger
import os

import hashlib
import json

logger = get_logger("upload_queue", "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,
                other_info=json.loads(task.other_info) if task.other_info else 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", "上传成功"),
                        )
                    )
                    
                    #请求巨量接口获取账户信息，如果存在则更新，否则插入新记录
                    param = {
                        "account_ids" : json.dumps([int(task.advertiser_id)]),
                    }
                    try:
                        account_info = await douyin_api.get_account_info(task.oauth_id, param)
                       
                        code = account_info.get("code", 0)
                        try:
                            code = int(code)
                        except (ValueError, TypeError):
                            code = -1
                        
                        if code != 0:
                            logger.error(f"Error getting account info: {json.dumps(account_info)}")
                        else:
                            existing_account = await db.execute(
                                select(UserOAuthAccount).where(
                                    UserOAuthAccount.oauth_id == task.oauth_id,
                                    UserOAuthAccount.advertiser_id == task.advertiser_id,
                                )
                            )
                            existing_account = existing_account.scalar_one_or_none()
                            
                            if existing_account:
                                await db.execute(
                                        update(UserOAuthAccount).where(
                                            UserOAuthAccount.id == existing_account.id,
                                        ).values(
                                            deleted_at=None,
                                            advertiser_name= account_info.get("data", {}).get("account_detail_list", [{}])[0].get("advertiser_name", ""),
                                            advertiser_role= "",
                                        )
                                    )      
                            else:
                                db.add(UserOAuthAccount(
                                    id=generate_id(),
                                    oauth_id=task.oauth_id,
                                    advertiser_id=task.advertiser_id,
                                    advertiser_name= account_info.get("data", {}).get("account_detail_list", [{}])[0].get("advertiser_name", ""),
                                    advertiser_role= "",
                                ))
                    except Exception as e:
                        logger.error(f"Exception getting account info: {e}")
                else:
                    await db.execute(
                        update(UploadTask).where(UploadTask.id == task_id).values(
                            status=4,
                            note=result.get("error", "上传失败"),
                        )
                    )
                await db.commit()

            message = result.get('message') or result.get('error', 'Unknown error')
            logger.info(f"Upload task {task_id} completed: {'success' if result.get('success') else 'failed'}. Message: {message}")

        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,
    other_info=None,
) -> dict:
    own_db = False
    if db is None:
        from app.models.base import async_session
        db = async_session()
        own_db = True
    
    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
        file_name = resource.file_name

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

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

        result = await _upload_to_juliang(
            oauth_id, storage_path, resource_type, advertiser_id, resource, db, user_id, file_name
        )

        if result.get("success") and resource_type == "video" and other_info and "is_pre_test" in other_info and other_info["is_pre_test"] == "1" and "pre_test_template" in other_info:
            await _pre_test_material(
                oauth_id,
                advertiser_id,
                [result.get("upload_id")],
                other_info["pre_test_template"],
                db,
            )
            
        return result
    finally:
        if own_db and 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,
    file_name: str,
) -> dict:
    #如果file_name不等于空，那么就是用file_name，否则用storage_path的文件名
    if file_name:
        filename = file_name
    else:
        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", "上传失败"),
            }

        #上传成功，添加视频id，素材id到素材库
        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}",
        }


async def _pre_test_material(
    oauth_id: str,
    advertiser_id: str,
    video_ids: list[str],
    pre_test_template_id: str,
    db: AsyncSession,
) -> any:
    oauth = await db.execute(
        select(UserOAuth).where(
            UserOAuth.id == oauth_id,
            UserOAuth.deleted_at.is_(None),
        )
    )
    oauth = oauth.scalar_one_or_none()
    if not oauth:
        note = "授权记录不存在"
        for video_id in video_ids:
            await _update_material_pre_test_status(
                db, oauth_id, advertiser_id, video_id,
                task_id=None,
                status="FAILED",
                note=note,
                pre_result=None,
                pre_test_template_id=pre_test_template_id,
            )
        await db.commit()
        return {"code": -1, "message": note, "data": {}}

    pre_test_template = await db.execute(
        select(PreTestTemplate).where(
            PreTestTemplate.id == pre_test_template_id,
            PreTestTemplate.deleted_at.is_(None),
            PreTestTemplate.user_id == oauth.user_id,
        )
    )
    pre_test_template = pre_test_template.scalar_one_or_none()
    if not pre_test_template:
        note = f"前测模板 {pre_test_template_id} 不存在"
        for video_id in video_ids:
            await _update_material_pre_test_status(
                db, oauth_id, advertiser_id, video_id,
                task_id=None,
                status="FAILED",
                note=note,
                pre_result=None,
                pre_test_template_id=pre_test_template_id,
            )
        await db.commit()
        return {"code": -1, "message": note, "data": {}}
    
    diagnose_config = {}
    if pre_test_template.platform:
        diagnose_config["platform"] = pre_test_template.platform
    if pre_test_template.external_action:
        diagnose_config["external_action"] = pre_test_template.external_action
    if pre_test_template.cpa_bid:
        diagnose_config["cpa_bid"] = pre_test_template.cpa_bid
    if pre_test_template.audience_gender:
        diagnose_config["audience_gender"] = pre_test_template.audience_gender
    if pre_test_template.audience_age:
        diagnose_config["audience_age"] = json.loads(pre_test_template.audience_age)
    if pre_test_template.audience_region:
        diagnose_config["audience_region"] = json.loads(pre_test_template.audience_region)
    if pre_test_template.audience_network:
        diagnose_config["audience_network"] = json.loads(pre_test_template.audience_network)
    if pre_test_template.cus_name:
        diagnose_config["cus_name"] = pre_test_template.cus_name
    if pre_test_template.pricing_type:
        diagnose_config["pricing_type"] = pre_test_template.pricing_type
    if pre_test_template.cost_cap:
        diagnose_config["cost_cap"] = pre_test_template.cost_cap
    if pre_test_template.target_cost:
        diagnose_config["target_cost"] = pre_test_template.target_cost
    if pre_test_template.nobid:
        diagnose_config["nobid"] = pre_test_template.nobid
    if pre_test_template.cpc_bid:
        diagnose_config["cpc_bid"] = pre_test_template.cpc_bid
    if pre_test_template.budget:
        diagnose_config["budget"] = pre_test_template.budget

    params = {
        "advertiser_id": int(advertiser_id),
        "video_ids": video_ids,
        "diagnose_config": diagnose_config,
    }
    response = await douyin_api.pre_test_material(oauth_id, params)
    
    code = response.get("code", -1)
    
    if code != 0:
        note = response.get("message", "未知错误")
        for video_id in video_ids:
            await _update_material_pre_test_status(
                db, oauth_id, advertiser_id, video_id,
                task_id=None,
                status="FAILED",
                note=note,
                pre_result=None,
                pre_test_template_id=pre_test_template_id,
            )
        await db.commit()
        return response
    
    data = response.get("data", {})
    task_ids = data.get("task_ids", [])
    fail_video_ids = data.get("fail_video_ids", {})
    
    success_count = 0
    for i, video_id in enumerate(video_ids):
        if video_id in fail_video_ids:
            fail_info = fail_video_ids[video_id]
            err_code = fail_info.get("err_code", "")
            err_message = fail_info.get("err_message", "未知错误")
            note = f"失败[{err_code}]: {err_message}"
            
            await _update_material_pre_test_status(
                db, oauth_id, advertiser_id, video_id,
                task_id=None,
                status="FAILED",
                note=note,
                pre_result=None,
                pre_test_template_id=pre_test_template_id,
            )
        else:
            task_id = str(task_ids[success_count]) if success_count < len(task_ids) else None
            
            await _update_material_pre_test_status(
                db, oauth_id, advertiser_id, video_id,
                task_id=task_id,
                status="PENDING",
                note="",
                pre_result=None,
                pre_test_template_id=pre_test_template_id,
            )
            success_count += 1
    
    await db.commit()
    
    return response


async def _update_material_pre_test_status(
    db: AsyncSession,
    oauth_id: str,
    advertiser_id: str,
    upload_id: str,
    task_id: str | None,
    status: str,
    note: str,
    pre_result: str | None,
    pre_test_template_id: str | None,
):
    await db.execute(
        update(ResourcesMaterial).where(
            ResourcesMaterial.oauth_id == oauth_id,
            ResourcesMaterial.advertiser_id == advertiser_id,
            ResourcesMaterial.upload_id == upload_id,
            ResourcesMaterial.deleted_at.is_(None),
        ).values(
            task_id=task_id,
            status=status,
            note=note,
            pre_result=pre_result,
            pre_test_template_id=pre_test_template_id,
        )
    )


upload_queue = UploadQueue()