"""Detect main regions in a label image via Qwen2.5-VL (DashScope).

Workflow
--------
1. Read the original image; record its exact dimensions (orig_w × orig_h).
2. Downscale a copy to fit within ``api_max_side`` for the API call
   (faster upload, lower token cost).  Record api_w × api_h.
3. Send the downscaled image to Qwen VL.
4. Parse the response coordinates (which are relative to the api image):
   a. Qwen2.5-VL grounding tokens  <|box_start|>(x1,y1),(x2,y2)<|box_end|>
      – normalised to [0, 1000] of the *api* image.
   b. Fallback: JSON array ``[{"label": "...", "bbox": [x1,y1,x2,y2]}, ...]``
      – pixel values in the *api* image coordinate space.
5. Scale coordinates back to the original image space:
   x_orig = round(x_api * orig_w / api_w)
6. Crop from the **original** high-resolution file → full-quality output.
"""
from __future__ import annotations

import base64
import json
import logging
import os
import re
from pathlib import Path
from typing import NamedTuple

logger = logging.getLogger(__name__)

_DASHSCOPE_BASE_URL_DEFAULT = "https://dashscope.aliyuncs.com/compatible-mode/v1"
# Model alias: the 7B variant is faster ("flash"); swap for 72B for higher accuracy
DEFAULT_MODEL = "qwen2.5-vl-7b-instruct"

_GROUNDING_RE = re.compile(
    r"<\|object_ref_start\|>(.*?)<\|object_ref_end\|>"
    r"<\|box_start\|>\((\d+),(\d+)\),\((\d+),(\d+)\)<\|box_end\|>",
    re.DOTALL,
)

_DEFAULT_PROMPT = (
    "请检测图像中食品包装标签的所有主要内容区域（如：主产品信息表格、"
    "营养成分表、标题、配料表、厂商信息、条码区等）。"
    "以JSON列表输出，格式为：\n"
    '[{"label": "区域名称", "bbox": [x1, y1, x2, y2]}, ...]'
    "\n坐标为实际像素值（整数），原点在左上角。"
)


class Region(NamedTuple):
    label: str
    x1: int
    y1: int
    x2: int
    y2: int

    @property
    def width(self) -> int:
        return self.x2 - self.x1

    @property
    def height(self) -> int:
        return self.y2 - self.y1


def _read_dotenv(path: Path) -> dict[str, str]:
    """Parse a simple KEY=VALUE .env file into a dict."""
    result: dict[str, str] = {}
    if not path.exists():
        return result
    for raw in path.read_text(encoding="utf-8").splitlines():
        line = raw.strip()
        if not line or line.startswith("#") or "=" not in line:
            continue
        k, v = line.split("=", 1)
        result[k.strip()] = v.strip().strip('"').strip("'")
    return result


def _load_env() -> dict[str, str]:
    """Merge .env files (project root → parent → home) into a single dict."""
    merged: dict[str, str] = {}
    for p in [
        Path(__file__).resolve().parents[2] / ".env",
        Path(__file__).resolve().parents[3] / ".env",
        Path.home() / ".env",
    ]:
        merged.update(_read_dotenv(p))
    return merged


def _get_api_key() -> str:
    """Read DASHSCOPE_API_KEY from env vars then .env files."""
    val = os.environ.get("DASHSCOPE_API_KEY", "").strip()
    if val:
        return val
    return _load_env().get("DASHSCOPE_API_KEY", "")


def _get_base_url() -> str:
    """Read DASHSCOPE_BASE_URL from env vars then .env files."""
    val = os.environ.get("DASHSCOPE_BASE_URL", "").strip()
    if val:
        return val
    return _load_env().get("DASHSCOPE_BASE_URL", _DASHSCOPE_BASE_URL_DEFAULT)


