#!/usr/bin/env python3
"""Batch vision inventory workflow with agent-assisted datasheet enrichment.

Option A workflow:
  1. Process every image in a folder with vision_inventory_mcp.py.
  2. Save one raw JSON file per image for auditability.
  3. Build parts_to_lookup.json for the agent/user to enrich from datasheets.
  4. If datasheet_cache.json exists, write a final enriched CSV.

The script intentionally does not browse the web. Fill datasheet_cache.json manually or with Pi
web-search assistance, then rerun this script with --skip-vision to regenerate the CSV.
"""

from __future__ import annotations

import argparse
import csv
import json
import re
import sys
from collections import Counter, defaultdict
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple

from PIL import Image, ImageOps

PROJECT_ROOT = Path(__file__).resolve().parents[1]
if str(PROJECT_ROOT) not in sys.path:
    sys.path.insert(0, str(PROJECT_ROOT))

import vision_inventory_mcp as vision  # noqa: E402

UNKNOWN_MARKINGS = {"", "unknown", "unreadable", "unclear", "none", "n/a"}
PART_PATTERN = re.compile(r"\b[A-Z]{1,4}[A-Z0-9]{2,}[A-Z0-9?\[\]-]*\b", re.IGNORECASE)
SEGMENT_CROP_MODEL_MAX_SIDE = 250


def safe_stem(path: Path) -> str:
    stem = re.sub(r"[^A-Za-z0-9_.-]+", "_", path.stem).strip("_")
    return stem or "image"


def load_json(path: Path, default: Any) -> Any:
    if not path.exists():
        return default
    return json.loads(path.read_text(encoding="utf-8"))


def write_json(path: Path, data: Any) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(json.dumps(data, indent=2, ensure_ascii=False), encoding="utf-8")


def is_supported_image(path: Path) -> bool:
    return path.is_file() and path.suffix.lower() in vision.SUPPORTED_EXTENSIONS


def iter_images(folder: Path, recursive: bool, limit: Optional[int]) -> List[Path]:
    iterator = folder.rglob("*") if recursive else folder.iterdir()
    images = sorted([p for p in iterator if is_supported_image(p)], key=lambda p: str(p).lower())
    return images[:limit] if limit is not None else images


def clamp(value: int, low: int, high: int) -> int:
    return max(low, min(high, value))


def boxes_overlap_fraction(a: Tuple[int, int, int, int], b: Tuple[int, int, int, int]) -> float:
    ax1, ay1, ax2, ay2 = a
    bx1, by1, bx2, by2 = b
    ix1, iy1 = max(ax1, bx1), max(ay1, by1)
    ix2, iy2 = min(ax2, bx2), min(ay2, by2)
    if ix2 <= ix1 or iy2 <= iy1:
        return 0.0
    intersection = (ix2 - ix1) * (iy2 - iy1)
    smaller = min((ax2 - ax1) * (ay2 - ay1), (bx2 - bx1) * (by2 - by1))
    return intersection / smaller if smaller else 0.0


