"""
Tags Direct API — Batch tagging thẳng vào PostgreSQL, không cần Redis/Worker.
- POST /api/tags-direct/run        → chạy batch tagging (background asyncio task)
- GET  /api/tags-direct/status     → poll tiến trình
- POST /api/tags-direct/stop       → dừng
- POST /api/tags-direct/single     → sinh tag cho 1 sản phẩm cụ thể (sync, trả về ngay)
"""

import ast
import asyncio
import json
import logging
import re
import time
from typing import Optional

import httpx
from fastapi import APIRouter, BackgroundTasks
from fastapi.responses import JSONResponse

from common.pool_wrapper import get_pooled_connection_compat

logger = logging.getLogger(__name__)
router = APIRouter(prefix="/api/tags-direct", tags=["Tags Direct"])

PG_TABLE = "dashboard_canifa.ultra_descriptions"

# ── Groq keys round-robin ──────────────────────────────────────────────────
GROQ_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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",
]
_key_idx = 0

def _next_key() -> str:
    global _key_idx
    k = GROQ_KEYS[_key_idx % len(GROQ_KEYS)]
    _key_idx += 1
    return k

GROQ_MODELS = [
    "openai/gpt-oss-120b",
    "openai/gpt-oss-20b",
    "llama-4-scout-17b-16e-instruct",
    "qwen/qwen3-32b",
]

VALID_PREFIXES = ("occ:", "wthr:", "func:", "style:", "fit:")

TAG_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 (bao gồm Tên sản phẩm chứa mã màu) và trích xuất 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 — Màu sắc & Tông màu ---
(Tự nhận diện màu từ tên SP. VD: "Đen/ Black" -> "Màu đen", "Xám/ Gray" -> "Màu xám")
- Màu sắc: Màu đen, Màu trắng, Màu xám, Màu xanh navy, Màu xanh da trời, Màu be, Màu nâu, Màu hồng, Màu đỏ, Màu vàng...
- Tông màu: Tông trung tính (Neutral), Tông màu pastel, Tông màu rực rỡ, Tông màu tối, Tông màu sáng.

--- TRỤC 2 — Hoàn cảnh & Dịp ---
Đi làm / Công sở, Đi chơi / dạo phố, Mặc nhà, Tập thể thao, Đi tiệc, Đi biển, Du lịch, Đi học.

--- TRỤC 3 — Mùa vụ ---
Mùa hè, Mùa đông, Cả 2 mùa, Giao mùa.

--- TRỤC 4 — Phong cách & Tính năng ---
Cơ bản (Basic), Năng động, Thanh lịch, Dễ thương, Cá tính, Tối giản, Smart Casual, Thoáng mát, Giữ ấm, Thấm hút, Nhanh khô, Chống UV, Cản gió.

--- TRỤC 5 — Dáng & Kiểu cắt ---
Oversize, Slim, Regular, Wide leg, Cropped, Relaxed.

QUY TẮC BẮT BUỘC:
1. Bắt buộc phải có tag Hoàn cảnh & Dịp, Mùa vụ (nếu có thể suy luận).
2. Nếu Tên sản phẩm có màu, ưu tiên trả về tag Màu sắc và Tông màu.
3. KHÔNG gán "Đi tiệc" cho áo phông/đồ trẻ em/đồ mặc nhà.
4. Chọn tối đa 6 tags, tối thiểu 3 tags.
5. Chỉ trả về một Python list duy nhất chứa các chuỗi (string), KHÔNG giải thích, KHÔNG markdown.
Ví dụ: ['Màu đen', 'Tông trung tính (Neutral)', 'Đi chơi / dạo phố', 'Cả 2 mùa', 'Cơ bản (Basic)']

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

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

# ── In-memory state (singleton per process) ───────────────────────────────
_state = {
    "is_running": False,
    "total": 0,
    "done": 0,
    "errors": 0,
    "current_code": None,
    "stop_requested": False,
    "started_at": None,
    "finished_at": None,
    "last_errors": [],   # last 10 error codes
}
_task: Optional[asyncio.Task] = None


# ── Helpers ────────────────────────────────────────────────────────────────

