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

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

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


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 bên dưới 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Ệ:
--- TRỤC 1 — Hoàn cảnh & Dịp (occ:) ---
occ:di_lam, occ:di_choi, occ:di_tiec, occ:di_hoc,
occ:mac_nha, occ:the_thao, occ:di_bien, occ:du_lich,
occ:da_ngoai, occ:di_ngu

--- TRỤC 2a — Thời tiết (wthr:) ---
wthr:mua_he, wthr:mua_dong, wthr:giao_mua,
wthr:troi_mua, wthr:troi_nang

--- TRỤC 2b — Tính năng vải (func:) ---
func:thoang_mat, func:giu_am, func:tham_hut,
func:nhanh_kho, func:chong_uv, func:can_gio

--- TRỤC 3 — Phong cách & Vibe (style:) ---
style:thanh_lich, style:nang_dong, style:basic,
style:ca_tinh, style:de_thuong, style:tre_trung,
style:toi_gian, style:smart_casual

--- TRỤC 4 — Dáng & Kiểu cắt (fit:) ---
fit:oversize, fit:slim, fit:regular,
fit:wide_leg, fit:cropped, fit:relaxed

QUY TẮC BẮT BUỘC:
1. Chỉ được dùng tags trong danh sách trên — KHÔNG được bịa tag mới
2. Chọn tối đa 6 tags, tối thiểu 2 tags
3. Cố gắng có ít nhất 1 tag occ: + 1 tag style: hoặc fit: trong kết quả
4. fit: và style: là 2 trục KHÁC NHAU — fit:oversize ≠ style:ca_tinh
5. Nếu mô tả không đủ thông tin để gán 1 trục nào đó, bỏ qua trục đó, KHÔNG đoán mò
6. Chỉ trả về một list Python duy nhất, KHÔNG giải thích, KHÔNG markdown
7. CẤM BẮT BUỘC: KHÔNG gán occ:di_tiec cho áo phông (thun), đồ trẻ em, đồ in hình hoạt hình (Stitch, Mickey...). Chỉ gán occ:di_tiec cho trang phục sang trọng (váy dạ hội, sơ mi dự tiệc, áo blazer, quần âu).

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
                if current_tags and len(current_tags) > 0:
                    done_count += 1
                    return

                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):
                    valid_prefixes = ("occ:", "wthr:", "func:", "style:", "fit:")
                    clean_tags = [t for t in tags if str(t).startswith(valid_prefixes)]
                    cur.execute(
                        f"UPDATE {PG_TABLE} SET tags = %s::jsonb, updated_at = NOW() WHERE internal_ref_code = %s",
                        (json.dumps(clean_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())
