"""
Batch Bedrock vision: one product reference image + instructions once, multiple shelf images.
Invoked by Step Functions Map with { product_id, items: [{ image_id } | { task_id }, ...] }.

Exactly one InvokeModel per group (primary profile only): no adaptive SDK retries beyond default,
no throttle backoff loop, no fallback inference profile — avoids burning quotas with duplicate calls.
"""
import base64
import copy
import json
import logging
import os
import re
from io import BytesIO
from typing import Any, Dict, List, Optional, Tuple
from urllib.parse import urlparse, unquote

import boto3
import requests
from botocore.config import Config
from PIL import Image, ImageOps
from supabase import Client, create_client

logger = logging.getLogger()
logger.setLevel(logging.INFO)

ssm = boto3.client("ssm")
bedrock_primary_once = boto3.client(
    "bedrock-runtime",
    config=Config(retries={"max_attempts": 1}, read_timeout=900, connect_timeout=30),
)

MAX_IMAGE_BYTES = int(5 * 1024 * 1024 * 3 / 4) - 100_000
MIN_RESIZE_HEIGHT = 720


def get_parameter(name: str) -> str:
    return ssm.get_parameter(Name=name, WithDecryption=True)["Parameter"]["Value"]


def get_supabase_client() -> Client:
    return create_client(get_parameter("/supabase/url"), get_parameter("/supabase/anon"))


def get_bedrock_invocation_target() -> Dict[str, str]:
    profile_arn = os.getenv("INFERENCE_PROFILE_ARN")
    if not profile_arn:
        raise RuntimeError("Missing INFERENCE_PROFILE_ARN")
    return {"modelId": profile_arn}


def invoke_bedrock_primary_only(body: bytes, content_type: str, accept: str) -> Tuple[Any, str]:
    """Single InvokeModel on the primary inference profile (no fallback; quota-friendly)."""
    primary = get_bedrock_invocation_target()
    resp = bedrock_primary_once.invoke_model(
        **primary,
        body=body,
        contentType=content_type,
        accept=accept,
    )
    return resp, primary["modelId"]


def get_image_data(supabase: Client, image_id: str) -> Dict:
    image_response = supabase.table("product_images").select("*").eq("id", image_id).single().execute()
    if not image_response.data:
        raise Exception(f"Image with ID {image_id} not found")
    image_record = image_response.data
    product = None
    product_id = image_record.get("product_id")
    if product_id is not None:
        product_response = supabase.table("product").select("*").eq("id", product_id).single().execute()
        if product_response.data:
            product = product_response.data
    task = None
    task_id = image_record.get("task_id")
    if task_id is not None:
        task_response = supabase.table("task").select("*").eq("id", task_id).single().execute()
        if task_response.data:
            task = task_response.data
    return {"image": image_record, "product": product, "task": task, "task_id": task_id}


def get_task_data(supabase: Client, task_id: str) -> Dict:
    task_response = supabase.table("task").select("*").eq("id", task_id).single().execute()
    if not task_response.data:
        raise Exception(f"Task with ID {task_id} not found")
    task = task_response.data
    product = None
    product_id = task.get("product_id")
    if product_id is not None:
        product_response = supabase.table("product").select("*").eq("id", product_id).single().execute()
        if product_response.data:
            product = product_response.data
    return {"image": None, "product": product, "task": task, "task_id": task_id}


def _supabase_public_url_bucket_key(url: str) -> Optional[tuple]:
    try:
        path = unquote((urlparse(url.strip()).path or ""))
        marker = "/storage/v1/object/public/"
        i = path.find(marker)
        if i == -1:
            return None
        rest = path[i + len(marker) :].lstrip("/")
        if "/" not in rest:
            return None
        bucket, key = rest.split("/", 1)
        if not bucket or not key:
            return None
        return (bucket, key)
    except Exception:
        return None