async def _call_ai(client: httpx.AsyncClient, text: str) -> Optional[list]:
    """Call Groq with model fallback, return list of valid tags or None."""
    global _state
    payload = {
        "messages": [{"role": "user", "content": text}],
        "temperature": 0.1,
        "max_tokens": 1000,
    }
    for model in GROQ_MODELS:
        payload["model"] = model
        for _ in range(len(GROQ_KEYS)):
            if _state.get("stop_requested"):
                return None
            key = _next_key()
            try:
                resp = await client.post(
                    "https://api.groq.com/openai/v1/chat/completions",
                    json=payload,
                    headers={"Authorization": f"Bearer {key}"},
                    timeout=25.0,
                )
                if resp.status_code == 200:
                    content = resp.json()["choices"][0]["message"]["content"].strip()
                    # strip code fences
                    content = re.sub(r"^```[a-z]*\n?", "", content)
                    content = re.sub(r"```$", "", content).strip()
                    try:
                        result = ast.literal_eval(content)
                    except Exception:
                        try:
                            result = json.loads(content)
                        except Exception as e:
                            print(f"[tags-direct] Parse Error. Content: {content}")
                            result = []
                    
                    if isinstance(result, list):
                        clean = [str(t) for t in result]
                        return clean
                    else:
                        with open("direct_tag.log", "a", encoding="utf-8") as f: f.write(f"[tags-direct] Expected list, got {type(result)}. Content: {content}\n")
                        return None
                elif resp.status_code == 429:
                    match = re.search(r"try again in ([0-9.]+)s", resp.text)
                    wait = float(match.group(1)) if match else 10.0
                    with open("direct_tag.log", "a", encoding="utf-8") as f: f.write(f"[tags-direct] Rate limit (429) model={model}. Wait {wait}s.\n")
                    wait = min(wait, 15.0) # Cap wait at 15s
                    for _ in range(int(wait)):
                        if _state.get("stop_requested"): return None
                        await asyncio.sleep(1)
                else:
                    with open("direct_tag.log", "a", encoding="utf-8") as f: f.write(f"[tags-direct] Groq Error {resp.status_code}: {resp.text}\n")
            except Exception as e:
                with open("direct_tag.log", "a", encoding="utf-8") as f: f.write(f"[tags-direct] Exception model={model}: {e}\n")
                for _ in range(2):
                    if _state.get("stop_requested"): return None
                    await asyncio.sleep(1)
    return None


def _fetch_untagged(limit: int = 2000, force_all: bool = False) -> list[tuple]:
    """Return list of (code, name, clean_description) for products."""
    conn = get_pooled_connection_compat()
    try:
        cur = conn.cursor()
        condition = "1=1" if force_all else "tags IS NULL"
        cur.execute(
            f"""SELECT internal_ref_code, product_name, clean_description
                FROM {PG_TABLE}
                WHERE {condition}
                  AND clean_description IS NOT NULL AND clean_description != ''
                LIMIT %s""",
            (limit,),
        )
        rows = cur.fetchall()
        cur.close()
        return rows
    finally:
        conn.close()


def _save_tags(code: str, tags: list) -> None:
    conn = get_pooled_connection_compat()
    try:
        cur = conn.cursor()
        cur.execute(
            f"UPDATE {PG_TABLE} SET tags = %s::jsonb, updated_at = NOW() WHERE internal_ref_code = %s",
            (json.dumps(tags, ensure_ascii=False), code),
        )
        conn.commit()
        cur.close()
    finally:
        conn.close()


# ── Core batch runner ──────────────────────────────────────────────────────

async def _run_batch(force_all: bool = False):
    global _state
    logger.info(f"🏷️ [tags-direct] Batch started (force_all={force_all})")

    try:
        rows = await asyncio.get_event_loop().run_in_executor(None, _fetch_untagged, 2000, force_all)
    except Exception as e:
        logger.error(f"[tags-direct] DB fetch error: {e}")
        _state["is_running"] = False
        return

    if not rows:
        logger.info("[tags-direct] Không có sản phẩm nào cần gán tag.")
        _state.update({"is_running": False, "total": 0, "done": 0, "finished_at": time.time()})
        return

    _state.update({
        "total": len(rows),
        "done": 0,
        "errors": 0,
        "last_errors": [],
        "current_code": rows[0][0] if rows else None,
    })

    sem = asyncio.Semaphore(2)  # Reduced to 2 to prevent Groq API rate limit exhaustion

    async def _process(code: str, name: str, desc: str):
        if _state["stop_requested"]:
            return
        async with sem:
            try:
                _state["current_code"] = code
                product_text = f"Sản phẩm: {name or ''}\nMô tả:\n{desc or ''}"
                prompt = TAG_PROMPT.format(product_description=product_text[:3000])

                tags = await _call_ai(client, prompt)
                if tags is not None:
                    await asyncio.get_event_loop().run_in_executor(None, _save_tags, code, tags)
                    _state["done"] += 1
                else:
                    _state["errors"] += 1
                    _state["last_errors"].append(code)
                    _state["last_errors"] = _state["last_errors"][-10:]
            except Exception as e:
                with open("direct_tag.log", "a", encoding="utf-8") as f: f.write(f"[tags-direct] error for {code}: {e}\n")
                _state["errors"] += 1

    async with httpx.AsyncClient() as client:
        tasks = [_process(r[0], r[1] or "", r[2] or "") for r in rows]
        await asyncio.gather(*tasks)

    _state.update({
        "is_running": False,
        "stop_requested": False,
        "finished_at": time.time(),
    })
    logger.info(f"✅ [tags-direct] Done: {_state['done']}/{_state['total']}, errors: {_state['errors']}")


