"""
Lists product-grouped batches for Bedrock batch analysis: same product reference image once,
multiple shelf images per Lambda. Same task universe as list_retry_tasks (fixed today SAST, incomplete GenAI).

Uses only the task table: each item is { task_id }; grouping uses task.product_id; shelf photo URL is
task.image_url only (no product_images lookup for listing or for BedrockBatchAnalyze task_id path).
"""
import logging
import os
from datetime import date, datetime, timedelta
from typing import Any, Dict, List, Optional
from zoneinfo import ZoneInfo

import boto3
from supabase import Client, create_client

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

ssm = boto3.client("ssm")

TZ = ZoneInfo("Africa/Johannesburg")
DEFAULT_LIST_PAGE = 80


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


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


def _sast_day_bounds_utc(for_day: Optional[date] = None) -> tuple:
    z = TZ
    d = datetime.now(z).date() if for_day is None else for_day
    day_start_local = datetime.combine(d, datetime.min.time(), tzinfo=z)
    day_end_local = day_start_local + timedelta(days=1)
    utc = ZoneInfo("UTC")
    start_utc = day_start_local.astimezone(utc)
    end_utc = day_end_local.astimezone(utc)
    return (
        start_utc.isoformat().replace("+00:00", "Z"),
        end_utc.isoformat().replace("+00:00", "Z"),
        d.isoformat(),
    )


def _chunks(xs: List[Any], n: int):
    for i in range(0, len(xs), n):
        yield xs[i : i + n]


def _task_incomplete(row: Dict[str, Any]) -> bool:
    return row.get("processed_genai") is not True


def lambda_handler(event: Dict[str, Any], context: Any) -> Dict[str, Any]:
    """
    Returns one page of batch groups for Step Functions Map.
    Shape: { groups, has_more, next_offset, run_date, timezone, count, total_count }

    Each group: { product_id, items: [ { task_id }, ... ] } with len(items) <= BATCH_MAX_SHELF_IMAGES.
    Tasks without product_id or task.image_url are skipped (analyzer downloads shelf from task.image_url only).
    """
    offset = 0
    if isinstance(event, dict) and event.get("offset") is not None:
        try:
            offset = max(0, int(event["offset"]))
        except (TypeError, ValueError):
            offset = 0

    run_date_str = event.get("run_date") if isinstance(event, dict) else None
    if offset > 0 and not run_date_str:
        raise ValueError("run_date is required when offset > 0")

    page_size = int(os.environ.get("BATCH_LIST_PAGE_SIZE") or DEFAULT_LIST_PAGE)
    max_shelf = int(os.environ.get("BATCH_MAX_SHELF_IMAGES") or "8")
    max_shelf = max(1, min(max_shelf, 20))

    for_day: Optional[date] = None
    if run_date_str:
        try:
            y, m, d = run_date_str.split("-")
            for_day = date(int(y), int(m), int(d))
        except (ValueError, AttributeError):
            for_day = None

    start_iso, end_iso, run_date_label = _sast_day_bounds_utc(for_day)
    logger.info("SAST window for fixed_at: %s <= fixed_at < %s (%s)", start_iso, end_iso, run_date_label)

    supabase = get_supabase_client()

    tr = (
        supabase.table("task")
        .select("id,image_url,processed_genai,fixed_at,product_id")
        .not_.is_("fixed_at", "null")
        .gte("fixed_at", start_iso)
        .lt("fixed_at", end_iso)
        .execute()
    )
    tasks = tr.data or []
    tasks_incomplete = [t for t in tasks if _task_incomplete(t)]

    # Group by product_id using task rows only; shelf URL for analysis is task.image_url (see BedrockBatchAnalyze).
    by_pid: Dict[str, List[Dict[str, str]]] = {}
    skipped_no_product = 0
    skipped_no_url = 0
    for t in tasks_incomplete:
        pid = t.get("product_id")
        tid = str(t["id"])
        if not pid:
            skipped_no_product += 1
            continue
        if not (t.get("image_url") or "").strip():
            skipped_no_url += 1
            continue
        by_pid.setdefault(str(pid), []).append({"task_id": tid})

    if skipped_no_product:
        logger.info("list_bedrock_batch_groups: skipped %s tasks without product_id", skipped_no_product)
    if skipped_no_url:
        logger.info("list_bedrock_batch_groups: skipped %s tasks without task.image_url", skipped_no_url)

    # Stable order: sort product ids, flatten chunks of shelf jobs
    flat_chunks: List[Dict[str, Any]] = []
    for pid in sorted(by_pid.keys(), key=lambda x: x):
        items = by_pid[pid]
        for batch in _chunks(items, max_shelf):
            flat_chunks.append({"product_id": pid, "items": batch})

    total_count = len(flat_chunks)
    if total_count == 0:
        return {
            "groups": [],
            "has_more": False,
            "next_offset": 0,
            "run_date": run_date_label,
            "timezone": "Africa/Johannesburg",
            "count": 0,
            "total_count": 0,
        }

    end = offset + page_size
    page = flat_chunks[offset:end]
    next_offset = offset + len(page)
    has_more = next_offset < total_count

    out = {
        "groups": page,
        "has_more": has_more,
        "next_offset": next_offset,
        "run_date": run_date_label,
        "timezone": "Africa/Johannesburg",
        "count": len(page),
        "total_count": total_count,
    }
    logger.info(
        "ListBedrockBatchGroups: offset=%s page_size=%s total_chunks=%s page=%s has_more=%s",
        offset,
        page_size,
        total_count,
        len(page),
        has_more,
    )
    return out