def download_image(supabase: Client, image_path: str) -> bytes:
    if image_path.startswith("http"):
        image_data = None
        parsed_storage = _supabase_public_url_bucket_key(image_path)
        if parsed_storage:
            bucket_name, object_key = parsed_storage
            try:
                image_data = supabase.storage.from_(bucket_name).download(object_key)
            except Exception:
                pass
        if image_data is None:
            headers = {
                "User-Agent": "Mozilla/5.0",
                "Accept": "image/avif,image/webp,image/*,*/*;q=0.8",
            }
            response = requests.get(image_path, headers=headers, timeout=30)
            response.raise_for_status()
            image_data = response.content
    else:
        image_path = image_path.lstrip("/")
        if image_path.startswith("public/"):
            image_path = image_path[7:]
        image_data = supabase.storage.from_("public").download(image_path)
    img = Image.open(BytesIO(image_data))
    img.verify()
    return image_data


def resize_image_to_fit(image_bytes: bytes, max_bytes: int = MAX_IMAGE_BYTES) -> bytes:
    img = Image.open(BytesIO(image_bytes))
    img = ImageOps.exif_transpose(img)
    img = img.convert("RGB")
    w, h = img.size
    for quality in (95, 90, 85, 80, 75, 70, 65, 60, 55, 50):
        out = BytesIO()
        img.save(out, format="JPEG", quality=quality, optimize=True, progressive=True)
        data = out.getvalue()
        if len(data) <= max_bytes:
            return data
    min_scale = max(320 / w, 320 / h, MIN_RESIZE_HEIGHT / h)
    for scale in (0.9, 0.8, 0.7, 0.6, 0.5, 0.4, 0.3, 0.2):
        scale = max(scale, min_scale)
        nw = max(320, int(w * scale))
        nh = max(320, int(h * scale))
        resized = img.resize((nw, nh), Image.Resampling.LANCZOS)
        for quality in (85, 75, 65, 55, 50):
            out = BytesIO()
            resized.save(out, format="JPEG", quality=quality, optimize=True, progressive=True)
            data = out.getvalue()
            if len(data) <= max_bytes:
                return data
    return data


def _encode_image_b64(image_data: bytes, label: str) -> Tuple[str, bytes]:
    if not isinstance(image_data, bytes):
        raise ValueError(f"{label} must be bytes")
    img = Image.open(BytesIO(image_data))
    img.verify()
    image_data = resize_image_to_fit(image_data, MAX_IMAGE_BYTES)
    b64 = base64.b64encode(image_data).decode("utf-8")
    pad = len(b64) % 4
    if pad:
        b64 += "=" * (4 - pad)
    return b64, image_data


def _parse_positions_json(raw: str) -> Optional[list]:
    raw = raw.strip()
    if not raw or raw.upper() == "N/A":
        return []
    start = raw.find("[")
    if start == -1:
        return None
    depth = 0
    end = -1
    for i in range(start, len(raw)):
        if raw[i] == "[":
            depth += 1
        elif raw[i] == "]":
            depth -= 1
            if depth == 0:
                end = i
                break
    if end == -1:
        return None
    try:
        arr = json.loads(raw[start : end + 1])
    except json.JSONDecodeError:
        return None
    if not isinstance(arr, list):
        return None
    out = []
    for item in arr:
        if not isinstance(item, dict):
            continue
        shelf = item.get("shelf")
        pos = item.get("position")
        facings = item.get("facings")
        if shelf is not None and pos is not None and facings is not None:
            out.append({"shelf": int(shelf), "position": str(pos), "facings": int(facings)})
    return out


def _positions_to_string(positions: list) -> str:
    return "; ".join(f"Shelf {p['shelf']}, {p['position']} ({p['facings']} facings)" for p in positions)


