"""
Batch Agent Orchestrator — DefectFill Edition (with VLM list selection)

Auto-discovers valid object classes and defect types from the checkpoint directory.
VLM selects object_class from the discovered list based on product description.
If defect_type is not provided, VLM selects from the per-object-class list.

Usage:
    # Auto-discover checkpoints, VLM selects everything from product description
    python batch_agent_orchestrator.py \
        --input-dir "C:/TestingImage" \
        --output-dir "C:/AgentOutput" \
        --checkpoint-dir "C:/.../checkpoints" \
        --product-desc "VCSEL laser diode with emission aperture" \
        --device cuda

    # Provide explicit valid lists (overrides auto-discovery)
    python batch_agent_orchestrator.py \
        --input-dir "C:/TestingImage" \
        --output-dir "C:/AgentOutput" \
        --checkpoint-dir "C:/.../checkpoints" \
        --valid-object-classes '["xray_PCB","vcsel"]' \
        --valid-defect-types '{"xray_PCB":["xray_die","bubble"],"vcsel":["scratch"]}' \
        --product-desc "VCSEL laser diode" \
        --device cuda

    # CSV manifest with per-image routing (overrides VLM selection for those images)
    python batch_agent_orchestrator.py \
        --input-dir "C:/TestingImage" \
        --output-dir "C:/AgentOutput" \
        --checkpoint-dir "C:/.../checkpoints" \
        --manifest "C:/products.csv" \
        --device cuda
"""

import os
import sys
import csv
import json
import argparse
import traceback
from pathlib import Path
from datetime import datetime
from typing import Dict, List, Optional
from tqdm import tqdm

import numpy as np
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 artiagent_orchestrator import ArtiAgentOrchestrator


def infer_product_from_path(image_path: Path) -> str:
    """Infer product description from folder structure or filename."""
    parent = image_path.parent.name.lower()
    if parent and parent not in ['.', '', 'images', 'imgs', 'data', 'input']:
        return parent.replace('_', ' ').replace('-', ' ')
    stem = image_path.stem.lower()
    for keyword in ['vcsel', 'lens', 'die', 'photodiode', 'sensor', 'chip', 'led', 'laser', 'optical']:
        if keyword in stem:
            return keyword
    return "electronic component"


def infer_object_class_from_path(image_path: Path) -> str:
    """Infer object_class from parent folder name."""
    parent = image_path.parent.name.lower()
    if parent and parent not in ['.', '', 'images', 'imgs', 'data', 'input']:
        return parent
    return ""


def discover_checkpoints(checkpoint_dir: str) -> Dict[str, List[str]]:
    """
    Scan checkpoint directory to discover available object_class -> defect_type mappings.
    Expected structure:
        checkpoint_dir/
          xray_PCB/
            xray_die/
              checkpoints/checkpoint_final.pt
            bubble/
              checkpoints/checkpoint_final.pt
    """
    cp = Path(checkpoint_dir)
    mapping = {}
    if not cp.exists():
        print(f"[Batch] Warning: checkpoint-dir does not exist: {checkpoint_dir}")
        return mapping

    for obj_dir in cp.iterdir():
        if not obj_dir.is_dir():
            continue
        defect_types = []
        for defect_dir in obj_dir.iterdir():
            if not defect_dir.is_dir():
                continue
            ckpt1 = defect_dir / "checkpoints" / "checkpoint_final.pt"
            ckpt2 = defect_dir / "checkpoint_final.pt"
            if ckpt1.exists() or ckpt2.exists():
                defect_types.append(defect_dir.name)
        if defect_types:
            mapping[obj_dir.name] = defect_types

    print(f"[Batch] Auto-discovered checkpoints: {json.dumps(mapping, indent=2)}")
    return mapping