def detect_dark_ic_like_boxes(
    image_path: Path,
    *,
    max_side: int = 1200,
    threshold: int = 95,
    min_area_ratio: float = 0.00045,
    max_area_ratio: float = 0.20,
    max_crops: int = 12,
) -> List[Tuple[int, int, int, int]]:
    """Find dark rectangular IC-like regions using Pillow-only connected components.

    This is intentionally conservative. It catches common black IC packages and
    returns original-image-coordinate boxes. If no reliable boxes are found, the
    caller should fall back to the original image.
    """
    image = ImageOps.exif_transpose(Image.open(image_path)).convert("RGB")
    original_w, original_h = image.size
    work = image.copy()
    scale = 1.0
    if max(original_w, original_h) > max_side:
        scale = max_side / float(max(original_w, original_h))
        resample = getattr(Image, "Resampling", Image).LANCZOS
        work.thumbnail((max_side, max_side), resample)

    gray = work.convert("L")
    w, h = gray.size
    pixels = list(gray.getdata())
    mask = [px <= threshold for px in pixels]
    visited = bytearray(w * h)
    min_area = max(80, int(w * h * min_area_ratio))
    max_area = max(min_area + 1, int(w * h * max_area_ratio))
    boxes: List[Tuple[int, int, int, int, int]] = []

    for start, is_dark in enumerate(mask):
        if not is_dark or visited[start]:
            continue
        stack = [start]
        visited[start] = 1
        area = 0
        min_x = w
        min_y = h
        max_x = 0
        max_y = 0
        while stack:
            idx = stack.pop()
            area += 1
            x = idx % w
            y = idx // w
            min_x = min(min_x, x)
            min_y = min(min_y, y)
            max_x = max(max_x, x)
            max_y = max(max_y, y)
            for nxt in (idx - 1, idx + 1, idx - w, idx + w):
                if nxt < 0 or nxt >= w * h or visited[nxt] or not mask[nxt]:
                    continue
                # Prevent left/right wrapping.
                if (nxt == idx - 1 and x == 0) or (nxt == idx + 1 and x == w - 1):
                    continue
                visited[nxt] = 1
                stack.append(nxt)

        box_w = max_x - min_x + 1
        box_h = max_y - min_y + 1
        if area < min_area or area > max_area:
            continue
        if box_w < 18 or box_h < 18:
            continue
        aspect = box_w / float(box_h)
        if aspect < 0.25 or aspect > 5.0:
            continue
        fill = area / float(box_w * box_h)
        if fill < 0.35:
            continue
        if box_w > w * 0.80 or box_h > h * 0.80:
            continue
        boxes.append((area, min_x, min_y, max_x + 1, max_y + 1))

    boxes.sort(reverse=True)
    selected: List[Tuple[int, int, int, int]] = []
    for _area, x1, y1, x2, y2 in boxes:
        candidate = (x1, y1, x2, y2)
        if any(boxes_overlap_fraction(candidate, existing) > 0.65 for existing in selected):
            continue
        selected.append(candidate)
        if len(selected) >= max_crops:
            break

    # Return top-to-bottom, left-to-right in original coordinates.
    scaled: List[Tuple[int, int, int, int]] = []
    for x1, y1, x2, y2 in selected:
        ox1 = int(x1 / scale)
        oy1 = int(y1 / scale)
        ox2 = int(x2 / scale)
        oy2 = int(y2 / scale)
        scaled.append((clamp(ox1, 0, original_w), clamp(oy1, 0, original_h), clamp(ox2, 0, original_w), clamp(oy2, 0, original_h)))
    return sorted(scaled, key=lambda box: (box[1], box[0]))


def write_ic_crops(image_path: Path, crops_dir: Path, args: argparse.Namespace) -> List[Dict[str, Any]]:
    boxes = detect_dark_ic_like_boxes(
        image_path,
        max_side=args.segment_detection_max_side,
        threshold=args.segment_dark_threshold,
        max_crops=args.segment_max_crops,
    )
    if not boxes:
        return []

    image = ImageOps.exif_transpose(Image.open(image_path)).convert("RGB")
    width, height = image.size
    crops_dir.mkdir(parents=True, exist_ok=True)
    crop_entries: List[Dict[str, Any]] = []
    for index, (x1, y1, x2, y2) in enumerate(boxes, start=1):
        pad = int(max(x2 - x1, y2 - y1) * args.segment_padding_ratio)
        crop_box = (
            clamp(x1 - pad, 0, width),
            clamp(y1 - pad, 0, height),
            clamp(x2 + pad, 0, width),
            clamp(y2 + pad, 0, height),
        )
        crop = image.crop(crop_box)
        if args.segment_mask_context:
            focus_pad = int(max(x2 - x1, y2 - y1) * args.segment_focus_padding_ratio)
            focus_box = (
                clamp(x1 - focus_pad, crop_box[0], crop_box[2]),
                clamp(y1 - focus_pad, crop_box[1], crop_box[3]),
                clamp(x2 + focus_pad, crop_box[0], crop_box[2]),
                clamp(y2 + focus_pad, crop_box[1], crop_box[3]),
            )
            rel_focus = (
                focus_box[0] - crop_box[0],
                focus_box[1] - crop_box[1],
                focus_box[2] - crop_box[0],
                focus_box[3] - crop_box[1],
            )
            masked = Image.new("RGB", crop.size, (255, 255, 255))
            masked.paste(crop.crop(rel_focus), rel_focus)
            crop = masked
        else:
            focus_box = (x1, y1, x2, y2)
        crop_path = crops_dir / f"{safe_stem(image_path)}__ic_{index:02d}.jpg"
        crop.save(crop_path, format="JPEG", quality=args.jpeg_quality, optimize=True)
        crop_entries.append({
            "crop_path": crop_path,
            "source_image": str(image_path),
            "crop_index": index,
            "bbox": crop_box,
            "focus_bbox": focus_box,
            "context_masked": args.segment_mask_context,
        })
    return crop_entries


def clean_marking(value: Any) -> str:
    text = str(value or "").strip()
    text = re.sub(r"\s+", " ", text)
    return text