def update_analysis_result(
    supabase: Client,
    task_id: str,
    analysis_result: bool,
    analysis_message: str,
    confidence: Optional[str] = None,
    facings_count: Optional[int] = None,
    position: Optional[str] = None,
    has_pi_label: Optional[bool] = None,
    image_id: Optional[str] = None,
    image_path: Optional[str] = None,
) -> None:
    if image_id is not None:
        update_data = {
            "processed_genai": True,
            "genai_result": analysis_result,
            "genai_message": analysis_message,
        }
        if image_path:
            update_data["last_processed_image_path"] = image_path
        supabase.table("product_images").update(update_data).eq("id", image_id).execute()
    task_update_payload = {
        "processed_genai": True,
        "genai_result": analysis_result,
        "genai_message": analysis_message,
        "genai_facing_count": int(facings_count) if facings_count is not None else None,
        "genai_position": position,
        "genai_has_label": bool(has_pi_label) if has_pi_label is not None else None,
    }
    if confidence is not None:
        try:
            conf_val = str(confidence).strip().replace("%", "")
            if float(conf_val) <= 1:
                task_update_payload["genai_confidence"] = int(round(float(conf_val) * 100))
            else:
                task_update_payload["genai_confidence"] = int(round(float(conf_val)))
        except Exception:
            task_update_payload["genai_confidence"] = None
    else:
        task_update_payload["genai_confidence"] = None
    if task_id and task_id != "N/A":
        supabase.table("task").update(task_update_payload).eq("id", task_id).execute()


def update_error_status(supabase: Client, image_id: str, error_msg: str) -> None:
    supabase.table("product_images").update({"processed_genai": False, "genai_error": error_msg}).eq(
        "id", image_id
    ).execute()


def _extract_json_object(text: str) -> Dict[str, Any]:
    """Parse first JSON object from model output (strip markdown fences)."""
    text = text.strip()
    if text.startswith("```"):
        lines = text.split("\n")
        if lines[0].startswith("```"):
            lines = lines[1:]
        if lines and lines[-1].strip() == "```":
            lines = lines[:-1]
        text = "\n".join(lines)
    start = text.find("{")
    end = text.rfind("}")
    if start == -1 or end == -1 or end <= start:
        raise ValueError("No JSON object in Bedrock response")
    return json.loads(text[start : end + 1])


