import os
import json
import base64
import random
import time
from datetime import datetime
from supabase import create_client, Client
import boto3
from botocore.config import Config
from botocore.exceptions import ClientError
from typing import Dict, Optional
import logging
from PIL import Image, ImageOps
from io import BytesIO
import requests
from urllib.parse import urlparse, unquote
import re
import copy

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

# Initialize AWS clients
ssm = boto3.client('ssm')
# Primary profile: exactly one SDK attempt (no adaptive stack); see invoke_bedrock_primary_then_fallback.
bedrock_primary_once = boto3.client(
    'bedrock-runtime',
    config=Config(
        retries={'max_attempts': 1},
        read_timeout=900,
        connect_timeout=30,
    ),
)
# Fallback path: adaptive SDK retries + optional outer backoff in invoke_bedrock_with_backoff
bedrock = boto3.client(
    'bedrock-runtime',
    config=Config(
        retries={'max_attempts': 8, 'mode': 'adaptive'},
        read_timeout=900,
        connect_timeout=30,
    ),
)
# S3 client created on demand when saving artifacts (optional)


def invoke_bedrock_with_backoff(*, max_outer: int = 10, client=None, **invoke_kwargs):
    """invoke_model with optional outer exponential backoff (throttling / transient errors).

    Primary (4.6) should call with max_outer=1 and client=bedrock_primary_once so we only
    hit the primary profile once per request; fallback (4.5) can use defaults.
    """
    invoke_client = client if client is not None else bedrock
    base = 2.0
    ceiling = 60.0
    last_exc = None
    for attempt in range(1, max_outer + 1):
        try:
            return invoke_client.invoke_model(**invoke_kwargs)
        except ClientError as e:
            last_exc = e
            code = (e.response or {}).get('Error', {}).get('Code', '') or ''
            if code not in (
                'ThrottlingException',
                'TooManyRequestsException',
                'ServiceUnavailable',
                'InternalServerException',
            ):
                raise
            if attempt >= max_outer:
                raise
            delay = min(base * (2 ** (attempt - 1)), ceiling)
            jitter = random.uniform(0, delay * 0.2)
            sleep_s = delay + jitter
            logger.warning(
                'Bedrock %s (outer attempt %s/%s); sleeping %.1fs before retry',
                code, attempt, max_outer, sleep_s,
            )
            time.sleep(sleep_s)
    raise last_exc  # pragma: no cover

# Constants
BEDROCK_MODEL_ID = "anthropic.claude-sonnet-4-6"
# Bedrock image limit 5 MB; base64 adds ~4/3, so keep raw JPEG under this
MAX_IMAGE_BYTES = int(5 * 1024 * 1024 * 3 / 4) - 100_000  # ~3.8 MB safe margin
# When resizing, don't go below this height so all shelf rows (e.g. 8) stay visible to the model
MIN_RESIZE_HEIGHT = 720

def get_parameter(name: str) -> str:
    """Get a parameter from AWS SSM Parameter Store"""
    try:
        response = ssm.get_parameter(
            Name=name,
            WithDecryption=True
        )
        return response['Parameter']['Value']
    except Exception as e:
        logger.error(f"Error getting parameter {name}: {str(e)}")
        raise

def get_supabase_client() -> Client:
    """Initialize and return Supabase client"""
    try:
        url: str = get_parameter('/supabase/url')
        key: str = get_parameter('/supabase/anon')
        supabase: Client = create_client(url, key)
        return supabase
    except Exception as e:
        logger.error(f"Error initializing Supabase client: {str(e)}")
        raise

def try_get_parameter(name: str) -> Optional[str]:
    """Best-effort fetch of an SSM parameter. Returns None if not found."""
    try:
        response = ssm.get_parameter(Name=name, WithDecryption=True)
        return response['Parameter']['Value']
    except ssm.exceptions.ParameterNotFound:
        return None
    except Exception as e:
        logger.error(f"Error getting optional parameter {name}: {str(e)}")
        return None

def get_bedrock_invocation_target() -> Dict[str, str]:
    """Return kwargs for Bedrock invocation using an inference profile.

    For Claude Sonnet 4 Vision, serverless on-demand is not available.
    We require an inference profile ARN to be provided via the
    INFERENCE_PROFILE_ARN environment variable (set by SAM template).
    """
    profile_arn = os.getenv('INFERENCE_PROFILE_ARN')
    if not profile_arn:
        raise RuntimeError(
            "Missing environment variable INFERENCE_PROFILE_ARN. "
            "Create an inference profile for anthropic.claude-sonnet-4-6 "
            "and set its ARN on the Lambda environment."
        )
    return {"modelId": profile_arn}