def likely_base_part(marking: str) -> str:
    """Extract the most likely device part from a longer package marking line."""
    cleaned = marking.replace("[?]", "?").upper()
    matches = PART_PATTERN.findall(cleaned)
    if not matches:
        return cleaned.strip()

    # Prefer tokens that look like common IC identifiers and contain digits.
    candidates = [m.strip("-_") for m in matches if any(ch.isdigit() for ch in m)]
    if not candidates:
        candidates = matches

    # Date/lot codes are usually short and mostly numeric; prefer longer alphanumeric tokens.
    candidates.sort(key=lambda s: (len(re.sub(r"[^A-Z]", "", s)) > 0, len(s)), reverse=True)
    return candidates[0]


def candidate_from_item(item: Dict[str, Any]) -> str:
    for key in ("likely_part", "package_marking"):
        value = clean_marking(item.get(key))
        if value.lower() not in UNKNOWN_MARKINGS:
            return likely_base_part(value)
    return "unknown"


def possible_same_as_candidate(item: Dict[str, Any]) -> str:
    confidence = str(item.get("possible_same_as_confidence", "none")).strip().lower()
    if confidence != "high":
        return "unknown"
    value = clean_marking(item.get("possible_same_as_likely_part"))
    if value.lower() in UNKNOWN_MARKINGS:
        return "unknown"
    return likely_base_part(value)


def extract_part_evidence(image_name: str, result: Dict[str, Any]) -> List[Dict[str, Any]]:
    evidence: List[Dict[str, Any]] = []

    items = result.get("items", [])
    if isinstance(items, list):
        for item in items:
            if not isinstance(item, dict):
                continue
            if str(item.get("item_type", "")).strip().lower() != "ic":
                continue
            marking = clean_marking(item.get("package_marking"))
            if marking.lower() not in UNKNOWN_MARKINGS:
                evidence.append({
                    "image": image_name,
                    "source": "items",
                    "position_hint": item.get("position_hint", "unknown"),
                    "observed_marking": marking,
                    "candidate_part": candidate_from_item(item),
                    "marking_confidence": item.get("marking_confidence", "unknown"),
                    "needs_review": bool(item.get("needs_review", True)),
                })
                continue

            possible_candidate = possible_same_as_candidate(item)
            if possible_candidate.lower() not in UNKNOWN_MARKINGS:
                reason = clean_marking(item.get("same_as_reason"))
                evidence.append({
                    "image": image_name,
                    "source": "visual_same_as_hypothesis",
                    "position_hint": item.get("position_hint", "unknown"),
                    "observed_marking": marking or "unreadable",
                    "candidate_part": possible_candidate,
                    "marking_confidence": item.get("marking_confidence", "unreadable"),
                    "needs_review": True,
                    "inference_type": "visual_same_as_hypothesis",
                    "visual_group_id": item.get("visual_group_id", "unknown"),
                    "possible_same_as_confidence": item.get("possible_same_as_confidence", "high"),
                    "same_as_reason": reason or "Unreadable IC visually matches nearby readable repeated ICs.",
                })

    observations = result.get("ic_marking_observations", [])
    if isinstance(observations, list):
        for obs in observations:
            if not isinstance(obs, dict):
                continue
            marking = clean_marking(obs.get("package_marking"))
            if marking.lower() in UNKNOWN_MARKINGS:
                continue
            evidence.append({
                "image": image_name,
                "source": "ic_marking_observations",
                "position_hint": obs.get("position_hint", "unknown"),
                "observed_marking": marking,
                "candidate_part": likely_base_part(marking),
                "marking_confidence": obs.get("marking_confidence", "unknown"),
                "needs_review": str(obs.get("marking_confidence", "")).lower() in {"low", "unreadable"},
            })

    return evidence


def preflight_credentials() -> None:
    _account_id, _api_token, credential_error = vision.get_cloudflare_credentials()
    if credential_error:
        raise SystemExit(
            f"{credential_error.get('message', 'Missing Cloudflare credentials.')} "
            "Set CLOUDFLARE_ACCOUNT_ID and CLOUDFLARE_AUTH_TOKEN/CLOUDFLARE_API_TOKEN, "
            "or run /vision-inventory-setup in Pi."
        )


def segmented_crop_prompt(crop_name: str, source_name: str) -> str:
    return vision.build_default_user_prompt(crop_name) + f"""

Additional crop/orientation instructions:
- This image is an automatically cropped IC/package candidate from source image: {source_name}.
- The crop may be rotated or upside-down relative to the full board.
- Mentally check 0, 90, 180, and 270 degree orientations before transcribing package text.
- Prefer the orientation that makes letters/numbers readable as normal IC top markings.
- Only analyze the central/unmasked IC package in this crop.
- Ignore partial ICs, partial packages, or cut-off markings near crop edges; they are context/noise and must not become items.
- Return the marking for this cropped IC/package only when possible.
""".strip()


