import asyncio
import json
import logging
import os
import sys
from typing import List, Optional

import httpx

sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from api.product_desc.bulk_ops_route import (_FIELD_LABELS, _call_codex, _render_description_text,
                                             get_pooled_connection_compat)
from common.cache import redis_cache

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

BATCH_CONCURRENCY = 100
PG_TABLE = "dashboard_canifa.ultra_descriptions"

# ═══ Groq API Keys Round-Robin for Tags ═══
GROQ_API_KEYS = [
    "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",
    "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",
    "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",
    "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",
    "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",
    "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",
]
_groq_key_index = 0

def _next_groq_key() -> str:
    global _groq_key_index
    key = GROQ_API_KEYS[_groq_key_index]
    _groq_key_index = (_groq_key_index + 1) % len(GROQ_API_KEYS)
    return key


# ═══ Tags Allowed (Plain, No Prefix) ═══
# Color tones + specific colors + seasons + occasions + styles + fits
TAGS_ALLOWED = [
    # Color tones
    "trung tính", "sáng", "đậm",
    # Specific colors
    "đen", "trắng", "xám", "nâu", "vàng", "hồng", "xanh da trời", "xanh than", "tím", "đỏ", "cam", "xanh lá",
    # Seasons
    "mùa hè", "mùa đông", "mùa xuân", "mùa thu",
    # Occasions
    "đi học", "đi chơi", "đi tiệc", "dạo phố", "mặc nhà", "đi ngủ",
    "thể thao", "đi dã ngoại", "đi biển", "đi làm", "giữ ấm", "thoáng mát",
    # Styles
    "thanh lịch", "năng động", "basic", "cá tính", "dễ thương", "trẻ trung", "tối giản", "smart casual",
    # Fit
    "oversize", "slim", "regular", "wide leg", "cropped", "relaxed"
]

TAG_GENERATION_PROMPT = """Bạn là chuyên gia gán tag thời trang cho sản phẩm Canifa.

NHIỆM VỤ:
Đọc mô tả sản phẩm và chọn TỐI ĐA 6 tags phù hợp nhất từ DANH SÁCH TAGS HỢP LỆ.

DANH SÁCH TAGS HỢP LỄ:
""" + "\n".join(TAGS_ALLOWED) + """

QUY TẮC:
1. CHỈ dùng tags trong danh sách trên — KHÔNG bịa tag mới
2. Chọn tối đa 6 tags, tối thiểu 2 tags
3. Có thể chọn từ bất kỳ nhóm nào (màu sắc, mùa, dịp, phong cách, dáng)
4. Chỉ trả về một list Python (ví dụ: ["đi làm", "thoáng mát", "mùa hè"])
5. KHÔNG giải thích, KHÔNG markdown

MÔ TẢ SẢN PHẨM:
{product_description}

OUTPUT (chỉ list, không gì khác):
"""

async def generate_tags_for_product(client: httpx.AsyncClient, code: str, name: str, clean_desc: str) -> Optional[List[str]]:
    product_description = f"Sản phẩm: {name or 'Trống'}\nMô tả:\n{clean_desc or 'Không có thông tin'}"
    prompt_hydrated = TAG_GENERATION_PROMPT.format(product_description=product_description)

    payload = {
        "model": "openai/gpt-oss-120b",
        "messages": [
            {"role": "user", "content": prompt_hydrated}
        ],
        "temperature": 0.1
    }

    base_url = "https://api.groq.com/openai/v1/chat/completions"

    models_to_try = [
        "gpt-oss-120b", "gpt-oss-20b", "gpt-oss-safeguard-20b",
        "llama-4-scout", "qwen3-32b"
    ]

    for model_name in models_to_try:
        payload["model"] = model_name
        try:
            api_key = _next_groq_key()
            auth_header = {"Authorization": f"Bearer {api_key}"}
            resp = await client.post(base_url, json=payload, headers=auth_header, timeout=30.0)
            if resp.status_code == 200:
                data = resp.json()
                content_str = data["choices"][0]["message"]["content"].strip()
                if content_str.startswith("```json"): content_str = content_str[7:]
                if content_str.startswith("```python"): content_str = content_str[9:]
                if content_str.startswith("```"): content_str = content_str[3:]
                if content_str.endswith("```"): content_str = content_str[:-3]
                import ast
                try:
                    res_list = ast.literal_eval(content_str.strip())
                    if isinstance(res_list, list): return res_list
                except:
                    try:
                        res_json = json.loads(content_str.strip())
                        if isinstance(res_json, list): return res_json
                    except: pass
            elif resp.status_code == 429:
                import re
                wait_time = 15.0
                match = re.search(r'try again in ([0-9.]+)s', resp.text)
                if match: wait_time = max(float(match.group(1)), 5.0)
                logger.warning(f"Rate limited (429) for {model_name}. Sleeping {wait_time}s")
                await asyncio.sleep(wait_time)
            # else: logger.error(f"HTTP {resp.status_code} for {model_name}: {resp.text}")
        except Exception as e:
            await asyncio.sleep(5)
            continue
    return None