def get_image_data(supabase: Client, image_id: str) -> Dict:
    """Get image, product, and task details from Supabase.

    From product_images we use: product_id, task_id, image_path (shelf image),
    processed_genai, last_processed_image_path; and we update processed_genai,
    genai_result, genai_message, last_processed_image_path, genai_error.
    Shelf image URL: prefer product_images.image_path, fallback to task.image_url.
    """
    try:
        # Get image record
        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

        # Get product details only if a valid product_id exists
        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 not product_response.data:
                raise Exception(f"Product with ID {product_id} not found")
            product = product_response.data

        # Get task details only if a valid task_id exists (needed for task name and image_url fallback)
        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 not task_response.data:
                raise Exception(f"Task with ID {task_id} not found")
            task = task_response.data

        return {
            "image": image_record,
            "product": product,
            "task": task,
            "task_id": task_id
        }
    except Exception as e:
        logger.error(f"Error getting data from Supabase: {str(e)}")
        raise

def get_task_data(supabase: Client, task_id: str) -> Dict:
    """Get task and product by task_id when no product_images row exists. Shelf image from task.image_url."""
    try:
        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}
    except Exception as e:
        logger.error(f"Error getting task data from Supabase: {str(e)}")
        raise

def _supabase_public_url_bucket_key(url: str) -> Optional[tuple]:
    """Parse Supabase Storage public URL -> (bucket, object_key). Works with custom domains."""
    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:
    """Download image from either Supabase storage path or direct URL"""
    try:
        # Check if the image_path is a URL
        if image_path.startswith('http'):
            image_data = None
            # Prefer Storage API for public object URLs (avoids 400s from CDN/WAF on Lambda egress)
            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)
                    logger.info(
                        f"Downloaded image via Supabase storage API bucket={bucket_name} key={object_key[:80]}..."
                    )
                except Exception as e:
                    logger.warning(
                        f"Supabase storage download failed ({bucket_name}/{object_key[:60]}...), trying HTTP: {e}"
                    )
            if image_data is None:
                # Fallback: raw HTTP (browser-like headers; some CDNs still return 400 from Lambda)
                headers = {
                    "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36",
                    "Accept": "image/avif,image/webp,image/apng,image/svg+xml,image/*,*/*;q=0.8",
                }
                response = requests.get(image_path, headers=headers, timeout=30)
                if not response.ok:
                    snippet = (response.text or "")[:400].replace("\n", " ")
                    logger.error(
                        f"HTTP {response.status_code} downloading URL (first 400 chars of body): {snippet!r}"
                    )
                response.raise_for_status()
                image_data = response.content
        else:
            # Parse path for Supabase storage
            # Remove leading slash if present
            image_path = image_path.lstrip('/')
            
            # If path starts with bucket name, remove it
            if image_path.startswith('public/'):
                image_path = image_path[7:]  # Remove 'public/'
                
            bucket_name = "public"
            image_data = supabase.storage.from_(bucket_name).download(image_path)

        # Validate image data
        try:
            # Try to open the image with PIL to validate it
            img = Image.open(BytesIO(image_data))
            img.verify()  # Verify it's a valid image
            logger.info(f"Successfully validated image from {image_path}, size: {len(image_data)} bytes, format: {img.format}")
            return image_data
        except Exception as e:
            logger.error(f"Invalid image data from {image_path}: {str(e)}")
            raise ValueError(f"Invalid image data received from {image_path}")
            
    except requests.RequestException as e:
        logger.error(f"Error downloading image from URL {image_path}: {str(e)}")
        raise
    except Exception as e:
        logger.error(f"Error downloading image {image_path}: {str(e)}")
        raise