def load_manifest(manifest_path: str) -> Dict[str, Dict]:
    """Load CSV manifest mapping image paths to object_class, defect_type, and product descriptions."""
    manifest = {}
    with open(manifest_path, 'r', encoding='utf-8') as f:
        reader = csv.DictReader(f)
        for row in reader:
            img_path = row.get('image_path', row.get('path', row.get('image', ''))).strip()
            desc = row.get('product_description', row.get('description', row.get('product', ''))).strip()
            obj_cls = row.get('object_class', row.get('object', '')).strip()
            dfc_type = row.get('defect_type', row.get('defect', '')).strip()
            if img_path:
                manifest[Path(img_path).resolve()] = {
                    'product_description': desc,
                    'object_class': obj_cls,
                    'defect_type': dfc_type
                }
    print(f"[Batch] Loaded manifest with {len(manifest)} entries")
    return manifest


def discover_images(input_dir: str, extensions=('.png', '.jpg', '.jpeg', '.bmp', '.tif', '.tiff')) -> List[Path]:
    """Recursively discover all images in input directory."""
    input_path = Path(input_dir)
    images = []
    for ext in extensions:
        images.extend(input_path.rglob(f"*{ext}"))
        images.extend(input_path.rglob(f"*{ext.upper()}"))
    unique = sorted(set(images))
    print(f"[Batch] Discovered {len(unique)} images in {input_dir}")
    return unique