def _encode_image_for_api(
    image_path: Path,
    max_side: int = 1024,
) -> tuple[str, int, int]:
    """Downscale image to fit within *max_side* × *max_side*, encode as PNG base64.

    Returns
    -------
    b64 : str
        Base64-encoded PNG of the (possibly resized) image.
    api_w : int
        Width  of the image that was actually sent to the API.
    api_h : int
        Height of the image that was actually sent to the API.
    """
    import io
    from PIL import Image

    with Image.open(image_path) as img:
        # Convert to RGB so PNG encoding always works
        img = img.convert("RGB")
        orig_w, orig_h = img.size

        if max(orig_w, orig_h) > max_side:
            scale = max_side / max(orig_w, orig_h)
            api_w = max(1, round(orig_w * scale))
            api_h = max(1, round(orig_h * scale))
            api_img = img.resize((api_w, api_h), Image.LANCZOS)
        else:
            api_w, api_h = orig_w, orig_h
            api_img = img

        buf = io.BytesIO()
        api_img.save(buf, format="PNG")

    b64 = base64.b64encode(buf.getvalue()).decode()
    return b64, api_w, api_h


def _parse_grounding_tokens(text: str, api_w: int, api_h: int) -> list[Region]:
    """Parse <|box_start|>(x1,y1),(x2,y2)<|box_end|> tokens.

    Qwen2.5-VL normalises coordinates to [0, 1000] of the *api* image.
    Returns pixel coordinates in the api image space.
    """
    regions: list[Region] = []
    for m in _GROUNDING_RE.finditer(text):
        label = m.group(1).strip()
        x1 = round(int(m.group(2)) * api_w / 1000)
        y1 = round(int(m.group(3)) * api_h / 1000)
        x2 = round(int(m.group(4)) * api_w / 1000)
        y2 = round(int(m.group(5)) * api_h / 1000)
        regions.append(Region(label, x1, y1, x2, y2))
    return regions


def _parse_json_regions(text: str) -> list[Region]:
    """Fallback: extract bbox from a JSON object or array in the response."""
    clean = re.sub(r"<\|[^|]+\|>", "", text)
    clean = re.sub(r"```[a-z]*", "", clean).strip("`").strip()

    def _extract_bbox(item: dict) -> list | None:
        """Try multiple known bbox key names, including nested dicts."""
        for key in ("bbox", "bbox_2d", "box", "coordinates", "bounding_box"):
            v = item.get(key)
            if isinstance(v, (list, tuple)) and len(v) >= 4:
                return list(v)
            # e.g. {"label": {"bbox_2d": [...]}}
            if isinstance(v, dict):
                inner = _extract_bbox(v)
                if inner:
                    return inner
        # Recurse into all dict values
        for v in item.values():
            if isinstance(v, dict):
                inner = _extract_bbox(v)
                if inner:
                    return inner
        return None

    def _region_from_dict(item: dict) -> Region | None:
        bbox = _extract_bbox(item)
        if not bbox or len(bbox) < 4:
            return None
        # Label: try common keys; skip if value is a nested dict
        raw_label = (item.get("label") or item.get("name") or item.get("type") or "主内容区")
        label = raw_label if isinstance(raw_label, str) else "主内容区"
        x1, y1, x2, y2 = (int(v) for v in bbox[:4])
        return Region(label, x1, y1, x2, y2)

    # Try single JSON object first (default prompt returns one object)
    obj_start, obj_end = clean.find("{"), clean.rfind("}")
    if obj_start != -1 and obj_end > obj_start:
        try:
            obj = json.loads(clean[obj_start : obj_end + 1])
            if isinstance(obj, dict):
                r = _region_from_dict(obj)
                if r:
                    return [r]
        except json.JSONDecodeError:
            pass

    # Fallback to JSON array
    arr_start, arr_end = clean.find("["), clean.rfind("]")
    if arr_start == -1 or arr_end <= arr_start:
        return []
    try:
        items = json.loads(clean[arr_start : arr_end + 1])
    except json.JSONDecodeError:
        return []
    regions: list[Region] = []
    for item in items:
        if isinstance(item, dict):
            r = _region_from_dict(item)
            if r:
                regions.append(r)
    return regions