async def handle_tags_job(redis, job_data, client, conn):
    job_id = job_data.get("job_id")
    codes = job_data.get("codes", [])
    if not job_id or not codes: return

    logger.info(f"🚀 Found Tags Job {job_id} with {len(codes)} items")
    progress_key = f"tags:progress:{job_id}"

    sem = asyncio.Semaphore(BATCH_CONCURRENCY)
    done_count = err_count = 0

    async def _worker(code: str):
        nonlocal done_count, err_count
        if not code: return
        async with sem:
            cur = None
            try:
                cur = conn.cursor()
                cur.execute(f"SELECT product_name, product_line, clean_description, tags FROM {PG_TABLE} WHERE internal_ref_code = %s", (code,))
                row = cur.fetchone()
                if not row:
                    err_count += 1
                    return
                name, p_line, clean_desc, current_tags = row

                # Parse current tags (may be JSON string or list)
                old_tags = []
                if current_tags:
                    if isinstance(current_tags, str):
                        try: old_tags = json.loads(current_tags)
                        except: old_tags = []
                    elif isinstance(current_tags, list):
                        old_tags = current_tags
                    # Ensure list
                    if not isinstance(old_tags, list):
                        old_tags = []

                full_name = name or p_line or ''
                tags = None
                while not tags:
                    tags = await generate_tags_for_product(client, code, full_name, clean_desc)
                    if not tags: await asyncio.sleep(3)

                if tags and isinstance(tags, list):
                    # Filter: only keep tags in TAGS_ALLOWED (no prefix check)
                    clean_tags = [t for t in tags if t in TAGS_ALLOWED]
                    # Merge: keep old tags, add new ones not already present
                    merged_tags = old_tags[:]
                    for t in clean_tags:
                        if t not in merged_tags:
                            merged_tags.append(t)

                    cur.execute(
                        f"UPDATE {PG_TABLE} SET tags = %s::jsonb, updated_at = NOW() WHERE internal_ref_code = %s",
                        (json.dumps(merged_tags, ensure_ascii=False), code)
                    )
                    conn.commit()
                else: err_count += 1
            except Exception as e:
                logger.error(f"❌ Batch tags failed for {code}: {e}")
                err_count += 1
                if conn: conn.rollback()
            finally:
                if cur: cur.close()
                done_count += 1
                if done_count % 10 == 0 or done_count == len(codes):
                    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"🎉 Tags Job {job_id} finished. Total: {len(codes)}, Errors: {err_count}")