def run_batch(
    supabase: Client,
    product_id: str,
    items: List[Dict[str, str]],
) -> Dict[str, Any]:
    """Load product once; one Bedrock call with ref + N shelves; apply DB updates."""
    pres = supabase.table("product").select("*").eq("id", product_id).single().execute()
    if not pres.data:
        raise ValueError(f"Product {product_id} not found")
    product_row = pres.data
    pname = product_row.get("product") or "N/A"
    pcat = product_row.get("category") or "N/A"
    pbrand = product_row.get("brand") or "N/A"
    pvariant = product_row.get("variant") or "N/A"
    psize = product_row.get("size") or "N/A"
    pref_url = (product_row.get("image_path") or "").strip()
    if not pref_url:
        raise ValueError(f"Product {product_id} has no image_path for reference")

    product_image_bytes = download_image(supabase, pref_url)
    ref_b64, ref_bytes = _encode_image_b64(product_image_bytes, "reference")

    shelf_jobs: List[Dict[str, Any]] = []
    for idx, item in enumerate(items, start=1):
        image_id = item.get("image_id")
        task_id_param = item.get("task_id")
        if image_id:
            data = get_image_data(supabase, image_id)
            image_path = (data["image"] or {}).get("image_path") or (data.get("task") or {}).get("image_url")
            tid = str(data.get("task_id") or "")
        else:
            data = get_task_data(supabase, task_id_param)
            image_path = (data.get("task") or {}).get("image_url")
            tid = str(task_id_param)
            image_id = None
        pdata = data.get("product") or {}
        if pdata and str(pdata.get("id")) != str(product_id):
            raise ValueError(f"product mismatch for task {tid}: expected {product_id}, got {pdata.get('id')}")
        if not image_path:
            raise ValueError(f"No shelf image for task {tid}")
        shelf_bytes = download_image(supabase, image_path)
        task_name = (data.get("task") or {}).get("name") or "N/A"
        shelf_jobs.append(
            {
                "shelf_index": idx,
                "task_id": tid,
                "image_id": image_id,
                "image_path": image_path,
                "task_name": task_name,
                "shelf_bytes": shelf_bytes,
            }
        )

    content: List[Dict[str, Any]] = []
    content.append(
        {
            "type": "text",
            "text": (
                "You are a shelf-verification assistant. Below: (1) ONE reference product image "
                "shared by all evaluations. (2) Multiple SHELF photos — each shelf photo is INDEPENDENT "
                "(different stores/visits). For each shelf_index, decide if the EXACT reference product "
                f"appears on that shelf.\nTARGET: Brand={pbrand}, Product={pname}, Category={pcat}, "
                f"Variant={pvariant}, Size={psize}.\nRules: Match packaging, logo, colors, and size "
                "(when Size is specified) to the reference only. Count front-row facings only; scan all shelves "
                "top-to-bottom in each photo.\nReturn ONLY valid JSON (no markdown), exactly this shape:\n"
                '{"results":[{"shelf_index":<int matching label>,"task_id":"<uuid>","image_id":"<uuid or empty string>",'
                '"FINAL_RESULT":true|false,"CONFIDENCE":"<0-100%> or N/A","FACINGS_COUNT":<int or null>,'
                '"POSITIONS":[{"shelf":N,"position":"left|center-left|center|center-right|right","facings":M}],'
                '"HAS_PI_LABEL":true|false|null,"ANALYSIS_SUMMARY":"<short>"}]}\n'
                "Include one object per shelf_index present. POSITIONS uses [] if not found.",
            ),
        }
    )
    content.append({"type": "text", "text": "REFERENCE PRODUCT IMAGE (same for all shelf checks below):"})
    content.append(
        {
            "type": "image",
            "source": {"type": "base64", "media_type": "image/jpeg", "data": ref_b64},
        }
    )

    for job in shelf_jobs:
        sb64, _ = _encode_image_b64(job["shelf_bytes"], f"shelf{job['shelf_index']}")
        label = (
            f"SHELF_INDEX={job['shelf_index']} TASK_ID={job['task_id']} "
            f"IMAGE_ID={job['image_id'] or ''} TASK_NAME={job['task_name']}"
        )
        content.append({"type": "text", "text": label})
        content.append({"type": "image", "source": {"type": "base64", "media_type": "image/jpeg", "data": sb64}})

    max_tokens = min(16000, 1500 + 800 * len(shelf_jobs))
    request_body = {
        "anthropic_version": "bedrock-2023-05-31",
        "max_tokens": max_tokens,
        "temperature": 0,
        "messages": [{"role": "user", "content": content}],
    }

    log_body = copy.deepcopy(request_body)
    for msg in log_body["messages"]:
        for it in msg["content"]:
            if it.get("type") == "image":
                it["source"]["data"] = "<omitted>"
    logger.info("Bedrock batch request (trimmed): %s", json.dumps(log_body)[:4000])

    response, model_id_used = invoke_bedrock_primary_only(
        body=json.dumps(request_body).encode("utf-8"),
        content_type="application/json",
        accept="application/json",
    )
    response_body = json.loads(response["body"].read())
    response_text = response_body["content"][0]["text"]
    usage = response_body.get("usage") or {}
    logger.info(
        "BedrockBatch usage: %s model=%s",
        json.dumps({"input_tokens": usage.get("input_tokens"), "output_tokens": usage.get("output_tokens")}),
        model_id_used,
    )

    parsed = _extract_json_object(response_text)
    raw_results = parsed.get("results")
    if not isinstance(raw_results, list):
        raise ValueError("Bedrock JSON missing results array")

    # Index parsed rows by shelf_index and task_id
    by_index: Dict[int, Dict[str, Any]] = {}
    by_tid: Dict[str, Dict[str, Any]] = {}
    for row in raw_results:
        if not isinstance(row, dict):
            continue
        si = row.get("shelf_index")
        if isinstance(si, int):
            by_index[si] = row
        tid = row.get("task_id")
        if tid:
            by_tid[str(tid)] = row

    applied = []
    errors = []

    for job in shelf_jobs:
        row = by_index.get(job["shelf_index"]) or by_tid.get(job["task_id"])
        if not row:
            msg = f"No JSON result for shelf_index={job['shelf_index']} task_id={job['task_id']}"
            logger.warning(msg)
            errors.append({"task_id": job["task_id"], "error": msg})
            if job["image_id"]:
                try:
                    update_error_status(supabase, job["image_id"], msg)
                except Exception as e:
                    logger.error("update_error_status failed: %s", e)
            continue

        try:
            fr = row.get("FINAL_RESULT")
            if isinstance(fr, str):
                final = fr.strip().upper() in ("TRUE", "1", "YES")
            else:
                final = bool(fr)
            confidence = row.get("CONFIDENCE")
            if confidence is not None:
                confidence = str(confidence).strip()
            fc = row.get("FACINGS_COUNT")
            facings_count = None
            if fc is not None and str(fc).strip().upper() != "N/A":
                try:
                    facings_count = int(fc)
                except (TypeError, ValueError):
                    facings_count = None
            positions_raw = row.get("POSITIONS")
            position_str = None
            positions_list = None
            if isinstance(positions_raw, list):
                positions_list = []
                for el in positions_raw:
                    if isinstance(el, dict):
                        shelf = el.get("shelf")
                        pos = el.get("position")
                        fac = el.get("facings")
                        if shelf is not None and pos is not None and fac is not None:
                            positions_list.append(
                                {"shelf": int(shelf), "position": str(pos), "facings": int(fac)}
                            )
                position_str = _positions_to_string(positions_list) if positions_list else None
            elif isinstance(positions_raw, str):
                positions_list = _parse_positions_json(positions_raw)
                position_str = _positions_to_string(positions_list) if positions_list else None

            hpi = row.get("HAS_PI_LABEL")
            if hpi is None or str(hpi).upper() == "N/A":
                has_pi = None
            elif isinstance(hpi, bool):
                has_pi = hpi
            else:
                has_pi = str(hpi).strip().upper() in ("TRUE", "1", "YES")

            summary = row.get("ANALYSIS_SUMMARY") or ""
            msg_out = json.dumps(row, ensure_ascii=False) if not summary else str(summary)

            update_analysis_result(
                supabase,
                job["task_id"],
                final,
                msg_out,
                confidence=str(confidence) if confidence is not None else None,
                facings_count=facings_count,
                position=position_str,
                has_pi_label=has_pi,
                image_id=job["image_id"],
                image_path=job["image_path"],
            )
            applied.append(job["task_id"])
        except Exception as e:
            logger.exception("Apply result failed for task %s", job["task_id"])
            errors.append({"task_id": job["task_id"], "error": str(e)})
            if job["image_id"]:
                try:
                    update_error_status(supabase, job["image_id"], str(e))
                except Exception:
                    pass

    return {
        "product_id": product_id,
        "applied_count": len(applied),
        "applied_task_ids": applied,
        "errors": errors,
        "model_id": model_id_used,
    }


def lambda_handler(event: Dict[str, Any], context: Any) -> Dict[str, Any]:
    """Step Functions passes one group: { product_id, items: [...] }."""
    product_id = event.get("product_id") or (event.get("Payload") or {}).get("product_id")
    items = event.get("items") or (event.get("Payload") or {}).get("items")
    if not product_id or not items:
        raise ValueError("product_id and items required")
    if not isinstance(items, list) or len(items) < 1:
        raise ValueError("items must be a non-empty list")

    supabase = get_supabase_client()
    return run_batch(supabase, str(product_id), items)