def resize_image_to_fit(image_bytes: bytes, max_bytes: int = MAX_IMAGE_BYTES) -> bytes:
    """Keep under max_bytes (Bedrock 5 MB base64 limit). Prefer COMPRESSION (lower JPEG quality, no resize);
    only resize as last resort so full resolution and all shelf rows stay visible.
    Applies EXIF orientation so images are always sent upright to Bedrock."""
    img = Image.open(BytesIO(image_bytes))
    img = ImageOps.exif_transpose(img)  # Apply EXIF orientation; no-op if none or already upright
    img = img.convert('RGB')
    w, h = img.size
    # 1) Try compression only (no resize): full size, lower quality until it fits. Good quality often at 70-80.
    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:
            logger.info(f"Image size {len(data)} bytes (max {max_bytes}), compressed only q={quality}, size {w}x{h}")
            return data
    # 2) Still over at quality 50: resize as last resort (min height 720 so all shelves visible)
    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:
                logger.info(f"Image size {len(data)} bytes (max {max_bytes}), resized {nw}x{nh} q={quality}")
                return data
    # Fallback: return last attempt
    logger.info(f"Image {len(data)} bytes; returning anyway")
    return data

def _parse_positions_json(raw: str) -> Optional[list]:
    """Parse POSITIONS line into list of {shelf, position, facings}. Returns None if invalid."""
    raw = raw.strip()
    if not raw or raw.upper() == 'N/A':
        return []
    # Extract JSON array: from first [ to matching ]
    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:
    """Convert positions list to backward-compatible string for genai_position."""
    return "; ".join(f"Shelf {p['shelf']}, {p['position']} ({p['facings']} facings)" for p in positions)

def analyze_image_with_bedrock(image_bytes: bytes, task_name: str, product_name: str, product_category: str, product_brand: str, product_variant: str, product_size: str, product_image_url: str, product_image_bytes: Optional[bytes] = None) -> Dict:
    """Analyze image using Bedrock's Claude Sonnet 4 Vision"""
    try:
        # Convert and validate images to base64
        try:
            # Function to encode and validate image
            def encode_image(image_data, image_name):
                # Convert image to bytes if it's not already
                if not isinstance(image_data, bytes):
                    logger.error(f"{image_name} data is not in bytes format")
                    raise ValueError(f"{image_name} data must be bytes")

                # Ensure image is valid
                try:
                    img = Image.open(BytesIO(image_data))
                    img.verify()
                    logger.info(f"{image_name} verified as valid {img.format} image")
                except Exception as e:
                    logger.error(f"{image_name} is not a valid image: {str(e)}")
                    raise

                # Resize/compress to stay under Bedrock 5 MB (base64) limit
                try:
                    image_data = resize_image_to_fit(image_data, MAX_IMAGE_BYTES)
                except Exception as e:
                    logger.error(f"Error resizing {image_name}: {str(e)}")
                    raise

                # Encode to base64; return (base64, processed_bytes) for S3 (save what model saw)
                try:
                    base64_data = base64.b64encode(image_data).decode('utf-8')
                    padding_needed = len(base64_data) % 4
                    if padding_needed:
                        base64_data += '=' * (4 - padding_needed)
                    logger.info(f"{image_name} encoded to base64, size: {len(image_data)} bytes, base64 length: {len(base64_data)}")
                    return (base64_data, image_data)
                except Exception as e:
                    logger.error(f"Error encoding {image_name} to base64: {str(e)}")
                    raise

            # Convert main image
            base64_image, shelf_image_bytes = encode_image(image_bytes, "Main image")

            # Convert product image if available
            base64_product_image = None
            product_image_bytes_sent = None
            if product_image_bytes:
                base64_product_image, product_image_bytes_sent = encode_image(product_image_bytes, "Product image")

        except Exception as e:
            logger.error(f"Error during image processing: {str(e)}")
            raise
        
        # Build size-specific instructions when size is provided (e.g. 100G, 200G, 430g)
        has_size = product_size and str(product_size).strip().upper() not in ('N/A', '')
        size_block = f"""
SIZE MUST MATCH — Target size is {product_size}:
- When Size is specified above, you MUST verify the size on the shelf product (e.g. "100G", "200G", "430g" printed on the package).
- Same brand, product, variant but DIFFERENT size is NOT the target: e.g. target 100G box—do NOT count 200G or 400G boxes. Target 430g jar—do NOT count 210g jars. Only count items where you can identify the size marking matches the target.
- If size is not legible on the shelf item, do NOT count it—when in doubt, exclude.""" if has_size else ""

        # Prepare the prompt: explicit "what to look for" for cans, jars, sachets; reference image is primary when provided
        prompt = 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}.

