"""
Dry-run Bedrock analysis endpoint for model testing.
Never writes task/product_images rows and never uploads artifacts.

Bedrock calls use **only** `invoke_bedrock_test_model` here — not `app.analyze_image_with_bedrock`.
If the request omits `model_id`, the default is **Qwen** (`DEFAULT_QWEN_TEST_MODEL_ID` / `BEDROCK_TEST_MODEL_ID`).
"""
import base64
import copy
import json
import logging
import os
import re
from io import BytesIO
from typing import Dict, List, Optional

import boto3
from PIL import Image

import app as analysis_app

logger = logging.getLogger()
logger.setLevel(logging.INFO)
bedrock = boto3.client("bedrock-runtime")

# When the request omits model_id: use BEDROCK_TEST_MODEL_ID on the Lambda, else this Qwen id (Converse).
DEFAULT_QWEN_TEST_MODEL_ID = "qwen.qwen3-vl-235b-a22b"


def _converse_max_image_bytes() -> int:
    """Smaller than app.MAX_IMAGE_BYTES: Converse sends 2 images + prompt; Qwen gateway rejects oversized total bodies."""
    raw = (os.getenv("BEDROCK_TEST_CONVERSE_MAX_IMAGE_BYTES") or "").strip()
    if raw.isdigit():
        return max(50_000, int(raw))
    return 1_500_000  # ~1.5 MiB per image (tune via env if Qwen still errors)


def _encode_image_for_anthropic(image_data: bytes, image_name: str) -> str:
    """Validate + compress image then return base64 for Anthropic Messages payload."""
    if not isinstance(image_data, bytes):
        raise ValueError(f"{image_name} data must be bytes")
    img = Image.open(BytesIO(image_data))
    img.verify()
    image_data = analysis_app.resize_image_to_fit(image_data, analysis_app.MAX_IMAGE_BYTES)
    base64_data = base64.b64encode(image_data).decode("utf-8")
    padding_needed = len(base64_data) % 4
    if padding_needed:
        base64_data += "=" * (4 - padding_needed)
    return base64_data


def _build_prompt(
    task_name: str,
    product_name: str,
    product_category: str,
    product_brand: str,
    product_variant: str,
    product_size: str,
) -> str:
    return f"""You are a shelf-verification assistant. Two images: (1) reference product, (2) current shelf.

TASK: {task_name}
TARGET: Brand={product_brand}, Product={product_name}, Category={product_category}, Variant={product_variant}, Size={product_size}.

WHAT TO LOOK FOR (from the reference image only):
1. Packaging form: Is the reference a SACHET/PACKET (soft pouch, rectangular), a CAN (metal cylinder), or a BOTTLE? Name it.
2. Main colors: e.g. red, green, white.
3. Distinctive logo/shape: e.g. green oval, star, specific pattern.

CRITICAL — How to search the shelf:
- The target can be on ANY shelf (row 1 = top, then 2, 3, …). Check every row.
- If the reference is a SACHET or PACKET, the shelf often also has CANS. Do not assume the shelf has only cans. Look specifically for soft pouches/sachets; they may be smaller, at the left or right of a row, or in a block. Scan the full width of each shelf for the reference's form.
- Match by form first (sachet vs can vs bottle), then by main colors and logo shape. You do NOT need to read small text.
- If you see the same form + same main colors + same logo shape as the reference on any shelf, say TRUE. Say FALSE only if you find no such item after checking every row.

RULE: Same packaging form and key visuals = TRUE. Rotated or partially visible is OK. When the main cues align, prefer TRUE.

OUTPUT DEFINITIONS:
- CONFIDENCE: 0-100% that the target product is on the shelf; N/A if unsure.
- FACINGS_COUNT: Count only the FRONT ROW (the row facing the customer). One facing = one unit in that front row. Do NOT count units stacked behind the front row. If uncertain (e.g. you think 6-8), report the LOWEST number (e.g. 6). Prefer under-count over over-count.
- POSITION: "Shelf N, left|center|right", Shelf 1 = top. Where the target product block is. N/A if not found.
- HAS_PI_LABEL: TRUE if a price tag, shelf-edge label, or price sticker is visible below or adjacent to the target product; FALSE if you can see that area but no label; N/A only if the label area is not visible (cropped or obscured). When the shelf and product are visible, prefer TRUE or FALSE over N/A.

End your reply with exactly these five lines (exact labels):
FINAL_RESULT: TRUE|FALSE
CONFIDENCE: <0-100% or N/A>
FACINGS_COUNT: <integer or N/A>
POSITION: Shelf <n>, <left|center|right> or N/A
HAS_PI_LABEL: TRUE|FALSE or N/A"""


