import asyncio
import json
import logging
import sys

from api.product_desc.product_desc_route import GenerateRequest, generate_description
from common.cache import redis_cache

logging.basicConfig(level=logging.INFO, format="20%(asctime)s [%(levelname)s] %(name)s: %(message)s", datefmt="%y-%m-%d %H:%M:%S")
logger = logging.getLogger("batch_worker")

BATCH_CONCURRENCY = 6  # 1 worker per Groq API key = max throughput

async def main():
    await redis_cache.initialize()
    redis = redis_cache.get_client()
    if not redis:
        logger.error("Redis client is not available. Please check REDIS_CACHE_TURN_ON in config.py.")
        sys.exit(1)

    # Initialize tables sequentially before starting concurrent tasks
    from common.ultra_desc_db import DescFieldConfig, UltraDescriptionDB
    UltraDescriptionDB.ensure_table()
    DescFieldConfig.ensure_table()

    logger.info("👷 Batch Worker started. Waiting for jobs on 'desc:queue'...")
    
    while True:
        try:
            # BRPOP blocks until an item is available, use timeout=2 to allow Windows Ctrl+C propagation!
            item = await redis.brpop("desc:queue", timeout=2)
            if not item:
                continue

            _, payload = item
            job_data = json.loads(payload)
            job_id = job_data.get("job_id")
            codes = job_data.get("codes", [])
            
            if not job_id or not codes:
                logger.warning(f"Invalid job payload: {job_data}")
                continue

            logger.info(f"🚀 Found Job {job_id} with {len(codes)} items.")
            
            # Setup job progress in Redis
            progress_key = f"desc:progress:{job_id}"
            await redis.set("desc:progress:latest_job", job_id)
            await redis.hset(progress_key, mapping={
                "is_running": "true",
                "total": len(codes),
                "done": 0,
                "errors": 0,
            })
            # Giữ trạng thái job trong 1 tiếng
            await redis.expire(progress_key, 3600)

            # Phân tách work và kiểm soát concurrency
            sem = asyncio.Semaphore(BATCH_CONCURRENCY)
            done_count = 0
            err_count = 0

            async def _worker(code: str):
                nonlocal done_count, err_count
                async with sem:
                    try:
                        logger.info(f"⚡ [Job {job_id}] [{done_count+1}/{len(codes)}] Generating {code}...")
                        req_data = GenerateRequest(internal_ref_code=code)
                        await generate_description(req_data)
                    except Exception as e:
                        logger.error(f"❌ Batch failed for {code}: {e}")
                        err_count += 1
                    finally:
                        done_count += 1
                        await redis.hset(progress_key, mapping={
                            "done": done_count,
                            "errors": err_count,
                            "current_code": code
                        })


            await asyncio.gather(*[_worker(code) for code in codes])
            
            await redis.hset(progress_key, "is_running", "false")
            logger.info(f"🎉 Job {job_id} finished. Total: {len(codes)}, Errors: {err_count}")

        except Exception as e:
            logger.error(f"Worker loop error: {e}")
            await asyncio.sleep(5)  # Backoff nếu Redis rớt

if __name__ == "__main__":
    # Workaround for Window's ProactorEventLoop not playing nice with some async teardowns
    if sys.platform == 'win32':
        asyncio.set_event_loop_policy(asyncio.WindowsSelectorEventLoopPolicy())
    asyncio.run(main())