REFERENCE IMAGE IS PRIMARY: When a reference product image is shown, it is the SOURCE OF TRUTH. Identify the target by visual match to the reference—packaging form, lid/label colors, logo, product type, and SIZE (when visible). If TARGET text conflicts with the reference image, trust the reference image.

EXACT PRODUCT ONLY — Do not count the wrong product:
- The target is the EXACT product shown in the reference and named above (e.g. "{product_name}").
- Same brand but different product is NOT the target. Same product in a different size is NOT the target when Size is specified.
- Other brands (even similar colors) are NOT the target. Similar-looking names (e.g. PAKCO vs PARDO) are different brands—compare exact logo and label.
{size_block}

WHAT TO LOOK FOR (from the reference image only):
1. Packaging form: SACHET/PACKET (soft pouch), CAN (metal cylinder), JAR (glass jar with lid), or BOTTLE/BOX?
2. Main colors: Exact layout (e.g. green band, orange lid, yellow label).
3. Distinctive logo/shape: Match the reference exactly. Only items matching THIS logo and THIS product type are the target.
4. Size (when specified): Look for the size marking (e.g. 100G, 200G, 430g) on the packaging—only count items where the size matches the target.

SHELF COUNT — Get the total number of shelves right:
- Look at the FULL image from top to BOTTOM. Many store images have 7 or 8 shelf rows; do NOT stop at 6. Count every horizontal row that holds products. State clearly: "There are exactly N shelves" (single integer).
- Shelf 1 = topmost row, Shelf 2 = next down, … last row = Shelf N. You MUST scan shelf 1, 2, 3, … through Shelf N. If you only counted 6, look again at the lower part of the image for more rows.

METHOD — Scan every shelf; only list locations where you're sure it's the exact product:
- For EACH shelf from 1 to N: Scan full width. Use position: left | center-left | center | center-right | right. Look only for the EXACT product (same product name/type as reference). Same brand but different product or different size do NOT count. The target may appear on ZERO, ONE, or a FEW shelves; do not assume similar-looking blocks are the target.
- Only add a location to POSITIONS when confident it is the exact target—same product, same variant, and same size (when Size is specified). If you cannot read the size on the shelf item, do NOT count it.

FACING COUNT — Count exactly; avoid over-counting:
- At each location, count only the FRONT ROW (units facing the viewer). Mentally number each unit: "1, 2, 3" (or 4, 5, …). Report that exact number. A common error is reporting 4 when there are 3—re-count before reporting. Do not count units stacked behind the front row.

RULE: Exact product + same form + same key colors + same logo + same size (when specified) = target. Other products or other sizes are NOT the target.

OUTPUT DEFINITIONS:
- CONFIDENCE: 0-100% that the target product is on the shelf; N/A if unsure.
- FACINGS_COUNT: Total front-row facings across ALL locations. Count only the FRONT ROW at each location (do not count units behind). At each block, number the visible units (1, 2, 3…) and report that exact count—e.g. if you see 3 cans in the front row, report 3, not 4. Sum the counts for all locations. Single integer. N/A only if not found.
- POSITIONS: A JSON array of every location where the target appears. Each element: {{"shelf": N, "position": "left"|"center-left"|"center"|"center-right"|"right", "facings": M}}. Position is from the VIEWER's perspective: left = left edge, center-left = left of center, center = middle, center-right = right of center, right = right edge. Example: [{{"shelf": 5, "position": "center-right", "facings": 3}}]. If not found use [].
- HAS_PI_LABEL: TRUE if a price/shelf label is visible below or adjacent to the target at any location; FALSE if visible but no label; N/A if label area not visible.

DOUBLE-CHECK (required before output):
1. Shelves: Did I look at the BOTTOM of the image? Many images have 7 or 8 rows—if I said 6, re-count. State exactly N shelves and scan 1 through N.
2. Product: For each location I listed, am I sure it is the EXACT product (same logo, label, packaging as reference)? If not sure, remove that location.
3. Size: When Size is specified (e.g. 100G), did I verify the size marking on each item? Same product in 200G or 400G is NOT the target—exclude if size does not match or is not legible.
4. Facings: At each location, did I count the front row only and report the exact number (e.g. 3 not 4)? Re-count: 1, 2, 3 = 3 facings.
5. Position: Which of left | center-left | center | center-right | right best describes the block's location on that shelf from the viewer's perspective? Do not invert.
6. Is FACINGS_COUNT the sum of the "facings" values in POSITIONS?