def _parse_response_text(response_text: str) -> Dict:
    text = response_text.strip()
    lines = [ln.strip() for ln in text.split("\n") if ln.strip()]
    structured_block = "\n".join(lines[-10:]) if len(lines) >= 10 else text
    result = False
    confidence = None
    facings_count = None
    position = None
    has_pi_label = None
    for raw_line in lines[-10:]:
        line = raw_line.strip()
        if line.upper().startswith("FINAL_RESULT:"):
            val = line.split(":", 1)[1].strip().upper()
            result = val == "TRUE"
        elif line.upper().startswith("CONFIDENCE:"):
            raw = line.split(":", 1)[1].strip().upper()
            confidence = None if raw == "N/A" else raw
        elif line.upper().startswith("FACINGS_COUNT:"):
            raw = line.split(":", 1)[1].strip().upper()
            if raw == "N/A":
                facings_count = None
            else:
                try:
                    facings_count = int(raw)
                except Exception:
                    facings_count = None
        elif line.upper().startswith("POSITION:"):
            raw = line.split(":", 1)[1].strip().upper()
            position = None if raw == "N/A" else line.split(":", 1)[1].strip()
        elif line.upper().startswith("HAS_PI_LABEL:"):
            val = line.split(":", 1)[1].strip().upper()
            has_pi_label = None if val == "N/A" else (val == "TRUE")
    if "FINAL_RESULT:" not in structured_block.upper():
        if re.search(r"\bTRUE\b", structured_block, re.IGNORECASE):
            result = True
        elif re.search(r"\bFALSE\b", structured_block, re.IGNORECASE):
            result = False
    return {
        "result": result,
        "message": text,
        "confidence": confidence,
        "facings_count": facings_count,
        "position": position,
        "has_pi_label": has_pi_label,
    }


def _run_anthropic_override(
    model_id: str,
    image_bytes: bytes,
    product_image_bytes: Optional[bytes],
    prompt: str,
) -> Dict:
    base64_image = _encode_image_for_anthropic(image_bytes, "Main image")
    base64_product_image = (
        _encode_image_for_anthropic(product_image_bytes, "Product image")
        if product_image_bytes
        else None
    )
    content = [{"type": "text", "text": "I will show you two images to compare:"}]
    if base64_product_image:
        content.extend(
            [
                {
                    "type": "text",
                    "text": "First, here is the reference image showing how the product should look:",
                },
                {
                    "type": "image",
                    "source": {
                        "type": "base64",
                        "media_type": "image/jpeg",
                        "data": base64_product_image,
                    },
                },
            ]
        )
    content.extend(
        [
            {"type": "text", "text": "Now, here is the current image we need to evaluate:"},
            {
                "type": "image",
                "source": {
                    "type": "base64",
                    "media_type": "image/jpeg",
                    "data": base64_image,
                },
            },
            {"type": "text", "text": prompt},
        ]
    )
    request_body = {
        "anthropic_version": "bedrock-2023-05-31",
        "max_tokens": 1000,
        "temperature": 0.2,
        "messages": [{"role": "user", "content": content}],
    }
    log_body = copy.deepcopy(request_body)
    for message in log_body["messages"]:
        for item in message["content"]:
            if item.get("type") == "image":
                item["source"]["data"] = "<image_data_omitted>"
    logger.info("Test invoke request body: %s", json.dumps(log_body)[:4000])
    response = bedrock.invoke_model(
        modelId=model_id,
        body=json.dumps(request_body).encode("utf-8"),
        contentType="application/json",
        accept="application/json",
    )
    response_body = json.loads(response["body"].read())
    response_text = response_body["content"][0]["text"]
    return _parse_response_text(response_text)


def _run_qwen_converse(
    model_id: str,
    image_bytes: bytes,
    product_image_bytes: Optional[bytes],
    prompt: str,
) -> Dict:
    cap = _converse_max_image_bytes()
    shelf_jpeg = analysis_app.resize_image_to_fit(image_bytes, cap)
    product_jpeg = (
        analysis_app.resize_image_to_fit(product_image_bytes, cap)
        if product_image_bytes
        else None
    )
    content: List[Dict] = [{"text": "I will show you two images to compare."}]
    if product_jpeg:
        content.append({"text": "First, this is the reference product image."})
        content.append({"image": {"format": "jpeg", "source": {"bytes": product_jpeg}}})
    content.append({"text": "Now, this is the current shelf image."})
    content.append({"image": {"format": "jpeg", "source": {"bytes": shelf_jpeg}}})
    content.append({"text": prompt})
    response = bedrock.converse(
        modelId=model_id,
        messages=[{"role": "user", "content": content}],
        inferenceConfig={"maxTokens": 4096, "temperature": 0},
    )
    out_message = (response.get("output") or {}).get("message") or {}
    texts = []
    for part in out_message.get("content") or []:
        if isinstance(part, dict) and part.get("text"):
            texts.append(str(part["text"]))
    response_text = "\n".join(texts).strip()
    logger.info("Qwen Converse response (trimmed): %s", response_text[:2000])
    return _parse_response_text(response_text)