def process_one_image(image_path: Path, raw_path: Path, args: argparse.Namespace, metadata: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
    model_max_side = SEGMENT_CROP_MODEL_MAX_SIDE if metadata else args.max_side
    result = vision.process_image_impl(
        image_path=str(image_path),
        max_side=model_max_side,
        jpeg_quality=args.jpeg_quality,
        custom_prompt=(segmented_crop_prompt(image_path.name, Path(metadata["source_image"]).name) if metadata else None),
    )
    if metadata and isinstance(result, dict):
        result["source_image"] = Path(metadata["source_image"]).name
        result["crop_index"] = metadata["crop_index"]
        result["crop_bbox"] = metadata["bbox"]
        result["crop_focus_bbox"] = metadata.get("focus_bbox")
        result["crop_context_masked"] = metadata.get("context_masked", False)
        result["crop_model_max_side"] = model_max_side
    write_json(raw_path, result)
    return {"image_path": str(image_path), "raw_json": str(raw_path), "result": result}


def process_images(args: argparse.Namespace, raw_dir: Path) -> List[Dict[str, Any]]:
    image_folder = Path(args.image_folder).expanduser().resolve()
    if not image_folder.is_dir():
        raise SystemExit(f"Image folder does not exist or is not a directory: {image_folder}")

    images = iter_images(image_folder, args.recursive, args.limit)
    if not images:
        print(f"No supported image files found in {image_folder}.")
        return []

    preflight_credentials()

    results: List[Dict[str, Any]] = []
    crops_dir = raw_dir.parent / "crops"
    for image_path in images:
        if args.segment_ics:
            crop_entries = write_ic_crops(image_path, crops_dir, args)
            if crop_entries:
                print(f"Processing {image_path} as {len(crop_entries)} segmented IC crop(s)")
                for crop in crop_entries:
                    crop_path = crop["crop_path"]
                    raw_path = raw_dir / f"{safe_stem(crop_path)}.json"
                    results.append(process_one_image(crop_path, raw_path, args, metadata=crop))
                continue
            print(f"No IC-like segments found for {image_path}; processing full image")

        print(f"Processing {image_path}")
        raw_path = raw_dir / f"{safe_stem(image_path)}.json"
        results.append(process_one_image(image_path, raw_path, args))

    return results


def load_raw_results(raw_dir: Path) -> List[Dict[str, Any]]:
    results = []
    for raw_path in sorted(raw_dir.glob("*.json")):
        result = load_json(raw_path, {})
        results.append({"image_path": result.get("image", raw_path.stem), "raw_json": str(raw_path), "result": result})
    return results


def classify_error(result: Dict[str, Any]) -> str:
    message = str(result.get("message") or result.get("error") or "Unknown error")
    lowered = message.lower()
    if "credential" in lowered or "cloudflare_account_id" in lowered or "api token" in lowered:
        return "credential errors"
    if "cloudflare" in lowered or "workers ai" in lowered:
        return "cloudflare api errors"
    if "prepare image" in lowered or "unsupported image" in lowered or "image file" in lowered:
        return "image preprocessing/input errors"
    if result.get("parse_error"):
        return "model json parse errors"
    return "other errors"


def result_error_summary(results: List[Dict[str, Any]]) -> Counter[str]:
    summary: Counter[str] = Counter()
    for entry in results:
        result = entry.get("result", {})
        if isinstance(result, dict) and (result.get("error") or result.get("parse_error")):
            summary[classify_error(result)] += 1
    return summary


def build_parts_to_lookup(results: List[Dict[str, Any]], output_dir: Path) -> Dict[str, Any]:
    grouped: Dict[str, List[Dict[str, Any]]] = defaultdict(list)
    all_evidence: List[Dict[str, Any]] = []

    for entry in results:
        result = entry["result"]
        image_name = str(result.get("image") or Path(entry["image_path"]).name)
        evidence = extract_part_evidence(image_name, result)
        all_evidence.extend(evidence)
        for row in evidence:
            part = row["candidate_part"].upper()
            if part and part.lower() not in UNKNOWN_MARKINGS:
                grouped[part].append(row)

    parts = []
    for part, evidence_rows in sorted(grouped.items()):
        observed = sorted({row["observed_marking"] for row in evidence_rows})
        images = sorted({row["image"] for row in evidence_rows})
        parts.append({
            "part": part,
            "query": f"{part} datasheet",
            "status": "needs_datasheet_lookup",
            "observed_markings": observed,
            "images": images,
            "has_visual_same_as_hypothesis": any(row.get("inference_type") == "visual_same_as_hypothesis" for row in evidence_rows),
            "evidence": evidence_rows,
            "enrichment_template": {
                "normalized_part": part,
                "description": "",
                "datasheet_url": "",
                "manufacturer": "",
                "verified": False,
                "notes": ""
            }
        })

    warnings: List[str] = []
    if not parts:
        warnings.append(
            "No candidate IC parts were extracted. Inspect raw JSON files for unreadable markings, "
            "image quality issues, or model/API errors."
        )

    return {
        "output_dir": str(output_dir),
        "datasheet_cache_path": str(output_dir / "datasheet_cache.json"),
        "datasheet_cache_template_path": str(output_dir / "datasheet_cache.template.json"),
        "instructions": [
            "Use web search to find each part datasheet, preferably from the manufacturer.",
            "Fill datasheet_cache.json in this same output directory, using datasheet_cache.template.json as the shape.",
            "Keep descriptions short, e.g. '74ls (4 bit) adder low power schottky ttl 5v DIP'.",
            "If exact candidate search fails but official results strongly indicate a likely OCR correction, keep the original candidate as this cache key and set normalized_part to the official datasheet part number.",
            "Example: if SN74AS283N appears to be an OCR error for official SN74LS283N, use key SN74AS283N with normalized_part SN74LS283N and explain the correction in notes.",
            "Only mark verified=true for a correction when the official datasheet and visual/package context make the correction highly likely; otherwise set verified=false and explain in notes.",
            "If the visual marking is uncertain, set verified=false and explain in notes."
        ],
        "warnings": warnings,
        "parts": parts,
        "all_evidence": all_evidence,
    }


def build_cache_template(parts_to_lookup: Dict[str, Any]) -> Dict[str, Any]:
    cache = {}
    for part in parts_to_lookup.get("parts", []):
        key = part["part"]
        cache[key] = part["enrichment_template"]
    return cache


def lookup_enrichment(part: str, cache: Dict[str, Any]) -> Dict[str, Any]:
    if part in cache and isinstance(cache[part], dict):
        return cache[part]
    upper = part.upper()
    if upper in cache and isinstance(cache[upper], dict):
        return cache[upper]
    return {}


def positive_int(value: Any) -> Optional[int]:
    try:
        parsed = int(value)
    except Exception:
        return None
    return parsed if parsed > 0 else None


def estimate_amount_for_candidate(result: Dict[str, Any], candidate: str, evidence_count: int = 1) -> int:
    """Estimate physical IC quantity for one candidate in one image.

    Prefer explicit visible_quantity when the model provides it. Older results only
    have count_index, which may be either an ordinal index or a grouped count, so
    the fallback remains heuristic and should be reviewed for important BOMs.
    """
    items = result.get("items", [])
    if not isinstance(items, list):
        return max(1, evidence_count)

    matched = 0
    count_values: List[int] = []
    visible_quantities: List[int] = []
    for item in items:
        if not isinstance(item, dict):
            continue
        if str(item.get("item_type", "")).strip().lower() != "ic":
            continue
        item_candidate = candidate_from_item(item).upper()
        if item_candidate != candidate.upper():
            possible_candidate = possible_same_as_candidate(item).upper()
            if possible_candidate != candidate.upper():
                continue
        matched += 1
        visible_quantity = positive_int(item.get("visible_quantity"))
        if visible_quantity is not None:
            visible_quantities.append(visible_quantity)
        count_index = positive_int(item.get("count_index"))
        if count_index is not None:
            count_values.append(count_index)

    if visible_quantities:
        return max(1, sum(visible_quantities))

    return max([1, evidence_count, matched, *count_values])


def image_part_rows(results: List[Dict[str, Any]], cache: Dict[str, Any]) -> List[Dict[str, Any]]:
    rows: List[Dict[str, Any]] = []
    for entry in results:
        result = entry["result"]
        image_name = str(result.get("image") or Path(entry["image_path"]).name)
        evidence = extract_part_evidence(image_name, result)
        if not evidence:
            notes = "No IC marking extracted"
            if isinstance(result, dict) and result.get("error"):
                notes = f"Vision processing error: {result.get('message', 'unknown error')}"
            elif isinstance(result, dict) and result.get("warnings"):
                notes = "; ".join(str(w) for w in result.get("warnings", []) if str(w).strip()) or notes
            rows.append({
                "image": image_name,
                "candidate_part": "",
                "normalized_part": "",
                "amount": 0,
                "description": "",
                "datasheet_url": "",
                "manufacturer": "",
                "verified": False,
                "vision_confidence": "unreadable",
                "needs_review": True,
                "observed_markings": "",
                "observations": "",
                "raw_json": entry["raw_json"],
                "notes": notes,
            })
            continue

        # One image may contain multiple different IC candidates. Emit one
        # evidence row per candidate instead of forcing a single image-level part.
        evidence_by_candidate: Dict[str, List[Dict[str, Any]]] = defaultdict(list)
        for row in evidence:
            candidate = row["candidate_part"].upper()
            if candidate and candidate.lower() not in UNKNOWN_MARKINGS:
                evidence_by_candidate[candidate].append(row)

        for candidate, candidate_evidence in sorted(evidence_by_candidate.items()):
            enrichment = lookup_enrichment(candidate, cache)
            likely_part = str(enrichment.get("normalized_part") or candidate).strip().upper()
            amount = estimate_amount_for_candidate(result, candidate, evidence_count=len(candidate_evidence))
            # Keep observed_markings normalized to the main visible part number, not full date/lot/package text.
            observed_markings = [likely_part]
            observations = "; ".join(
                f"{row['position_hint']}: {row['observed_marking']} ({row['marking_confidence']})"
                + (
                    f" [visual same-as {row.get('possible_same_as_confidence')}: {row.get('same_as_reason')}]"
                    if row.get("inference_type") == "visual_same_as_hypothesis"
                    else ""
                )
                for row in candidate_evidence
            )
            confidence_values = [str(row.get("marking_confidence", "unknown")) for row in candidate_evidence]
            visual_same_as_notes = [
                str(row.get("same_as_reason", "")).strip()
                for row in candidate_evidence
                if row.get("inference_type") == "visual_same_as_hypothesis" and str(row.get("same_as_reason", "")).strip()
            ]
            needs_review = any(row.get("needs_review", True) for row in candidate_evidence) or not enrichment.get("verified", False)
            notes = enrichment.get("notes", "Missing datasheet enrichment")
            if visual_same_as_notes:
                notes = f"{notes}; Visual same-as hypothesis for unreadable IC(s): {' | '.join(sorted(set(visual_same_as_notes)))}"

            rows.append({
                "image": image_name,
                "candidate_part": candidate,
                "normalized_part": likely_part,
                "amount": amount,
                "description": enrichment.get("description", ""),
                "datasheet_url": enrichment.get("datasheet_url", ""),
                "manufacturer": enrichment.get("manufacturer", ""),
                "verified": bool(enrichment.get("verified", False)),
                "vision_confidence": "/".join(sorted(set(confidence_values))),
                "needs_review": needs_review,
                "observed_markings": " | ".join(observed_markings),
                "observations": observations,
                "raw_json": entry["raw_json"],
                "notes": notes,
            })
    return rows


def write_csv(path: Path, fieldnames: List[str], rows: List[Dict[str, Any]]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("w", newline="", encoding="utf-8") as csv_file:
        writer = csv.DictWriter(csv_file, fieldnames=fieldnames)
        writer.writeheader()
        writer.writerows(rows)


def write_final_csv(results: List[Dict[str, Any]], cache: Dict[str, Any], output_csv: Path) -> None:
    """Write the default deduplicated BOM CSV and a per-image evidence CSV."""
    evidence_rows = image_part_rows(results, cache)
    evidence_fieldnames = [
        "image",
        "candidate_part",
        "normalized_part",
        "amount",
        "description",
        "datasheet_url",
        "manufacturer",
        "verified",
        "vision_confidence",
        "needs_review",
        "observed_markings",
        "observations",
        "raw_json",
        "notes",
    ]

    grouped: Dict[str, List[Dict[str, Any]]] = defaultdict(list)
    no_part_rows: List[Dict[str, Any]] = []
    for row in evidence_rows:
        part = str(row.get("normalized_part") or row.get("candidate_part") or "").strip().upper()
        if not part:
            no_part_rows.append(row)
        else:
            grouped[part].append(row)

    bom_rows: List[Dict[str, Any]] = []
    for part, rows_for_part in sorted(grouped.items()):
        first = rows_for_part[0]
        images = sorted({str(row["image"]) for row in rows_for_part})
        observed_markings = sorted({marking for row in rows_for_part for marking in str(row["observed_markings"]).split(" | ") if marking})
        raw_json_files = sorted({str(row["raw_json"]) for row in rows_for_part})
        confidence_values = sorted({value for row in rows_for_part for value in str(row["vision_confidence"]).split("/") if value})
        notes = sorted({str(row.get("notes", "")) for row in rows_for_part if str(row.get("notes", "")).strip()})
        amount = sum(int(row.get("amount", 0) or 0) for row in rows_for_part)

        bom_rows.append({
            "normalized_part": part,
            "candidate_parts": ", ".join(sorted({str(row["candidate_part"]) for row in rows_for_part if row.get("candidate_part")})),
            "amount": amount,
            "sighting_count": len(rows_for_part),
            "description": first.get("description", ""),
            "datasheet_url": first.get("datasheet_url", ""),
            "manufacturer": first.get("manufacturer", ""),
            "verified": all(bool(row.get("verified", False)) for row in rows_for_part),
            "vision_confidence": "/".join(confidence_values),
            "needs_review": any(bool(row.get("needs_review", True)) for row in rows_for_part),
            "images": " | ".join(images),
            "observed_markings": " | ".join(observed_markings),
            "raw_json": " | ".join(raw_json_files),
            "notes": " | ".join(notes),
        })

    for row in no_part_rows:
        bom_rows.append({
            "normalized_part": "",
            "candidate_parts": "",
            "amount": 0,
            "sighting_count": 1,
            "description": "",
            "datasheet_url": "",
            "manufacturer": "",
            "verified": False,
            "vision_confidence": row.get("vision_confidence", "unreadable"),
            "needs_review": True,
            "images": row.get("image", ""),
            "observed_markings": "",
            "raw_json": row.get("raw_json", ""),
            "notes": row.get("notes", "No IC marking extracted"),
        })

    bom_fieldnames = [
        "normalized_part",
        "candidate_parts",
        "amount",
        "sighting_count",
        "description",
        "datasheet_url",
        "manufacturer",
        "verified",
        "vision_confidence",
        "needs_review",
        "images",
        "observed_markings",
        "raw_json",
        "notes",
    ]
    write_csv(output_csv, bom_fieldnames, bom_rows)
    write_csv(output_csv.with_name(f"{output_csv.stem}_evidence{output_csv.suffix}"), evidence_fieldnames, evidence_rows)


def validate_setup(args: argparse.Namespace, output_dir: Path) -> None:
    image_folder = Path(args.image_folder).expanduser().resolve()
    print("Setup validation:")
    print(f"- Python executable: {sys.executable}")
    print("- Required imports: ok")

    if image_folder.is_dir():
        images = iter_images(image_folder, args.recursive, args.limit)
        print(f"- Image folder: ok ({image_folder})")
        print(f"- Supported images found: {len(images)}")
        if args.segment_ics and images:
            sample_boxes = detect_dark_ic_like_boxes(
                images[0],
                max_side=args.segment_detection_max_side,
                threshold=args.segment_dark_threshold,
                max_crops=args.segment_max_crops,
            )
            print(f"- IC segmentation enabled: first image would produce {len(sample_boxes)} crop(s)")
        heic_images = [p for p in images if p.suffix.lower() in {".heic", ".heif"}]
        if heic_images:
            try:
                import pillow_heif  # noqa: F401
                print("- HEIC/HEIF support: ok")
            except Exception:
                print("- HEIC/HEIF support: missing pillow-heif; install it to process HEIC/HEIF images")
    else:
        print(f"- Image folder: missing or not a directory ({image_folder})")

    try:
        output_dir.mkdir(parents=True, exist_ok=True)
        probe = output_dir / ".vision_inventory_write_test"
        probe.write_text("ok", encoding="utf-8")
        probe.unlink(missing_ok=True)
        print(f"- Output directory writable: ok ({output_dir})")
    except Exception as exc:
        print(f"- Output directory writable: failed ({exc})")

    if args.skip_vision:
        raw_dir = output_dir / "raw"
        raw_count = len(list(raw_dir.glob("*.json"))) if raw_dir.exists() else 0
        print(f"- Raw JSON files for --skip-vision: {raw_count}")
    else:
        _account_id, _api_token, credential_error = vision.get_cloudflare_credentials()
        if credential_error:
            print(f"- Cloudflare credentials: missing ({credential_error.get('message')})")
        else:
            print("- Cloudflare credentials: present")


def print_workflow_summary(results: List[Dict[str, Any]], parts_to_lookup: Dict[str, Any]) -> None:
    error_summary = result_error_summary(results)
    evidence_count = len(parts_to_lookup.get("all_evidence", []))
    part_count = len(parts_to_lookup.get("parts", []))
    print("Workflow summary:")
    print(f"- Processed/raw result files: {len(results)}")
    print(f"- IC marking evidence rows: {evidence_count}")
    print(f"- Candidate parts extracted: {part_count}")
    if error_summary:
        print("- Processing errors:")
        for label, count in sorted(error_summary.items()):
            print(f"  - {label}: {count}")
    if not part_count:
        print("No candidate IC parts were extracted. Inspect output/raw/*.json for unreadable markings, image quality issues, or API errors.")


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Process electronics images and prepare datasheet-enriched CSV workflow.")
    parser.add_argument("image_folder", help="Folder containing electronics/PCB images")
    parser.add_argument("output_dir", help="Output directory for raw JSON, lookup files, and CSV")
    parser.add_argument("--csv", default="inventory.csv", help="CSV filename/path, relative to output_dir unless absolute")
    parser.add_argument("--recursive", action="store_true", help="Scan image_folder recursively")
    parser.add_argument("--limit", type=int, default=None, help="Maximum number of images to process")
    parser.add_argument("--skip-vision", action="store_true", help="Reuse existing output_dir/raw/*.json instead of calling vision AI")
    parser.add_argument("--validate-setup", action="store_true", help="Check dependencies, paths, credentials, and image discovery without processing images")
    parser.add_argument("--segment-ics", action="store_true", help="Detect dark IC-like packages, crop them individually, and process crops instead of the full image when crops are found")
    parser.add_argument("--segment-max-crops", type=int, default=12, help="Maximum IC/package crops to process per source image when --segment-ics is enabled")
    parser.add_argument("--segment-dark-threshold", type=int, default=95, help="Grayscale threshold for dark IC/package segmentation (0-255, lower is stricter)")
    parser.add_argument("--segment-detection-max-side", type=int, default=1200, help="Maximum image side used internally for segmentation detection")
    parser.add_argument("--segment-padding-ratio", type=float, default=0.22, help="Padding around each detected IC crop as a fraction of its largest side")
    parser.add_argument("--segment-focus-padding-ratio", type=float, default=0.04, help="Extra unmasked padding around the detected IC body when masking crop context")
    parser.add_argument("--no-segment-mask-context", dest="segment_mask_context", action="store_false", help="Do not whiten context outside the detected IC body in segmented crops")
    parser.set_defaults(segment_mask_context=True)
    parser.add_argument("--max-side", type=int, default=vision.DEFAULT_MAX_SIDE, help="Maximum resized image side; use 0 for full resolution (default)")
    parser.add_argument("--jpeg-quality", type=int, default=vision.DEFAULT_JPEG_QUALITY, help="JPEG quality for model input")
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    output_dir = Path(args.output_dir).expanduser().resolve()

    if args.validate_setup:
        validate_setup(args, output_dir)
        return

    raw_dir = output_dir / "raw"
    output_dir.mkdir(parents=True, exist_ok=True)
    raw_dir.mkdir(parents=True, exist_ok=True)

    if args.skip_vision:
        results = load_raw_results(raw_dir)
        if not results:
            raise SystemExit(f"No raw JSON files found in {raw_dir}")
    else:
        results = process_images(args, raw_dir)

    parts_to_lookup = build_parts_to_lookup(results, output_dir)
    parts_path = output_dir / "parts_to_lookup.json"
    template_path = output_dir / "datasheet_cache.template.json"
    cache_path = output_dir / "datasheet_cache.json"
    write_json(parts_path, parts_to_lookup)
    write_json(template_path, build_cache_template(parts_to_lookup))

    cache = load_json(cache_path, {})
    csv_path = Path(args.csv)
    if not csv_path.is_absolute():
        csv_path = output_dir / csv_path
    write_final_csv(results, cache, csv_path)

    print(f"Raw results: {raw_dir}")
    print(f"Parts to lookup: {parts_path}")
    print(f"Datasheet cache template: {template_path}")
    print(f"Datasheet cache used: {cache_path if cache_path.exists() else 'not found yet'}")
    print(f"CSV written: {csv_path}")
    print_workflow_summary(results, parts_to_lookup)
    if not cache_path.exists():
        print("Next step: copy datasheet_cache.template.json to datasheet_cache.json in the output directory, enrich it via web search, then rerun with --skip-vision.")

    errors = result_error_summary(results)
    if results and sum(errors.values()) == len(results):
        raise SystemExit("All processed images returned errors; see the summary above and inspect raw JSON files.")


if __name__ == "__main__":
    main()