def run_batch(
    input_dir: str,
    output_dir: str,
    checkpoint_dir: str,
    object_class: Optional[str] = None,
    defect_type: Optional[str] = None,
    product_desc: Optional[str] = None,
    manifest_path: Optional[str] = None,
    valid_object_classes: Optional[List[str]] = None,
    valid_defect_types: Optional[Dict[str, List[str]]] = None,
    max_defects_per_image: int = 3,
    device: str = 'cuda',
    vlm_model: str = 'gemma3:12b',
    image_size: int = 512,
    num_steps: int = 50,
    guidance_scale: float = 7.5,
    resume: bool = False,
    save_failed: bool = True
):
    """Run agent orchestrator over all images in input directory."""

    timestamp = datetime.now().strftime("%Y%m%d_%H%M")
    output_path = Path(output_dir) / timestamp
    output_path.mkdir(parents=True, exist_ok=True)
    print(f"[Batch] Output folder: {output_path}")

    # ------------------------------------------------------------------
    # Auto-discover valid lists from checkpoint directory if not provided
    # ------------------------------------------------------------------
    if valid_defect_types is None:
        valid_defect_types = discover_checkpoints(checkpoint_dir)
    if valid_object_classes is None:
        valid_object_classes = list(valid_defect_types.keys())

    print(f"[Batch] Valid object classes: {valid_object_classes}")
    print(f"[Batch] Valid defect types mapping: {json.dumps(valid_defect_types, indent=2)}")

    # ------------------------------------------------------------------
    # Load manifest if provided
    # ------------------------------------------------------------------
    manifest = {}
    if manifest_path and os.path.exists(manifest_path):
        manifest = load_manifest(manifest_path)
        print(f"[Batch] Scenario A Active: CSV Manifest ({len(manifest)} entries)")
    elif product_desc:
        print(f"[Batch] Scenario B Active: Global Description = '{product_desc}'")
        if object_class:
            print(f"[Batch] Global object_class = '{object_class}'")
        if defect_type:
            print(f"[Batch] Global defect_type = '{defect_type}'")
    else:
        print("[Batch] Scenario C Active: Folder Name Inference (no description provided)")

    images = discover_images(input_dir)
    if not images:
        print("[Batch] No images found. Exiting.")
        return

    progress_file = output_path / "batch_progress.json"
    processed_ids = set()
    if resume and progress_file.exists():
        with open(progress_file, 'r') as f:
            progress = json.load(f)
            processed_ids = set(progress.get('processed_paths', []))
        print(f"[Batch] Resuming: {len(processed_ids)} images already processed")

    # ------------------------------------------------------------------
    # Initialize orchestrator once with the valid lists
    # ------------------------------------------------------------------
    orchestrator = ArtiAgentOrchestrator(
        device=device,
        output_dir=str(output_path),
        vlm_model=vlm_model,
        checkpoint_dir=checkpoint_dir,
        object_class=object_class or "",
        defect_type=defect_type or "",
        valid_object_classes=valid_object_classes,
        valid_defect_types=valid_defect_types,
        image_size=image_size,
        num_steps=num_steps,
        guidance_scale=guidance_scale
    )

    stats = {
        'total': len(images),
        'processed': 0,
        'successful': 0,
        'failed': 0,
        'defects_generated': 0,
        'start_time': datetime.now().isoformat(),
        'processed_paths': [],
        'failed_images': []
    }

    if resume:
        images = [img for img in images if str(img.resolve()) not in processed_ids]

    print(f"[Batch] Processing {len(images)} images...")
    print(f"[Batch] Max defects per image: {max_defects_per_image}")
    print("=" * 70)

    for img_path in tqdm(images, desc="Agent Batch Processing"):
        img_key = str(img_path.resolve())

        try:
            # Resolve description, object_class, and defect_type per image
            if img_key in manifest:
                entry = manifest[img_key]
                desc = entry.get('product_description') or product_desc or infer_product_from_path(img_path)
                obj_cls = entry.get('object_class') or object_class or None
                dfc_type = entry.get('defect_type') or defect_type or None
                source = "manifest"
            elif product_desc:
                desc = product_desc
                obj_cls = object_class or None
                dfc_type = defect_type or None
                source = "global"
            else:
                desc = infer_product_from_path(img_path)
                obj_cls = object_class or infer_object_class_from_path(img_path) or None
                dfc_type = defect_type or None
                source = "inferred"

            # If object_class is still None, VLM will select from valid_object_classes
            # If defect_type is still None, VLM will select from valid_defect_types for the chosen object_class
            # If object_class is provided but not in valid list, warn
            if obj_cls and valid_object_classes and obj_cls not in valid_object_classes:
                print(f"[Batch] Warning: object_class '{obj_cls}' not in valid list {valid_object_classes}. "
                      f"VLM will select a valid one.")
                obj_cls = None

            print(f"\n[Batch] Processing: {img_path.name} | desc source: {source}")
            if obj_cls:
                print(f"[Batch] object_class: '{obj_cls}' (user-provided)")
            else:
                print(f"[Batch] object_class: <VLM will select from {valid_object_classes}>")
            if dfc_type:
                print(f"[Batch] defect_type: '{dfc_type}' (user-provided)")
            else:
                print(f"[Batch] defect_type: <VLM will select from valid list>")
            if source in ['inferred', 'global']:
                print(f"[Batch] Using description: '{desc}'")

            result = orchestrator.run(
                product_description=desc,
                image_path=str(img_path),
                max_defects=max_defects_per_image,
                defect_type=dfc_type,
                object_class=obj_cls
            )

            successful_defects = sum(1 for r in result['results'] if r['success'])

            stats['processed'] += 1
            stats['successful'] += 1 if successful_defects > 0 else 0
            stats['defects_generated'] += successful_defects
            stats['processed_paths'].append(img_key)

            if successful_defects == 0:
                stats['failed'] += 1
                stats['failed_images'].append({'path': img_key, 'reason': 'no_defects_generated'})

            if stats['processed'] % 5 == 0:
                with open(progress_file, 'w') as f:
                    json.dump(stats, f, indent=2)

        except Exception as e:
            stats['failed'] += 1
            stats['failed_images'].append({'path': img_key, 'reason': str(e)})
            print(f"[Batch] FAILED: {img_path.name} -> {str(e)}")
            if save_failed:
                fail_dir = output_path / "_failed" / img_path.stem
                fail_dir.mkdir(parents=True, exist_ok=True)
                with open(fail_dir / "error.txt", 'w') as f:
                    f.write(traceback.format_exc())

    with open(progress_file, 'w') as f:
        json.dump(stats, f, indent=2)

    orchestrator.cleanup()

    elapsed = (datetime.now() - datetime.fromisoformat(stats['start_time'])).total_seconds()
    hours = int(elapsed // 3600)
    minutes = int((elapsed % 3600) // 60)
    seconds = int(elapsed % 60)

    print("\n" + "=" * 70)
    print("BATCH ORCHESTRATION COMPLETE")
    print("=" * 70)
    print(f"Total images:      {stats['total']}")
    print(f"Processed:         {stats['processed']}")
    print(f"Successful:        {stats['successful']}")
    print(f"Failed:            {stats['failed']}")
    print(f"Defects generated: {stats['defects_generated']}")
    print(f"Total time:        {hours}h {minutes}m {seconds}s")
    print(f"Output directory:  {output_path}")
    print("=" * 70)


def main():
    parser = argparse.ArgumentParser(description='Batch Agent-Driven Defect Generation (DefectFill with VLM selection)')

    # Input / Output
    parser.add_argument('--input-dir', required=True, help='Directory containing clean product images')
    parser.add_argument('--output-dir', required=True, help='Output directory for all defect images')

    # DefectFill checkpoint routing
    parser.add_argument('--checkpoint-dir', required=True,
                        help='Root directory containing object_class/defect_type checkpoint subfolders')

    # Description / routing sources
    parser.add_argument('--product-desc', default=None,
                        help='Global product description applied to ALL images (drives VLM selection)')
    parser.add_argument('--object-class', default=None,
                        help='Global object_class for ALL images (optional; VLM selects if omitted)')
    parser.add_argument('--manifest', default=None,
                        help='CSV manifest with columns: image_path,object_class,defect_type,product_description')

    # Defect control
    parser.add_argument('--defect-type', default=None,
                        help='Global defect type for ALL images (optional; VLM selects if omitted)')
    parser.add_argument('--max-defects-per-image', type=int, default=3,
                        help='Maximum defects to generate per image (default: 3)')

    # Valid lists (auto-discovered from checkpoint-dir if not provided)
    parser.add_argument('--valid-object-classes', default=None,
                        help="JSON array of valid object classes, e.g., '[\"xray_PCB\",\"vcsel\"]'" )
    parser.add_argument('--valid-defect-types', default=None,
                        help="JSON dict mapping object_class to defect types, e.g., '{\"xray_PCB\":[\"xray_die\",\"bubble\"]}'" )
    # Generation control
    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='Generation resolution')
    parser.add_argument('--num-steps', type=int, default=50, help='Denoising steps')
    parser.add_argument('--guidance-scale', type=float, default=7.5, help='CFG scale')
    parser.add_argument('--resume', action='store_true', help='Resume from previous batch run')
    parser.add_argument('--no-save-failed', action='store_true', help='Do not save failed case logs')

    args = parser.parse_args()

    # Parse valid lists from CLI
    valid_object_classes = None
    valid_defect_types = None
    if args.valid_object_classes:
        valid_object_classes = json.loads(args.valid_object_classes)
    if args.valid_defect_types:
        valid_defect_types = json.loads(args.valid_defect_types)

    run_batch(
        input_dir=args.input_dir,
        output_dir=args.output_dir,
        checkpoint_dir=args.checkpoint_dir,
        object_class=args.object_class,
        defect_type=args.defect_type,
        product_desc=args.product_desc,
        manifest_path=args.manifest,
        valid_object_classes=valid_object_classes,
        valid_defect_types=valid_defect_types,
        max_defects_per_image=args.max_defects_per_image,
        device=args.device,
        vlm_model=args.vlm_model,
        image_size=args.image_size,
        num_steps=args.num_steps,
        guidance_scale=args.guidance_scale,
        resume=args.resume,
        save_failed=not args.no_save_failed
    )


if __name__ == "__main__":
    main()