def invoke_bedrock_test_model(
    model_id: str,
    image_bytes: bytes,
    product_image_bytes: Optional[bytes],
    prompt: str,
) -> Dict:
    """
    Single test-only Bedrock entrypoint. Does not call production `analyze_image_with_bedrock`.
    Routes by model id prefix: qwen.* → Converse; otherwise InvokeModel (Anthropic Messages).
    """
    mid = (model_id or "").strip()
    if not mid:
        raise ValueError("model_id is required for invoke_bedrock_test_model")
    path = "test_converse" if mid.lower().startswith("qwen.") else "test_invoke_model"
    logger.info("Bedrock test invoke path=%s model_id=%s", path, mid)
    if path == "test_converse":
        parsed = _run_qwen_converse(mid, image_bytes, product_image_bytes, prompt)
    else:
        parsed = _run_anthropic_override(mid, image_bytes, product_image_bytes, prompt)
    parsed["_test_invoke_path"] = path
    return parsed


def resolve_test_model_id(body: Dict) -> str:
    """Body model_id/modelId wins; else BEDROCK_TEST_MODEL_ID env; else Qwen default (test route is for Qwen by default)."""
    explicit = (body.get("model_id") or body.get("modelId") or "").strip()
    if explicit:
        return explicit
    env_default = (os.getenv("BEDROCK_TEST_MODEL_ID") or "").strip()
    if env_default:
        return env_default
    return DEFAULT_QWEN_TEST_MODEL_ID


def lambda_handler(event, context):
    """POST body: same as /analyze-image plus optional model_id for dry-run testing."""
    try:
        if event.get("httpMethod") == "OPTIONS":
            return analysis_app.create_api_response(200, {"message": "CORS preflight request successful"})

        try:
            if isinstance(event.get("body"), str):
                body = json.loads(event["body"])
            else:
                body = event.get("body") or {}
        except json.JSONDecodeError:
            return analysis_app.create_api_response(400, {"error": "Invalid JSON in request body"})

        # API Gateway sends a string "body"; the Lambda console and direct invokes use root keys only.
        if isinstance(event, dict):
            body = dict(body)
            for key in ("image_id", "task_id", "model_id", "modelId"):
                if key not in body and event.get(key) not in (None, ""):
                    body[key] = event[key]

        image_id = body.get("image_id")
        task_id_param = body.get("task_id")
        if not image_id and not task_id_param:
            return analysis_app.create_api_response(400, {"error": "image_id or task_id is required in request body"})
        if image_id and task_id_param:
            task_id_param = None

        supabase = analysis_app.get_supabase_client()
        if image_id:
            data = analysis_app.get_image_data(supabase, image_id)
            image_path = (data["image"] or {}).get("image_path") or (data.get("task") or {}).get("image_url")
        else:
            data = analysis_app.get_task_data(supabase, task_id_param)
            image_path = (data.get("task") or {}).get("image_url")

        if not image_path:
            return analysis_app.create_api_response(400, {"error": "No shelf image available: set product_images.image_path or task.image_url"})

        image_bytes = analysis_app.download_image(supabase, image_path)
        task_name = (data.get("task") or {}).get("name") or "N/A"
        task_id = data.get("task_id") or "N/A"
        product_name = (data.get("product") or {}).get("product") or "N/A"
        product_category = (data.get("product") or {}).get("category") or "N/A"
        product_brand = (data.get("product") or {}).get("brand") or "N/A"
        product_variant = (data.get("product") or {}).get("variant") or "N/A"
        product_size = (data.get("product") or {}).get("size") or "N/A"
        product_image_url = (data.get("product") or {}).get("image_path") or "N/A"

        product_image_bytes = None
        if product_image_url != "N/A":
            try:
                product_image_bytes = analysis_app.download_image(supabase, product_image_url)
            except Exception as e:
                logger.warning("Test analyze: failed to download product image: %s", e)

        try:
            model_id = resolve_test_model_id(body)
        except ValueError as e:
            return analysis_app.create_api_response(400, {"error": str(e)})

        prompt = _build_prompt(
            task_name,
            product_name,
            product_category,
            product_brand,
            product_variant,
            product_size,
        )
        analysis = invoke_bedrock_test_model(model_id, image_bytes, product_image_bytes, prompt)

        invoke_path = analysis.get("_test_invoke_path")
        result = {k: v for k, v in analysis.items() if not (isinstance(k, str) and k.startswith("_"))}
        payload = {
            "message": "Dry-run analysis completed (no database writes; test-only Bedrock path)",
            "dry_run": True,
            "model_id": model_id,
            "test_invoke_path": invoke_path,
            "task_id": str(task_id) if task_id else None,
            "image_id": str(image_id) if image_id else None,
            "result": result,
        }
        return analysis_app.create_api_response(200, payload)
    except Exception as e:
        logger.exception("test_app lambda_handler failed")
        return analysis_app.create_api_response(500, {"error": str(e)})