"""
ArtiAgent Orchestrator — DefectDiffu Edition

Adapted from the FLUX-based pipeline to use DefectDiffu (ECCV 2024).

KEY ARCHITECTURAL CHANGES vs. FLUX version:
  1. DefectDiffu generates NEW images from noise (text-to-image), rather than
     editing an existing image via inversion-injection.
  2. Three disentangled text prompts drive generation:
       c_p = background/product consistency prompt
       c_d = defect consistency prompt
       c_f = fusion prompt
  3. Double-free strategy controls defect strength (w_d) and product fidelity (w_p).
  4. Masks are generated automatically from defect-block cross-attention maps.
  5. Patch-based artifact mappings (16x16 FLUX patches) are REMOVED entirely.
  6. The input "clean image" is used for PLANNING and VERIFICATION only.

Usage:
    python artiagent_orchestrator.py \\
        --product-desc "VCSEL laser diode with glass lens cap" \\
        --image ./clean_chip.png \\
        --output-dir ./defect_output \\
        --defectdiffu-ckpt ./defectdiffu_ckpt.pt \\
        --vae-path ./sd-vae-ft-mse \\
        --device cuda
"""

import os
import sys
import json
import argparse
import uuid
import traceback
from pathlib import Path
from typing import Dict, List, Optional, Tuple
from datetime import datetime

import numpy as np
import torch
from PIL import Image

SCRIPT_DIR = Path(__file__).parent
if str(SCRIPT_DIR) not in sys.path:
    sys.path.insert(0, str(SCRIPT_DIR))

from pipeline.local_vlm_client import LocalVLMClient
from pipeline.prompts import (
    plan_defects_for_product,
    artifact_description,
    MoneyManager
)
from pipeline.gsam_detector import GSAMDetector
from pipeline.defectdiffu_generator import DefectDiffuGenerator, DefectDiffuConfig  # NEW
from pipeline.instance_processor import InstanceProcessor
from pipeline.defect_rag import get_rag
from pipeline.domain_router import get_router

import cv2