# ── Endpoints ──────────────────────────────────────────────────────────────

@router.post("/run", summary="Chạy batch gán tags thẳng vào PostgreSQL (không cần Redis)")
async def run_tags_direct(background_tasks: BackgroundTasks, force_all: bool = False):
    global _state, _task

    if _state["is_running"]:
        return JSONResponse(
            status_code=409,
            content={"status": "error", "message": "Đang chạy rồi! Đợi xong hoặc bấm Stop."},
        )

    _state.update({
        "is_running": True,
        "stop_requested": False,
        "started_at": time.time(),
        "finished_at": None,
        "total": 0,
        "done": 0,
        "errors": 0,
        "current_code": None,
        "last_errors": [],
    })

    background_tasks.add_task(_run_batch, force_all)
    return {"status": "success", "message": "🚀 Batch tagging đã bắt đầu chạy nền!"}


@router.get("/status", summary="Poll tiến trình batch tagging")
async def tags_direct_status():
    s = _state.copy()
    pct = round(s["done"] / s["total"] * 100, 1) if s["total"] > 0 else 0
    elapsed = round(time.time() - s["started_at"], 0) if s["started_at"] else 0
    return {
        "is_running": s["is_running"],
        "total": s["total"],
        "done": s["done"],
        "errors": s["errors"],
        "current_code": s["current_code"],
        "percent": pct,
        "elapsed_sec": elapsed,
        "last_errors": s["last_errors"],
    }


@router.post("/stop", summary="Dừng batch tagging")
async def stop_tags_direct():
    if not _state["is_running"]:
        return {"status": "ok", "message": "Không có task nào đang chạy."}
    _state["stop_requested"] = True
    return {"status": "ok", "message": "✋ Đã yêu cầu dừng. Các job đang chạy sẽ kết thúc sau giây lát."}


@router.post("/single", summary="Sinh tag cho 1 sản phẩm cụ thể (sync)")
async def single_tag(body: dict):
    code = (body.get("internal_ref_code") or "").strip()
    if not code:
        return JSONResponse(status_code=400, content={"status": "error", "message": "Thiếu internal_ref_code"})

    try:
        conn = get_pooled_connection_compat()
        cur = conn.cursor()
        cur.execute(
            f"SELECT product_name, clean_description FROM {PG_TABLE} WHERE internal_ref_code = %s",
            (code,),
        )
        row = cur.fetchone()
        cur.close()
        conn.close()
    except Exception as e:
        return JSONResponse(status_code=500, content={"status": "error", "message": str(e)})

    if not row:
        return JSONResponse(status_code=404, content={"status": "error", "message": f"Không tìm thấy {code}"})

    name, desc = row
    prompt = TAG_PROMPT.format(product_description=f"Sản phẩm: {name or ''}\nMô tả:\n{(desc or '')[:3000]}")

    async with httpx.AsyncClient() as client:
        tags = await _call_ai(client, prompt)

    if not tags:
        return JSONResponse(status_code=500, content={"status": "error", "message": "AI không trả về tags hợp lệ"})

    try:
        _save_tags(code, tags)
    except Exception as e:
        return JSONResponse(status_code=500, content={"status": "error", "message": f"Lưu DB lỗi: {e}"})

    return {"status": "success", "code": code, "tags": tags}


@router.get("/count-untagged", summary="Đếm số sản phẩm chưa có tag")
async def count_untagged():
    try:
        conn = get_pooled_connection_compat()
        cur = conn.cursor()
        cur.execute(
            f"SELECT COUNT(*) FROM {PG_TABLE} WHERE tags IS NULL"
        )
        count = cur.fetchone()[0]
        cur.close()
        conn.close()
        return {"status": "success", "untagged": count}
    except Exception as e:
        return JSONResponse(status_code=500, content={"status": "error", "message": str(e)})