def detect_regions(
    image_path: Path,
    api_key: str | None = None,
    model: str = DEFAULT_MODEL,
    prompt: str = _DEFAULT_PROMPT,
    api_max_side: int = 1024,
) -> tuple[list[Region], str]:
    """Call Qwen VL to detect main regions.

    The image is downscaled to *api_max_side* before the API call for speed
    and cost efficiency.  Returned ``Region`` coordinates are always mapped
    back to the **original** image pixel space, so ``crop_and_save`` will
    produce full-resolution output.

    Parameters
    ----------
    api_max_side:
        Maximum side length (px) of the image sent to the API.
        Increase for very large originals where detection needs more detail.

    Returns
    -------
    regions : list[Region]
        Bounding boxes in **original** image coordinates.
    raw_response : str
        Full model text (for debugging).
    """
    from openai import OpenAI
    from PIL import Image

    key = api_key or _get_api_key()
    if not key:
        raise RuntimeError(
            "DASHSCOPE_API_KEY not set. "
            "Add it to the project .env or set the environment variable."
        )

    # ── 1. Original dimensions ────────────────────────────────────────────
    with Image.open(image_path) as img:
        orig_w, orig_h = img.size

    # ── 2. Downscale for API; remember api dims for coordinate mapping ─────
    b64, api_w, api_h = _encode_image_for_api(image_path, max_side=api_max_side)
    logger.info(
        "Calling %s on %s  orig=%dx%d → api=%dx%d (scale=%.3f)",
        model, image_path.name, orig_w, orig_h, api_w, api_h, api_w / orig_w,
    )

    # ── 3. API call ───────────────────────────────────────────────────────
    client = OpenAI(api_key=key, base_url=_get_base_url())
    response = client.chat.completions.create(
        model=model,
        messages=[
            {
                "role": "user",
                "content": [
                    {
                        "type": "image_url",
                        "image_url": {"url": f"data:image/png;base64,{b64}"},
                    },
                    {"type": "text", "text": prompt},
                ],
            }
        ],
    )

    raw = response.choices[0].message.content or ""
    logger.debug("Qwen VL raw response:\n%s", raw)

    # ── 4. Parse — coordinates are in api-image space ─────────────────────
    regions = _parse_grounding_tokens(raw, api_w, api_h)
    if regions:
        logger.info("Parsed %d region(s) from grounding tokens", len(regions))
    else:
        regions = _parse_json_regions(raw)
        if regions:
            logger.info("Parsed %d region(s) from JSON fallback", len(regions))

    if not regions:
        logger.warning("No regions parsed from response:\n%s", raw[:400])
        return [], raw

    # ── 5. Scale coordinates back to original image space ─────────────────
    sx, sy = orig_w / api_w, orig_h / api_h
    original_regions = [
        Region(r.label,
               round(r.x1 * sx), round(r.y1 * sy),
               round(r.x2 * sx), round(r.y2 * sy))
        for r in regions
    ]
    logger.info(
        "Coordinates remapped  api(%dx%d) → orig(%dx%d)",
        api_w, api_h, orig_w, orig_h,
    )
    return original_regions, raw


def merge_regions(regions: list[Region], label: str = "主内容区") -> Region:
    """Return the union bounding box of all regions as a single Region."""
    if not regions:
        raise ValueError("Cannot merge empty region list")
    x1 = min(r.x1 for r in regions)
    y1 = min(r.y1 for r in regions)
    x2 = max(r.x2 for r in regions)
    y2 = max(r.y2 for r in regions)
    return Region(label, x1, y1, x2, y2)


def crop_and_save(
    image_path: Path,
    regions: list[Region],
    output_dir: Path,
) -> list[dict]:
    """Crop each region and save as PNG.

    Returns
    -------
    list of dicts with keys: label, bbox, path
    """
    from PIL import Image

    output_dir.mkdir(parents=True, exist_ok=True)
    results: list[dict] = []

    with Image.open(image_path) as img:
        img_w, img_h = img.size
        for i, region in enumerate(regions, start=1):
            # Clamp to image bounds
            x1 = max(0, region.x1)
            y1 = max(0, region.y1)
            x2 = min(img_w, region.x2)
            y2 = min(img_h, region.y2)
            if x2 <= x1 or y2 <= y1:
                logger.warning("Skipping zero-area region: %s", region.label)
                continue
            cropped = img.crop((x1, y1, x2, y2))
            safe_name = re.sub(r"[^\w\u4e00-\u9fff-]", "_", region.label)[:40]
            out_path = output_dir / f"{i:02d}_{safe_name}.png"
            cropped.save(out_path)
            logger.info("Saved region [%s] → %s", region.label, out_path.name)
            results.append({
                "label": region.label,
                "bbox": [x1, y1, x2, y2],
                "path": str(out_path),
            })

    return results