async def handle_ai_edit_job(redis, job_data, client, conn):
    job_id = job_data.get("job_id")
    codes = job_data.get("codes", [])
    target_field = job_data.get("target_field")
    instruction = job_data.get("instruction")
    mode = job_data.get("mode")

    if not job_id or not codes: return
    logger.info(f"🚀 Found AI Edit Job {job_id} with {len(codes)} items.")

    progress_key = f"ai_edit:progress:{job_id}"
    field_label = _FIELD_LABELS.get(target_field, target_field)

    done_count = err_count = 0
    results_history = []

    # We do sequential for high safety so we don't spam codex completely.
    # 5 concurrent workers max for ai-edit
    sem = asyncio.Semaphore(5)

    async def _worker(code: str):
        nonlocal done_count, err_count, results_history
        if not code: return
        async with sem:
            cur = None
            try:
                cur = conn.cursor()
                cur.execute(f"SELECT product_name, product_line, description_data FROM {PG_TABLE} WHERE internal_ref_code = %s", (code,))
                row = cur.fetchone()
                if not row:
                    err_count += 1
                    results_history.append({"code": code, "status": "not_found"})
                    return

                name, p_line, desc_data = row

                if isinstance(desc_data, str):
                    try: desc_data = json.loads(desc_data)
                    except: desc_data = {}
                elif desc_data is None:
                    desc_data = {}

                current_value = desc_data.get(target_field, "")

                context_fields = {
                    "ten_san_pham": desc_data.get("ten_san_pham", name or ""),
                    "phong_cach": desc_data.get("phong_cach", ""),
                    "chat_lieu": desc_data.get("chat_lieu", ""),
                    "gioi_tinh": desc_data.get("gioi_tinh", ""),
                    "mua": desc_data.get("mua", ""),
                    "do_tuoi": desc_data.get("do_tuoi", ""),
                    "product_line": p_line or "",
                }
                context_str = "\\n".join(f"- {k}: {v}" for k, v in context_fields.items() if v)

                if target_field == "huong_dan_size":
                    # Specialized prompt for AI Size Guide
                    base_prompt = f"""Bạn là Stylist thời trang chuyên nghiệp của Canifa. Dựa trên thông tin sản phẩm và {('nội dung hiện tại cùng ' if mode == 'append' else '')}Bảng thông số size (được cung cấp trong YÊU CẦU), hãy sinh ra đoạn văn bản Hướng dẫn chọn size tự nhiên, có tâm và chính xác nhất.

THÔNG TIN SẢN PHẨM:
{context_str}

{f"GIÁ TRỊ HIỆN TẠI:\\n{current_value}\\n" if mode == 'append' and current_value else ""}
YÊU CẦU (Chứa bảng size thô / hướng dẫn custom):
{instruction}

QUY TẮC BẮT BUỘC ('Hướng dẫn chọn size'):
- LIỆT KÊ các size có sẵn dựa trên bảng thông số. Trình bày LẠI bảng kích cỡ một cách dễ nhìn (bao gồm Size, Chiều cao, Cân nặng, Rộng ngực, Rộng mông nếu có trong dữ liệu thô).
- ĐƯA RA LỜI KHUYÊN DỄ HIỂU (Ví dụ: Bạn cao 170cm, nặng 65kg -> nên chọn size M).
- MỤC TIPS MẶC ĐẸP: Nhất thiết phải có mục "Tips mặc đẹp từ Stylist Canifa:". Dựa vào "phong cách", "chất liệu", "thông số ngực/mông/đùi" kết hợp với "Yêu cầu", để đưa ra khuyên về form dáng (VD: Dáng Regular Fit thì nên chọn đúng size, nếu có vòng bụng lớn hoặc thích mặc oversize thì khuyên tăng 1 size).
- LƯU Ý KÍCH THƯỚC: Luôn kèm một lưu ý nhỏ cuối cùng (VD: Số đo cụ thể vòng ngực/mông có thể chênh lệch 1-2cm tùy độ co giãn thực tế của chất liệu).
- FORMAT GIAO DIỆN: Dùng bullet point rõ ràng dễ đọc.
- {'GIỮ NGUYÊN nội dung cũ, CHỈ THÊM phần phân tích mới vào. ' if mode == 'append' and current_value else ''}Trả về chỉ nội dung trường, KHÔNG giải thích.
"""
                    prompt = base_prompt
                elif mode == "append" and current_value:
                    prompt = f"""Bạn là Stylist thời trang Canifa. Dựa trên thông tin sản phẩm dưới đây, hãy BỔ SUNG thêm nội dung cho trường "{field_label}".

THÔNG TIN SẢN PHẨM:
{context_str}

GIÁ TRỊ HIỆN TẠI của trường "{field_label}":
{current_value}

YÊU CẦU BỔ SUNG:
{instruction}

QUY TẮC:
- GIỮ NGUYÊN nội dung cũ, CHỈ THÊM phần mới vào
- Viết tự nhiên, phù hợp với sản phẩm cụ thể này
- Dùng cùng format/style với nội dung cũ (nếu dùng dấu · thì thêm cũng dùng dấu ·)
- Trả về TOÀN BỘ nội dung trường (cũ + mới), KHÔNG giải thích

Trả về chỉ nội dung trường, không có gì khác."""
                else:
                    prompt = f"""Bạn là Stylist thời trang Canifa. Dựa trên thông tin sản phẩm dưới đây, hãy VIẾT nội dung cho trường "{field_label}".

THÔNG TIN SẢN PHẨM:
{context_str}

YÊU CẦU:
{instruction}

QUY TẮC:
- Viết tự nhiên, phù hợp với sản phẩm cụ thể này
- Phong cách giống Stylist thời trang cao cấp
- Trả về chỉ nội dung trường, KHÔNG giải thích

Trả về chỉ nội dung trường, không có gì khác."""

                messages = [
                    {"role": "system", "content": "Bạn là AI phụ trách nội dung sản phẩm thời trang Canifa. Trả lời ngắn gọn, chỉ trả về nội dung được yêu cầu."},
                    {"role": "user", "content": prompt},
                ]

                # We import _call_codex from bulk_ops_route, which uses its own _next_groq_key internally, that's fine.
                ai_output = await _call_codex(messages, max_tokens=500, temperature=0.7)
                new_value = ai_output.strip().strip('"').strip("'")

                desc_data[target_field] = new_value
                new_clean = _render_description_text(desc_data, name or "")

                cur.execute(
                    f"UPDATE {PG_TABLE} SET description_data = %s::jsonb, clean_description = %s, updated_at = NOW() WHERE internal_ref_code = %s",
                    (json.dumps(desc_data, ensure_ascii=False), new_clean, code)
                )
                conn.commit()

                done_count += 1
                results_history.append({
                    "code": code,
                    "name": name or "",
                    "status": "success",
                    "field": target_field,
                    "old_value": current_value[:100] if current_value else "",
                    "new_value": new_value[:150]
                })

            except Exception as e:
                logger.error(f"❌ AI edit error for {code}: {e}")
                err_count += 1
                results_history.append({"code": code, "status": "error", "error": str(e)})
                if conn: conn.rollback()
            finally:
                if cur: cur.close()
                if done_count % 5 == 0 or (done_count + err_count) == len(codes):
                    # We only store last 20 results to avoid massive redis payload
                    trim_history = results_history[-20:]
                    await redis.hset(progress_key, mapping={
                        "done": done_count,
                        "errors": err_count,
                        "current_code": code,
                        "results": json.dumps(trim_history, ensure_ascii=False)
                    })

    # create tasks and run
    await asyncio.gather(*[_worker(code) for code in codes])

    await redis.hset(progress_key, "is_running", "false")
    logger.info(f"🏁 AI Edit Job {job_id} completed: {done_count}/{len(codes)} (errors: {err_count})")


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

    logger.info("👷 Unified AI Worker started (Tags & AI Edit).")
    logger.info("Waiting for jobs on 'tags:queue' & 'ai_edit:queue'...")

    async with httpx.AsyncClient() as client:
        while True:
            try:
                item = await redis.brpop(["tags:queue", "ai_edit:queue"], timeout=2)
                if not item: continue

                queue_bytes, payload = item
                queue_name = queue_bytes.decode('utf-8') if isinstance(queue_bytes, bytes) else queue_bytes
                job_data = json.loads(payload)

                conn = get_pooled_connection_compat()

                try:
                    if queue_name == "tags:queue":
                        await handle_tags_job(redis, job_data, client, conn)
                    elif queue_name == "ai_edit:queue":
                        await handle_ai_edit_job(redis, job_data, client, conn)
                finally:
                    conn.close()

            except Exception as e:
                logger.error(f"Worker loop error: {e}")
                await asyncio.sleep(5)

if __name__ == "__main__":
    if sys.platform == 'win32':
        asyncio.set_event_loop_policy(asyncio.WindowsSelectorEventLoopPolicy())
    asyncio.run(main())