def blend_defect_onto_real_image(
    real_image: np.ndarray,      # Clean real VCSEL factory photo [H, W, 3]
    defect_image: Image.Image,   # Pure DefectDiffu generated output (PIL)
    defect_mask: np.ndarray,     # Binary mask from DefectDiffu [512, 512]
    target_bbox: List[int],      # [x1, y1, x2, y2] from VLM perception
    max_defect_ratio: Optional[float] = None, # None or >= 1.0 means 100% full ROI coverage
    mask_shape: str = "free"               # "circle", "square", "rectangle", or "free"
) -> Tuple[np.ndarray, np.ndarray]:
    """
    Injects a DefectDiffu defect patch onto a real clean factory image at 
    target_bbox using Poisson Seamless Cloning (Method 2).
    """
    defect_np = np.array(defect_image)
    mask_uint8 = (defect_mask.astype(np.uint8) * 255) if defect_mask.dtype == bool else defect_mask.astype(np.uint8)

    # 1. Crop tight patch around the generated defect mask
    ys, xs = np.where(mask_uint8 > 0)
    if len(ys) == 0 or len(xs) == 0:
        return real_image.copy(), np.zeros(real_image.shape[:2], dtype=np.uint8)

    y1_d, y2_d = ys.min(), ys.max()
    x1_d, x2_d = xs.min(), xs.max()

    defect_patch = defect_np[y1_d:y2_d + 1, x1_d:x2_d + 1]
    mask_patch = mask_uint8[y1_d:y2_d + 1, x1_d:x2_d + 1]

    # 2. Extract VLM target ROI dimensions
    x1_t, y1_t, x2_t, y2_t = target_bbox
    target_w = max(1, x2_t - x1_t)
    target_h = max(1, y2_t - y1_t)

    # 3. Dynamic Sizing Logic
    # max_defect_ratio = 1.0
    if max_defect_ratio < 0.2:
        max_defect_ratio = 0.2

    if max_defect_ratio is None or max_defect_ratio >= 1.0:
        # OPTION A: 100% Full target ROI fit (for smudges, contamination, large scratches)
        final_w = target_w
        final_h = target_h
    else:
        # OPTION B: Scaled down relative to target ROI (for bubbles, pinholes, particles)
        scale_factor = np.sqrt(max_defect_ratio)
        scaled_w = int(target_w * scale_factor)
        scaled_h = int(target_h * scale_factor)
        
        patch_h, patch_w = defect_patch.shape[:2]
        aspect_ratio = patch_w / max(1, patch_h)
        
        if aspect_ratio > 1:
            final_w = max(15, scaled_w)
            final_h = max(15, int(final_w / aspect_ratio))
        else:
            final_h = max(15, scaled_h)
            final_w = max(15, int(final_h * aspect_ratio))

    # 4. Resize patch and mask to fit VLM bounding box
    defect_patch_resized = cv2.resize(defect_patch, (final_w, final_h), interpolation=cv2.INTER_AREA)
    mask_patch_resized = cv2.resize(mask_patch, (final_w, final_h), interpolation=cv2.INTER_NEAREST)

    # =========================================================================
    # 🎨 VLM DYNAMIC MASK SHAPE GENERATION
    # =========================================================================
    # shape_type = mask_shape.lower().strip()
    shape_type = "free"

    if shape_type == "circle":
        # Draw a perfect filled circle
        geom_mask = np.zeros((final_h, final_w), dtype=np.uint8)
        center = (final_w // 2, final_h // 2)
        radius = max(1, min(final_w, final_h) // 2 - 1)
        cv2.circle(geom_mask, center, radius, 255, thickness=-1)
        mask_patch_resized = geom_mask

    elif shape_type == "square":
        # Draw a centered square
        geom_mask = np.zeros((final_h, final_w), dtype=np.uint8)
        side = max(1, min(final_w, final_h) - 2)
        top_left_x = (final_w - side) // 2
        top_left_y = (final_h - side) // 2
        cv2.rectangle(
            geom_mask, 
            (top_left_x, top_left_y), 
            (top_left_x + side, top_left_y + side), 
            255, 
            thickness=-1
        )
        mask_patch_resized = geom_mask

    elif shape_type == "rectangle":
        # Fill the entire patch as a solid rectangle
        mask_patch_resized = np.full((final_h, final_w), 255, dtype=np.uint8)

    elif shape_type == "free" or shape_type == "irregular":
        # Keep original organic AI-generated mask from DefectDiffu
        pass
    # =========================================================================

    # 5. Calculate center point for cv2.seamlessClone
    center_x = x1_t + target_w // 2
    center_y = y1_t + target_h // 2
    center = (center_x, center_y)

    # 6. Convert RGB -> BGR for OpenCV Poisson Blending
    real_bgr = cv2.cvtColor(real_image, cv2.COLOR_RGB2BGR)
    patch_bgr = cv2.cvtColor(defect_patch_resized, cv2.COLOR_RGB2BGR)

    # Inside blend_defect_onto_real_image:
    patch_mean = np.mean(defect_patch_resized)

    if patch_mean < 30:
        # Use NORMAL_CLONE for dark/subtle features to prevent Poisson smoothing from erasing them
        clone_mode = cv2.NORMAL_CLONE
    else:
        clone_mode = cv2.MIXED_CLONE # Preserves underlying substrate structure while injecting defect texture

    blended_bgr = cv2.seamlessClone(
        patch_bgr,
        real_bgr,
        mask_patch_resized,
        center,
        clone_mode 
    )
    blended_rgb = cv2.cvtColor(blended_bgr, cv2.COLOR_BGR2RGB)

    # 7. Map binary mask to full real image resolution for segmentation ground-truth
    full_mask = np.zeros(real_image.shape[:2], dtype=np.uint8)
    top_left_x = max(0, center_x - final_w // 2)
    top_left_y = max(0, center_y - final_h // 2)
    
    h_end = min(real_image.shape[0], top_left_y + final_h)
    w_end = min(real_image.shape[1], top_left_x + final_w)
    
    mask_crop_h = h_end - top_left_y
    mask_crop_w = w_end - top_left_x

    if mask_crop_h > 0 and mask_crop_w > 0:
        full_mask[top_left_y:h_end, top_left_x:w_end] = (
            mask_patch_resized[:mask_crop_h, :mask_crop_w] > 128
        ).astype(np.uint8)

    return blended_rgb, full_mask

def create_visual_prompt_image(full_image: np.ndarray, bbox: list) -> np.ndarray:
    """Draws a bright neon bounding box on the full image around the target ROI."""
    viz_img = full_image.copy()
    x1, y1, x2, y2 = bbox
    
    # Draw a 2px bright neon green or red rectangle
    cv2.rectangle(viz_img, (x1, y1), (x2, y2), (0, 255, 0), thickness=2)
    return viz_img

class ArtiAgentOrchestrator:
    """Agentic orchestrator for directed defect generation with DefectDiffu."""

    def __init__(
        self,
        device='cuda',
        output_dir='./defect_output',
        vlm_model='gemma3:12b',
        defectdiffu_ckpt: str = "",
        vae_path: str = "",
        image_size: int = 512,
        num_steps: int = 50
    ):
        self.device = device
        self.output_dir = Path(output_dir)
        self.output_dir.mkdir(parents=True, exist_ok=True)

        self.vlm_client = LocalVLMClient(model=vlm_model)
        self.money_manager = MoneyManager(model="gpt-4o")

        self.gsam_detector = None
        self.defectdiffu_generator = None

        # DefectDiffu config
        self.defectdiffu_ckpt = defectdiffu_ckpt
        self.vae_path = vae_path
        self.image_size = image_size
        self.num_steps = num_steps

        self.rag = get_rag()
        self.router = get_router()

    # ------------------------------------------------------------------
    # Lazy initializers
    # ------------------------------------------------------------------

    def _init_gsam(self):
        if self.gsam_detector is None or getattr(self.gsam_detector, 'sam_predictor', None) is None:
            print("[Agent] Initializing GSAM detector...")
            self.gsam_detector = GSAMDetector(
                device=self.device,
                openai_client=self.vlm_client
            )

    def _init_defectdiffu(self):
        if self.defectdiffu_generator is None:
            print("[Agent] Initializing DefectDiffu generator...")
            config = DefectDiffuConfig(
                ckpt_path=self.defectdiffu_ckpt,
                vae_path=self.vae_path,
                image_size=self.image_size,
                num_steps=self.num_steps,
                device=self.device,
                seed=42
            )
            self.defectdiffu_generator = DefectDiffuGenerator(config)

    # ------------------------------------------------------------------
    # Step 1: Planning
    # ------------------------------------------------------------------

    def plan(self, product_description: str, image: np.ndarray,
             defect_type: Optional[str] = None, num_defects: int = 3):
        """Agent plans defects based on product knowledge."""
        print(f"\\n{'='*60}")
        print("[Agent] Step 1: Planning defects from product description...")
        print(f"Product: {product_description}")
        if defect_type:
            print(f"[Agent] Target user defect type requested: {defect_type}")

        plan = plan_defects_for_product(
            self.vlm_client,
            product_description,
            image,
            money_manager=self.money_manager,
            target_defect_type=defect_type,
            num_defects=num_defects
        )

        if plan is None:
            raise RuntimeError("Defect planning failed")

        print(f"[Agent] Product type: {plan.product_type}")
        print(f"[Agent] Analysis: {plan.analysis}")
        print(f"[Agent] Proposed {len(plan.possible_defects)} defects:")
        for i, d in enumerate(plan.possible_defects, 1):
            print(f"  {i}. [{d.defect_type.upper()}] {d.description}")
            print(f"     Target: {d.target_entity} / {d.target_subentity or '(whole)'}"
                  f" | Location: {d.location_hint}")
            if hasattr(d, 'c_p'):
                print(f"     c_p: {d.c_p}")
                print(f"     c_d: {d.c_d}")
                print(f"     w_d: {d.w_d}")

        return plan

    # ------------------------------------------------------------------
    # Step 2: Perception (optional — for verification bbox only)
    # ------------------------------------------------------------------

    def perceive(self, image: np.ndarray, defect_plan):
        """
        Directed perception — detect target entity for verification cropping.
        With DefectDiffu this is OPTIONAL; the generator does not need patches.
        We keep it to obtain a bbox for the VLM verification step.
        """
        print(f"\\n{'='*60}")
        print("[Agent] Step 2: Directed perception (verification bbox)...")

        self._init_gsam()

        entity = defect_plan.target_entity
        synonym_map = {
            "metal can package": ["TO-can", "metal can", "can body", "package body", "metal ring"],
            "lens cap": ["glass lens cap", "lens", "glass dome", "optical window"],
            "electrode bars": ["vertical bars", "electrodes", "metal lines"],
        }
        search_terms = [entity] + synonym_map.get(entity.lower(), [])

        predictions = []
        for term in search_terms:
            preds, _, viz = self.gsam_detector.detect_parts(
                image=image,
                entities=[term],
                subentities=[defect_plan.target_subentity] if defect_plan.target_subentity else [],
                entity_subentity_mapping={},
                min_area_ratio=0.005,
                max_area_ratio=0.5,
                openai_client=self.vlm_client
            )
            if len(preds) > 0:
                predictions = preds
                if term != entity:
                    print(f"[Agent] Fallback: detected '{term}' instead of '{entity}'")
                break

        if not predictions:
            print(f"[Agent] Warning: No detections for {entity}; verification will use full image")
            h, w = image.shape[:2]
            best_pred = {
                'bbox': [0, 0, w, h],
                'pred_mask': torch.ones((h, w), dtype=torch.bool)
            }
        if predictions:
            best_pred = max(predictions, key=lambda p: p.get('area_ratio', 0))
            bbox = best_pred['bbox']
            h, w = image.shape[:2]
        
            # Check if detected bbox center is too close to border (< 10% margin)
            cx = (bbox[0] + bbox[2]) / 2
            cy = (bbox[1] + bbox[3]) / 2
            if cx < 0.1 * w or cx > 0.9 * w or cy < 0.1 * h or cy > 0.9 * h:
                print(f"[Agent] Warning: Bbox {bbox} on edge. Falling back to image center.")
                best_pred['bbox'] = [int(0.25 * w), int(0.25 * h), int(0.75 * w), int(0.75 * h)]
            
        return best_pred

    # ------------------------------------------------------------------
    # Step 3: Prepare DefectDiffu generation conditions
    # ------------------------------------------------------------------

    def prepare_generation_conditions(
        self,
        defect_plan,
        product_description: str,
        rag_k: int = 3
    ) -> Dict:
        """
        Build the three DefectDiffu text prompts + double-free scales.
        Retrieves RAG examples to enrich the defect prompt c_d and set w_d.
        """
        print(f"\\n{'='*60}")
        print("[Agent] Step 3: Preparing DefectDiffu generation conditions...")

        domain = self.router.route(product_description)

        # --- RAG: Retrieve in-context defect examples ---
        rag_examples = []
        try:
            # DefectRAG expects an object with artifact_type, description, target_entity
            rag_query_obj = type('RAGQuery', (), {
                'artifact_type': defect_plan.defect_type,
                'description': defect_plan.description,
                'target_entity': defect_plan.target_entity
            })()
            rag_examples = self.rag.retrieve(
                rag_query_obj,
                k=rag_k,
                domain_filter=domain if domain != "general" else None,
                commercial_only=True
            )
            if rag_examples:
                print(f"[RAG] Retrieved {len(rag_examples)} example(s) for '{defect_plan.description}'")
                for ex in rag_examples:
                    print(f"       → {ex.get('defect_name', 'unknown')} ({ex.get('domain', 'unknown')})")
            else:
                print(f"[RAG] No examples found for '{defect_plan.description}'")
        except Exception as e:
            print(f"[RAG] Retrieval failed: {e}")

        # --- Build prompts ---
        # c_p: product / background consistency
        c_p = f"A photo of {product_description}"

        # c_d: defect consistency — enrich with RAG if available
        clean_defect_plan_description = defect_plan.description.replace("A photo of ", "")
        defect_name = clean_defect_plan_description
        if rag_examples:
            # Use the most similar example's caption to enrich
            rag_defect_desc = rag_examples[0].get('caption', '')
            if rag_defect_desc:
                defect_name = f"{clean_defect_plan_description}, {rag_defect_desc}"
        c_d = f"A photo of {defect_name}"

        # c_f: fusion prompt
        c_f = f"A photo of {product_description} with {clean_defect_plan_description}"

        # --- Double-free scales ---
        severity_to_wd = {"low": 0.6, "minor": 0.6,
                          "medium": 1.0, "moderate": 1.0,
                          "high": 1.5, "severe": 1.5}
        w_d = severity_to_wd.get(getattr(defect_plan, 'severity', 'medium').lower(), 1.0)

        # If RAG suggests a strength adjustment, apply it
        for ex in rag_examples:
            meta_wd = ex.get('metadata', {}).get('recommended_wd')
            if meta_wd is not None:
                w_d = float(meta_wd)
                print(f"[RAG] Adjusted w_d to {w_d} from retrieved example")
                break

        w_p = 1.0  # Default product consistency; increase if background drifts

        print(f"[Agent] c_p: {c_p}")
        print(f"[Agent] c_d: {c_d}")
        print(f"[Agent] c_f: {c_f}")
        print(f"[Agent] w_d={w_d}, w_p={w_p}")

        return {
            'c_p': c_p,
            'c_d': c_d,
            'c_f': c_f,
            'w_d': w_d,
            'w_p': w_p,
            'rag_examples': rag_examples,
            'rag_domain': domain,
            'defect_plan': defect_plan
        }

    # ------------------------------------------------------------------
    # Step 4: Synthesize with DefectDiffu
    # ------------------------------------------------------------------

    def synthesize(self, gen_conditions: Dict) -> Tuple[Image.Image, np.ndarray]:
        """Generate defect image + mask with DefectDiffu."""
        print(f"\\n{'='*60}")
        print("[Agent] Step 4: Synthesizing defect with DefectDiffu...")

        self._init_defectdiffu()

        img, mask, meta = self.defectdiffu_generator.generate_from_plan(
            product_description=gen_conditions['c_p'].replace("A photo of ", ""),
            defect_description=gen_conditions['c_d'].replace("A photo of ", ""),
            w_d=gen_conditions['w_d'],
            w_p=gen_conditions['w_p'],
            seed=42
        )

        print("[Agent] DefectDiffu synthesis complete")
        return img, mask

    # ------------------------------------------------------------------
    # Step 5: Verification
    # ------------------------------------------------------------------

    def verify(self, original_image: np.ndarray, generated_image: Image.Image,
               defect_mask: np.ndarray, defect_plan) -> Dict:
        """
        VLM verification of generated defect.
        Since DefectDiffu generates from noise (not editing the original),
        we verify that the generated image contains the planned defect
        in a plausible location.
        """
        print(f"\\n{'='*60}")
        print("[Agent] Step 5: Verifying generated defect...")

        # Crop to the mask region (or full image if mask is empty)
        mask_bool = defect_mask.astype(bool)
        if mask_bool.sum() == 0:
            print("[Agent] Warning: Empty mask; verifying full image")
            y1, x1 = 0, 0
            y2, x2 = original_image.shape[0], original_image.shape[1]
        else:
            ys, xs = np.where(mask_bool)
            y1, y2 = ys.min(), ys.max()
            x1, x2 = xs.min(), xs.max()
            # Add margin
            margin = 32
            h, w = original_image.shape[:2]
            y1 = max(0, y1 - margin)
            x1 = max(0, x1 - margin)
            y2 = min(h, y2 + margin)
            x2 = min(w, x2 + margin)

        gen_crop = np.array(generated_image)[y1:y2, x1:x2]
        orig_crop = original_image[y1:y2, x1:x2]

        obj_name = f"a {defect_plan.target_subentity or defect_plan.target_entity}"

        result = artifact_description(
            self.vlm_client,
            original_image,      # masked original (full image)
            orig_crop,             # original crop
            gen_crop,              # generated crop
            obj_name,
            defect_plan.defect_type,
            self.money_manager
        )

        print(f"[Agent] Verification result: has_artifact={result.has_artifact}")
        print(f"[Agent] Explanation: {result.explanation}")
        print(f"[Agent] Label: {result.label}")

        return {
            'passed': result.has_artifact,
            'explanation': result.explanation,
            'label': result.label
        }

    # ------------------------------------------------------------------
    # Main pipeline
    # ------------------------------------------------------------------

    def run(self, product_description: str, image_path: str,
            caption: Optional[str] = None, num_defects: int = 3,
            defect_type: Optional[str] = None) -> Dict:
        """Run the full agentic pipeline with DefectDiffu."""
        start_time = datetime.now()
        exp_id = str(uuid.uuid4())[:8]

        image = np.array(Image.open(image_path).convert('RGB'))
        # Resize to DefectDiffu resolution if needed
        if image.shape[0] != self.image_size or image.shape[1] != self.image_size:
            image_pil = Image.fromarray(image).resize((self.image_size, self.image_size), Image.LANCZOS)
            image = np.array(image_pil)
            print(f"[Agent] Resized image to {self.image_size}x{self.image_size} for DefectDiffu")

        print(f"[Agent] Loaded image: {image.shape}")

        plan = self.plan(product_description, image, defect_type=defect_type, num_defects=num_defects)

        results = []
        defects_to_process = plan.possible_defects[:num_defects]
        print(f"[Agent] Processing top {len(defects_to_process)} of {len(plan.possible_defects)} planned defects")

        for i, defect_plan in enumerate(defects_to_process):
            print(f"\\n{'='*80}")
            print(f"[Agent] Defect {i+1}/{len(defects_to_process)}: [{defect_plan.defect_type.upper()}] {defect_plan.description}")
            print(f"{'='*80}")

            try:
                # Step 2: Perceive (VLM ROI bounding box for placement)
                prediction = self.perceive(image, defect_plan)
                target_bbox = prediction.get('bbox', [0, 0, image.shape[1], image.shape[0]])

                # Step 3: Prepare conditions
                gen_conditions = self.prepare_generation_conditions(
                    defect_plan, product_description=product_description
                )

                # Step 4: Generate patch from DefectDiffu
                generated_image, defect_mask = self.synthesize(gen_conditions)

                # Default ratio presets per defect type
                DEFECT_RATIO_PRESETS = {
                    "bubble": 0.25,        # Small localized bubble
                    "pinhole": 0.15,       # Very small point defect
                    "particle": 0.20,      # Dust / particle
                    "scratch": 0.40,       # Medium line scratch
                    "crack": 0.50,         # Medium crack
                    "smudge": 0.85,        # Large surface coverage
                    "contamination": 1.0, # 100% full ROI coverage
                    "discoloration": 1.0, # 100% full ROI coverage
                    "residue": 0.80        # Large area residue
                }

                # Get defect type from defect_plan
                current_defect_type = getattr(defect_plan, 'defect_type', '').lower()

                #    Pick preset ratio (defaults to None / 100% if defect type isn't in presets)
                # selected_ratio = DEFECT_RATIO_PRESETS.get(current_defect_type, None)

                # Extract VLM dynamic scale & shape decisions
                vlm_ratio = getattr(defect_plan, 'defect_coverage_ratio', None)
                vlm_shape = getattr(defect_plan, 'mask_shape', 'free')

                # Fallback ratio if VLM didn't specify
                if vlm_ratio is None:
                    current_defect_type = getattr(defect_plan, 'defect_type', '').lower()
                    vlm_ratio = DEFECT_RATIO_PRESETS.get(current_defect_type, 1.0)

                # Step 4.5: Blend defect onto REAL image using dynamic ratio
                blended_image, full_defect_mask = blend_defect_onto_real_image(
                    real_image=image,
                    defect_image=generated_image,
                    defect_mask=defect_mask,
                    target_bbox=target_bbox,
                    max_defect_ratio=vlm_ratio,
                    mask_shape=vlm_shape
                )

                # Extract the exact pixel bounding box from full_defect_mask
                ys, xs = np.where(full_defect_mask > 0)

                if len(ys) > 0 and len(xs) > 0:
                    # Exact coordinates of the blended defect
                    mask_bbox = [int(xs.min()), int(ys.min()), int(xs.max()), int(ys.max())]
                else:
                    # Fallback to target_bbox if mask is empty
                    mask_bbox = target_bbox

                # Draw the green verification box around the EXACT mask coordinates
                blended_with_green_box = create_visual_prompt_image(blended_image, mask_bbox)

                # Step 5: Verify blended result
                verification = self.verify(image, Image.fromarray(blended_image), full_defect_mask, defect_plan)

                # Save outputs
                # Dynamically set target folder depending on VLM verification result
                if not verification['passed']:
                    defect_dir = self.output_dir / "failed" / f"{exp_id}_defect_{i}_{defect_plan.defect_type}"
                else:
                    defect_dir = self.output_dir / f"{exp_id}_defect_{i}_{defect_plan.defect_type}"

                defect_dir.mkdir(parents=True, exist_ok=True)

                Image.fromarray(image).save(defect_dir / "real_clean_image.png")
                Image.fromarray(blended_image).save(defect_dir / "blended_factory_defect.png") # Final injected photo
                generated_image.save(defect_dir / "raw_defectdiffu_patch.png")

                # Save pixel-exact segmentation mask for model training
                mask_img = Image.fromarray(full_defect_mask * 255)
                mask_img.save(defect_dir / "defect_mask.png")

                metadata = {
                    'experiment_id': exp_id,
                    'product_description': product_description,
                    'product_type': plan.product_type,
                    'defect_plan': defect_plan.dict() if hasattr(defect_plan, 'dict') else vars(defect_plan),
                    'generation_conditions': {k: v for k, v in gen_conditions.items() if k != 'defect_plan'},
                    'verification': verification,
                    'timestamp': datetime.now().isoformat()
                }
                with open(defect_dir / "metadata.json", 'w') as f:
                    json.dump(metadata, f, indent=2, default=str)

                if not verification['passed']:
                    results.append({
                        'defect_type': defect_plan.defect_type,
                        'success': False,
                        'verification_passed': False,
                        'error': f"VLM verification failed: {verification.get('explanation', 'no explanation')}"
                    })
                    print(f"[Agent] Verification FAILED for {defect_plan.defect_type}.")
                else:
                    results.append({
                        'defect_type': defect_plan.defect_type,
                        'success': True,
                        'verification_passed': verification['passed'],
                        'output_dir': str(defect_dir)
                    })
                    print(f"[Agent] Defect {i+1} complete. Saved to {defect_dir}")

            except Exception as e:
                print(f"[Agent] ERROR processing defect {i+1}: {str(e)}")
                traceback.print_exc()
                results.append({
                    'defect_type': defect_plan.defect_type,
                    'success': False,
                    'error': str(e)
                })

        # Cleanup caches
        if hasattr(self, '_gsam_cache'):
            keys_to_remove = [k for k in self._gsam_cache if k.startswith(str(image_path) + "::")]
            for k in keys_to_remove:
                self._gsam_cache.pop(k, None)

        elapsed = (datetime.now() - start_time).total_seconds()
        print(f"\\n{'='*60}")
        print(f"[Agent] Pipeline complete in {elapsed:.1f}s")
        print(f"[Agent] Results: {sum(1 for r in results if r['success'])}/{len(results)} succeeded")

        return {
            'experiment_id': exp_id,
            'product_type': plan.product_type,
            'results': results,
            'output_dir': str(self.output_dir),
            'elapsed_time': elapsed
        }

    def cleanup(self):
        """Call once after all batch processing is done."""
        if self.gsam_detector:
            self.gsam_detector.cleanup()
            self.gsam_detector = None
        if self.defectdiffu_generator:
            self.defectdiffu_generator.unload_models()
            self.defectdiffu_generator = None
        print("[Agent] All models cleaned up.")


# =============================================================================
# CLI
# =============================================================================

def main():
    parser = argparse.ArgumentParser(description='ArtiAgent — DefectDiffu Edition')
    parser.add_argument('--product-desc', required=True, help='Product description')
    parser.add_argument('--image', required=True, help='Path to clean product image (for planning/verification)')

    # DefectDiffu model paths (REQUIRED)
    parser.add_argument('--defectdiffu-ckpt', required=True,
                        help='Path to trained DefectDiffu checkpoint (.pt)')
    parser.add_argument('--vae-path', required=True,
                        help='Path to Stable Diffusion VAE (e.g. stabilityai/sd-vae-ft-mse)')

    # Generation control
    parser.add_argument('--defect-type', default=None,
                        help='Specific defect type to generate (e.g., bubble, scratch)')
    parser.add_argument('--output-dir', default='./defect_output', help='Output directory')
    parser.add_argument('--caption', default=None, help='Optional image caption')
    parser.add_argument('--num-defects', type=int, default=3,
                        help='Number of defects to generate (default: 3)')
    parser.add_argument('--device', default='cuda', help='Device (cuda/cpu)')
    parser.add_argument('--vlm-model', default='gemma3:12b', help='Local VLM model')
    parser.add_argument('--image-size', type=int, default=512,
                        help='DefectDiffu generation resolution (default: 512)')
    parser.add_argument('--num-steps', type=int, default=50,
                        help='Denoising steps for DefectDiffu (default: 50)')

    args = parser.parse_args()

    orchestrator = ArtiAgentOrchestrator(
        device=args.device,
        output_dir=args.output_dir,
        vlm_model=args.vlm_model,
        defectdiffu_ckpt=args.defectdiffu_ckpt,
        vae_path=args.vae_path,
        image_size=args.image_size,
        num_steps=args.num_steps
    )

    result = orchestrator.run(
        product_description=args.product_desc,
        image_path=args.image,
        caption=args.caption,
        num_defects=args.num_defects,
        defect_type=args.defect_type
    )

    print(f"\\nFinal output saved to: {result['output_dir']}")


if __name__ == "__main__":
    main()