End your reply with exactly these six lines (exact labels):
FINAL_RESULT: TRUE|FALSE
CONFIDENCE: <0-100% or N/A>
FACINGS_COUNT: <integer or N/A>
POSITIONS: <JSON array of {{"shelf", "position" ("left"|"center-left"|"center"|"center-right"|"right"), "facings"}} or []>
HAS_PI_LABEL: TRUE|FALSE or N/A"""

        # Start with the text prompt to set context
        content = [
            {
                "type": "text",
                "text": "I will show you two images to compare:"
            }
        ]

        # Add the reference image first if available
        if base64_product_image:
            content.append({
                "type": "text",
                "text": "First, here is the reference image showing how the product should look:"
            })
            content.append({
                "type": "image",
                "source": {
                    "type": "base64",
                    "media_type": "image/jpeg",
                    "data": base64_product_image
                }
            })

        # Add the current image
        content.append({
            "type": "text",
            "text": "Now, here is the current image we need to evaluate:"
        })
        content.append({
                "type": "image",
                "source": {
                    "type": "base64",
                    "media_type": "image/jpeg",
                    "data": base64_image
                }
        })
        
        # Add the main analysis prompt
        content.append({
            "type": "text",
            "text": prompt
        })

        # Log the prompt for debugging
        logger.info("Generated prompt:")
        logger.info(prompt)

        # Log content structure
        logger.info("Content structure:")
        for item in content:
            if item.get("type") == "text":
                logger.info(f"Text content: {item.get('text')}")
            elif item.get("type") == "image":
                logger.info(f"Image content: {item.get('source', {}).get('type')} image included")

        # Prepare request body for Claude Sonnet 4
        request_body = {
            "anthropic_version": "bedrock-2023-05-31",
            "max_tokens": 1000,
            "temperature": 0,
            "messages": [
                {
                    "role": "user",
                    "content": content
                }
            ]
        }
        
        # Log the full request body (excluding actual image data for brevity)
        # Use a deep copy so we do NOT mutate the actual request payload
        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("Request body:")
        logger.info(json.dumps(log_body, indent=2))

        # Invoke Bedrock (prefer inference profile if configured); backoff on throttling under parallel Step Functions
        target_kwargs = get_bedrock_invocation_target()
        response = invoke_bedrock_with_backoff(
            **target_kwargs,
            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']
        # Log token usage for cost tracking (queryable in CloudWatch Logs Insights)
        usage = response_body.get('usage') or {}
        inp = usage.get('input_tokens', 0)
        out = usage.get('output_tokens', 0)
        logger.info("BedrockUsage: " + json.dumps({"input_tokens": inp, "output_tokens": out}))

        # Log the response
        logger.info("Bedrock response:")
        logger.info(response_text)

        # Parse response: extract structured block (last 15 lines) for FINAL_RESULT etc.; also search full text for POSITIONS JSON
        text = response_text.strip()
        lines = [ln.strip() for ln in text.split('\n') if ln.strip()]
        structured_block = '\n'.join(lines[-15:]) if len(lines) >= 15 else text
        result = False
        confidence = None
        facings_count = None
        position = None
        positions = None  # list of {"shelf", "position", "facings"}
        has_pi_label = None
        for raw_line in lines[-15:]:
            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('POSITIONS:'):
                raw = line.split(':', 1)[1].strip()
                positions = _parse_positions_json(raw)
                if positions is not None:
                    position = _positions_to_string(positions) if positions else None
            elif line.upper().startswith('POSITION:'):
                if position is None:
                    raw = line.split(':', 1)[1].strip()
                    position = None if raw.upper() == '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 POSITIONS was multi-line, try to extract from full text
        if positions is None and 'POSITIONS:' in text:
            idx = text.upper().find('POSITIONS:')
            rest = text[idx + len('POSITIONS:'):].strip()
            positions = _parse_positions_json(rest)
            if positions is not None:
                position = _positions_to_string(positions) if positions else None
        # Heuristic fallback only in structured block so HAS_PI_LABEL: TRUE does not set result
        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
        message = text
        out = {
            "result": result,
            "message": message.strip(),
            "confidence": confidence,
            "facings_count": facings_count,
            "position": position,
            "has_pi_label": has_pi_label
        }
        if positions is not None:
            out["positions"] = positions
        # Optional: for S3 artifacts — exactly what was sent to the model (do not persist to DB)
        out["_prompt_text"] = prompt
        out["_shelf_image_bytes"] = shelf_image_bytes
        out["_product_image_bytes"] = product_image_bytes_sent
        return out
    except Exception as e:
        logger.error(f"Error analyzing image with Bedrock: {str(e)}")
        raise

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:
    """Update the analysis result in the database. If image_id is set, update product_images and task; else only task."""
    try:
        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()
        # Update the task table. Always set all genai columns so a FALSE run clears old facings/position/label/confidence.
        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()
    except Exception as e:
        logger.error(f"Error updating analysis result: {str(e)}")
        raise

def update_error_status(supabase: Client, image_id: str, error_msg: str) -> None:
    """Update error status in the database"""
    try:
        supabase.table("product_images").update({
            "processed_genai": False,
            "genai_error": error_msg
        }).eq("id", image_id).execute()
    except Exception as e:
        logger.error(f"Error updating error status: {str(e)}")
        raise

# --- Optional S3 artifacts (images sent to model + prompt + response). Enable via env GENAI_ARTIFACTS_S3_BUCKET. ---
def save_artifacts_to_s3(bucket: str, key_prefix: str, shelf_image_bytes: bytes, prompt_text: str, response_text: str, product_image_bytes: Optional[bytes] = None) -> None:
    """Upload shelf and product images (exactly as sent to Bedrock), prompt, and response to S3. No-op if bucket empty. Failures are logged only."""
    if not bucket or not bucket.strip():
        return
    try:
        s3 = boto3.client('s3')
        s3.put_object(Bucket=bucket, Key=f"{key_prefix}shelf_image.jpg", Body=shelf_image_bytes, ContentType="image/jpeg")
        if product_image_bytes:
            s3.put_object(Bucket=bucket, Key=f"{key_prefix}product_image.jpg", Body=product_image_bytes, ContentType="image/jpeg")
        s3.put_object(Bucket=bucket, Key=f"{key_prefix}prompt.txt", Body=prompt_text.encode("utf-8"), ContentType="text/plain; charset=utf-8")
        s3.put_object(Bucket=bucket, Key=f"{key_prefix}response.txt", Body=response_text.encode("utf-8"), ContentType="text/plain; charset=utf-8")
        artifacts = ["shelf_image.jpg"] + (["product_image.jpg"] if product_image_bytes else []) + ["prompt.txt", "response.txt"]
        logger.info(f"Saved artifacts to s3://{bucket}/{key_prefix} ({', '.join(artifacts)})")
    except Exception as e:
        logger.warning(f"Failed to save artifacts to S3 (bucket={bucket}): {str(e)}", exc_info=True)

def create_api_response(status_code: int, body: Dict) -> Dict:
    """Create a standardized API Gateway response"""
    return {
        'statusCode': status_code,
        'headers': {
            'Content-Type': 'application/json',
            'Access-Control-Allow-Origin': '*',
            'Access-Control-Allow-Methods': 'POST,OPTIONS',
            'Access-Control-Allow-Headers': 'Content-Type,X-Amz-Date,Authorization,X-Api-Key,X-Amz-Security-Token'
        },
        'body': json.dumps(body)
    }

def lambda_handler(event, context):
    """Main Lambda handler function (API Gateway or direct Step Functions / Lambda invoke)."""
    image_id = None  # set below; used in except for update_error_status
    try:
        # Handle OPTIONS request for CORS
        if event.get('httpMethod') == 'OPTIONS':
            return create_api_response(200, {'message': 'CORS preflight request successful'})

        # Direct invoke (e.g. Step Functions Map): payload is { image_id } or { task_id } at root
        if not event.get('httpMethod') and (event.get('image_id') or event.get('task_id')):
            body = {k: event[k] for k in ('image_id', 'task_id', 'from_trigger', 'product_id') if k in event}
            print("Direct invoke body:", body)
        else:
            # Parse API Gateway body
            print("Incoming event:", event)  # Log full event
            try:
                if isinstance(event.get('body'), str):
                    body = json.loads(event['body'])
                else:
                    body = event.get('body', {})
                print("Parsed body:", body)  # Log parsed body
            except json.JSONDecodeError:
                print("Failed to parse JSON body:", event.get('body'))  # Log failed parsing
                return create_api_response(400, {'error': 'Invalid JSON in request body'})

        # Validate input: require image_id (product_images) or task_id (task table only)
        image_id = body.get('image_id')
        task_id_param = body.get('task_id')
        if not image_id and not task_id_param:
            return create_api_response(400, {'error': 'image_id or task_id is required in request body'})
        if image_id and task_id_param:
            # Prefer image_id when both provided
            task_id_param = None

        # Check if request is from trigger
        from_trigger = body.get('from_trigger', False)

        # Initialize Supabase client
        supabase = get_supabase_client()

        # Get data: from product_images (image_id) or from task only (task_id)
        if image_id:
            data = get_image_data(supabase, image_id)
            image_record = data['image']
            image_path = image_record.get('image_path') or (data.get('task') or {}).get('image_url')
            processed_genai = image_record.get('processed_genai', False)
            last_processed_image_path = image_record.get('last_processed_image_path')
        else:
            data = get_task_data(supabase, task_id_param)
            image_path = (data.get('task') or {}).get('image_url')
            processed_genai = False
            last_processed_image_path = None

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

        # If from trigger and we have product_images, check skip
        if from_trigger and image_id:
            if processed_genai and last_processed_image_path == image_path:
                logger.info(f"Skipping processing for image_id {image_id}: already processed with same image_path")
                return create_api_response(200, {
                    'message': 'Processing skipped: already processed with same image_path',
                    'skipped': True
                })
            logger.info(f"Processing from trigger: processed_genai={processed_genai}, image_path changed={last_processed_image_path != image_path}")
        
        # Download and validate main image
        logger.info(f"Attempting to download main image from: {image_path}")
        image_bytes = download_image(supabase, image_path)
        logger.info(f"Successfully downloaded main image, size: {len(image_bytes)} bytes")
        
        # Analyze image
        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'
        # Get product_image from product data (maps to product.image_path)
        product_image_url = (data.get('product') or {}).get('image_path') or 'N/A'
        
        # Download the product image if available
        product_image_bytes = None
        if product_image_url != 'N/A':
            try:
                logger.info(f"Attempting to download product image from: {product_image_url}")
                product_image_bytes = download_image(supabase, product_image_url)
                logger.info("Successfully downloaded product image")
            except Exception as e:
                logger.warning(f"Failed to download product image from {product_image_url}: {str(e)}")
        else:
            logger.warning("No product image URL provided (N/A)")
        
        analysis = analyze_image_with_bedrock(
            image_bytes,
            task_name,
            product_name,
            product_category,
            product_brand,
            product_variant,
            product_size,
            product_image_url,
            product_image_bytes
        )
        
        # Update results (product_images only when image_id present)
        update_analysis_result(
            supabase,
            task_id,
            analysis['result'],
            analysis['message'],
            analysis.get('confidence'),
            analysis.get('facings_count'),
            analysis.get('position'),
            analysis.get('has_pi_label'),
            image_id=image_id,
            image_path=image_path
        )

        # Optional: save image + prompt + response to S3 (set GENAI_ARTIFACTS_S3_BUCKET to enable)
        bucket =  "datafy-bedrock-artifacts"
        if bucket:
            run_id = image_id or task_id or "unknown"
            ts = datetime.utcnow().strftime("%Y%m%dT%H%M%SZ")
            key_prefix = f"artifacts/{product_name}-{product_variant}-{product_size}-{ts}/"
            logger.info(f"Saving artifacts to s3://{bucket}/{key_prefix}")
            save_artifacts_to_s3(
                bucket,
                key_prefix,
                analysis.get("_shelf_image_bytes") or image_bytes,
                analysis.get("_prompt_text") or "",
                analysis.get("message") or "",
                analysis.get("_product_image_bytes"),
            )
        else:
            logger.info("GENAI_ARTIFACTS_S3_BUCKET not set; skipping S3 artifacts")

        # Exclude internal keys from API response
        result = {k: v for k, v in analysis.items() if not (isinstance(k, str) and k.startswith("_"))}
        return create_api_response(200, {
            'message': 'Image analysis completed successfully',
            'result': result
        })

    except Exception as e:
        error_message = str(e)
        logger.error(f"Error in lambda_handler: {error_message}")
        
        # Update error status only when we have product_images (image_id)
        if image_id is not None and 'supabase' in locals():
            try:
                update_error_status(supabase, image_id, error_message)
            except Exception as update_error:
                logger.error(f"Failed to update error status: {str(update_error)}")
        
        return create_api_response(500, {'error': error_message})