app.py

"""
DefectFill & DefectDiffu Unified Web UI - Flask Backend
=========================================================
Supports both DefectFill (few-shot LoRA training) and DefectDiffu (text-guided generation).

Run from project root:
    python defectfill_webui/app.py

Requirements:
    pip install flask flask-cors pillow numpy
"""

import os
import sys
import json
import time
import shutil
import threading
import subprocess
import traceback
from pathlib import Path
from datetime import datetime
from PIL import Image, ImageDraw
import numpy as np
import uuid
import io
import re
import importlib.util


# =============================================================================
# Configuration
# =============================================================================

BASE_DIR = Path(__file__).parent
PROJECT_ROOT = BASE_DIR.parent  # ArtiAgent_with_WebUI/

# DefectFill paths
DEFECTFILL_ROOT = PROJECT_ROOT / "ArtiAgent - DefectFill"
TRAIN_PY = DEFECTFILL_ROOT / 'engine' / 'DefectFill' / 'train.py'
TRAINING_DATA_DIR = DEFECTFILL_ROOT / 'src' / 'data'
CHECKPOINT_DIR = DEFECTFILL_ROOT / 'engine' / 'DefectFill' / 'checkpoints'
TRAINED_GALLERY_DIR = DEFECTFILL_ROOT / "storage" / "trained_products_gallery"
TRAINED_GALLERY_DIR.mkdir(parents=True, exist_ok=True)
TEMP_IMAGE_DIR = DEFECTFILL_ROOT / 'src' / 'data' / 'tempImage'
TEMP_IMAGE_DIR.mkdir(parents=True, exist_ok=True)

# DefectDiffu paths
DEFECTDIFFU_ROOT = PROJECT_ROOT / "ArtiAgent - DefectDiffu"
DEFECTDIFFU_ENGINE = PROJECT_ROOT / "engine" / "DefectDiffu"

# Web UI dirs
UPLOAD_FOLDER = BASE_DIR / "uploads"
OUTPUT_DIR = BASE_DIR / "outputs"
CONFIG_DIR = BASE_DIR / "config"

# OpenCV availability check
try:
    import cv2
    CV2_AVAILABLE = True
except ImportError:
    CV2_AVAILABLE = False
    print("[WARNING] opencv-python not installed. Image alignment unavailable.")

# Ensure directories exist
for d in [UPLOAD_FOLDER, OUTPUT_DIR, CONFIG_DIR]:
    d.mkdir(parents=True, exist_ok=True)

# Domain instructions config
DOMAIN_CONFIG_PATH = CONFIG_DIR / "domain_instructions.json"
if not DOMAIN_CONFIG_PATH.exists():
    with open(DOMAIN_CONFIG_PATH, "w") as f:
        json.dump({"examples": {}, "active": {}}, f, indent=2)

# DefectDiffu config
DEFECTDIFFU_CONFIG_PATH = CONFIG_DIR / "defectdiffu_config.json"
if not DEFECTDIFFU_CONFIG_PATH.exists():
    with open(DEFECTDIFFU_CONFIG_PATH, "w") as f:
        json.dump({"ckpt_path": "", "vae_path": "", "vlm_model": "gemma3:12b"}, f, indent=2)

# Setup sys.path for imports
# sys.path.insert(0, str(DEFECTFILL_ROOT / 'src' / 'segment_anything'))
# sys.path.insert(0, str(DEFECTFILL_ROOT / 'src'))
# sys.path.insert(0, str(DEFECTFILL_ROOT / 'pipeline'))
# sys.path.insert(0, str(DEFECTFILL_ROOT))
# sys.path.insert(0, str(DEFECTDIFFU_ROOT))
# sys.path.insert(0, str(DEFECTDIFFU_ROOT / 'src'))
# sys.path.insert(0, str(DEFECTDIFFU_ROOT / 'pipeline'))

# =============================================================================
# Dynamic Import Helpers (avoid module name collisions)
# =============================================================================

def load_orchestrator_from_path(module_path, class_name="ArtiAgentOrchestrator", extra_paths=None):
    """
    Dynamically import a class from a Python file, temporarily adding
    extra_paths to sys.path so that the module's internal imports
    resolve to the correct project.
    """
    if not module_path.exists():
        print(f"[ERROR] Module file not found: {module_path}")
        return None

    original_sys_path = sys.path[:]   # save current state

    try:
        # Insert project‑specific paths at the front
        if extra_paths:
            for p in reversed(extra_paths):   # reverse to keep the given order
                p_str = str(p)
                if p_str not in sys.path:
                    sys.path.insert(0, p_str)

        # Also ensure the module's own directory is available
        module_dir = module_path.parent
        if module_dir not in sys.path:
            sys.path.insert(0, str(module_dir))

        # Use a unique module name to avoid caching collisions
        module_id = f"_dynamic_{class_name}_{uuid.uuid4().hex[:8]}"
        spec = importlib.util.spec_from_file_location(module_id, str(module_path))
        module = importlib.util.module_from_spec(spec)
        spec.loader.exec_module(module)

        cls = getattr(module, class_name, None)
        if cls is None:
            print(f"[ERROR] Class '{class_name}' not found in {module_path}")
        return cls

    except Exception as e:
        print(f"[ERROR] Failed to load {module_path}: {e}")
        traceback.print_exc()
        return None

    finally:
        # Restore original sys.path to avoid leaking
        sys.path[:] = original_sys_path


# DefectFill imports — load directly from file to avoid sys.path collision with DefectFill
extra_paths_df = [
    DEFECTFILL_ROOT / 'src' / 'segment_anything',
    DEFECTFILL_ROOT / 'src',
    DEFECTFILL_ROOT / 'pipeline',
    DEFECTFILL_ROOT,
]

df_orchestrator_path = DEFECTFILL_ROOT / "src" / "artiagent_orchestrator.py"
DefectFillOrchestrator = load_orchestrator_from_path(
    df_orchestrator_path,
    "ArtiAgentOrchestrator",
    extra_paths_df
)
print("[INFO] DefectFill orchestrator is ready.", flush=True)


# DefectDiffu imports — load directly from file to avoid name collision with DefectDiffu
extra_paths_dd = [
    DEFECTDIFFU_ROOT / 'src',
    DEFECTDIFFU_ROOT / "engine" / "DefectDiffu",
    DEFECTDIFFU_ROOT,
]

dd_orchestrator_path = DEFECTDIFFU_ROOT / 'src' / "artiagent_orchestrator.py"
DefectDiffuOrchestrator = load_orchestrator_from_path(
    dd_orchestrator_path,
    "ArtiAgentOrchestrator",
    extra_paths_dd
)
print("[INFO] DefectDiffu orchestrator is ready.")

def load_helper_from_path(module_path, function_name):
    # Use a generic loader (can reuse load_orchestrator_from_path with no class)
    # But for functions, you can load the module and extract the function.
    original_sys_path = sys.path[:]
    try:
        # Add DefectFill paths (or pass as extra_paths)
        sys.path.insert(0, str(DEFECTFILL_ROOT / 'src'))
        sys.path.insert(0, str(DEFECTFILL_ROOT))
        spec = importlib.util.spec_from_file_location(f"_helper_{uuid.uuid4().hex[:8]}", str(module_path))
        module = importlib.util.module_from_spec(spec)
        spec.loader.exec_module(module)
        return getattr(module, function_name, None)
    finally:
        sys.path[:] = original_sys_path

# Usage:
align_image_to_reference = load_helper_from_path(
    DEFECTFILL_ROOT / 'src' / 'align_image_to_reference.py', 'align_image_to_reference'
)
check_product_training_status = load_helper_from_path(
    DEFECTFILL_ROOT / 'src' / 'detect_similar_product.py', 'check_product_training_status'
)

from flask import Flask, render_template, request, jsonify, send_from_directory
from flask_cors import CORS

app = Flask(__name__, template_folder=str(BASE_DIR / "templates"), static_folder=str(BASE_DIR / "static"))
app.config["MAX_CONTENT_LENGTH"] = 100 * 1024 * 1024  # 100MB upload limit
CORS(app)

# =============================================================================
# Global State
# =============================================================================

training_state = {
    "running": False,
    "process": None,
    "object_class": None,
    "defect_type": None,
    "output_dir": None,
    "log_lines": [],
    "start_time": None,
    "completed": False,
    "error": None
}

training_lock = threading.Lock()

# --- Generation Job Tracking ---
generation_jobs = {}
jobs_lock = threading.Lock()


class ProgressCapture:
    """Captures BOTH stdout and stderr, parses [Agent] Step X: markers."""
    def __init__(self, job_id):
        self.job_id = job_id
        self._buffer = ""
        self._step_re = re.compile(r'\[Agent\] Step (\d+):\s*(.+)')
        self._defect_re = re.compile(r'\[Agent\] Defect (\d+)/(\d+):')

    def write(self, s):
        self._buffer += s
        while '\n' in self._buffer:
            line, self._buffer = self._buffer.split('\n', 1)
            self._process_line(line)

    def _process_line(self, line):
        m = self._step_re.search(line)
        if m:
            with jobs_lock:
                if self.job_id in generation_jobs:
                    generation_jobs[self.job_id]['step'] = int(m.group(1))
                    generation_jobs[self.job_id]['step_text'] = m.group(2).strip()
        d = self._defect_re.search(line)
        if d:
            with jobs_lock:
                if self.job_id in generation_jobs:
                    job = generation_jobs[self.job_id]
                    job['defect_current'] = int(d.group(1))
                    job['defect_total'] = int(d.group(2))
                    # --- ADD THIS BLOCK BELOW ---
                    if 'unit_total' in job and job['unit_total'] > 0:
                        defects_per_image = job['defect_total']
                        image_current = job['image_current']
                        # Formula: (Current Image - 1) * Defects per Image + Current Defect
                        unit_current = (image_current - 1) * defects_per_image + job['defect_current']
                        job['unit_current'] = min(unit_current, job['unit_total'])
        if '[Agent] ERROR' in line or 'Traceback (most recent call last):' in line:
            with jobs_lock:
                if self.job_id in generation_jobs:
                    generation_jobs[self.job_id]['has_error_trace'] = True

    def flush(self):
        pass

    def isatty(self):
        return False


class TeeCapture:
    """Duplicates output to capture AND original stream."""
    def __init__(self, capture, original):
        self.capture = capture
        self.original = original

    def write(self, s):
        self.capture.write(s)
        self.original.write(s)

    def flush(self):
        self.capture.flush()
        self.original.flush()

    def isatty(self):
        return False


def run_generation_job(job_id, params, mode='defectfill'):
    """Background thread: runs generation + captures progress."""
    with jobs_lock:
        generation_jobs[job_id]['status'] = 'running'
        generation_jobs[job_id]['step'] = 1
        generation_jobs[job_id]['step_text'] = 'Planning defects from product description...'
        generation_jobs[job_id]['started_at'] = time.time()

        generation_jobs[job_id]['unit_total'] = len(params['clean_paths']) * params['num_defects']
        generation_jobs[job_id]['unit_current'] = 0

    capture = ProgressCapture(job_id)
    old_stdout = sys.stdout
    old_stderr = sys.stderr
    sys.stdout = TeeCapture(capture, old_stdout)
    sys.stderr = TeeCapture(capture, old_stderr)

    try:
        if mode == 'defectfill':
            orchestrator = DefectFillOrchestrator(
                device=params['device'],
                output_dir=str(params['gen_output_dir']),
                checkpoint_dir=str(params['ckpt_dir']),
                object_class=params['object_class'],
                defect_type=params['defect_type'] or "",
                valid_object_classes=params['valid_object_classes'],
                valid_defect_types=params['valid_defect_types'],
                image_size=params['image_size'],
                num_steps=params['num_steps'],
                guidance_scale=params['guidance_scale'],
                domain_hint=params.get('domain_hint', '')
            )
        else:  # defectdiffu
            orchestrator = DefectDiffuOrchestrator(
                device=params['device'],
                output_dir=str(params['gen_output_dir']),
                vlm_model=params.get('vlm_model', 'gemma3:12b'),
                defectdiffu_ckpt=params['defectdiffu_ckpt'],
                vae_path=params['vae_path'],
                image_size=params['image_size'],
                num_steps=params['num_steps']
            )

        all_output_images = []
        all_grouped_results = []
        total_images = len(params['clean_paths'])

        for img_idx, clean_path in enumerate(params['clean_paths']):
            with jobs_lock:
                generation_jobs[job_id]['image_current'] = img_idx + 1
                generation_jobs[job_id]['image_total'] = total_images
                generation_jobs[job_id]['step_text'] = f'Image {img_idx+1}/{total_images}: Planning defects...'

            if mode == 'defectfill':
                result = orchestrator.run(
                    product_description=params['product_desc'],
                    image_path=clean_path,
                    num_defects=params['num_defects'],
                    defect_type=params['defect_type'],
                    object_class=params['object_class']
                )
            else:  # defectdiffu
                result = orchestrator.run(
                    product_description=params['product_desc'],
                    image_path=clean_path,
                    num_defects=params['num_defects'],
                    defect_type=params['defect_type']
                )

            source_name = Path(clean_path).name
            clean_rel = Path(clean_path).relative_to(BASE_DIR)
            clean_url = f"/api/output/image/{clean_rel}"

            defects_for_image = []

            for r in result.get("results", []):
                if r.get("success") and "output_dir" in r:
                    out_d = Path(r["output_dir"])
                    if out_d.exists():
                        mask_file = None
                        output_file = None
                        blended_file = None
                        raw_file = None
                        for img_file in sorted(out_d.glob("*.png")):
                            fname = img_file.name.lower()
                            if "mask" in fname:
                                mask_file = img_file
                            elif "blended" in fname:
                                blended_file = img_file
                            elif "raw_defectdiffu" in fname or "output" in fname:
                                raw_file = img_file
                            elif "defect" in fname and blended_file is None:
                                output_file = img_file

                            rel_path = img_file.relative_to(BASE_DIR)
                            all_output_images.append({
                                "filename": img_file.name,
                                "url": f"/api/output/image/{rel_path}",
                                "defect_type": r.get("defect_type", "unknown"),
                                "object_class": r.get("object_class", params.get('object_class', 'unknown')),
                                "source_index": img_idx,
                                "source_name": source_name
                            })

                        # DefectDiffu has different output files
                        if mode == 'defectdiffu':
                            defects_for_image.append({
                                "defect_type": r.get("defect_type", "unknown"),
                                "defect_label": f"Defect {len(defects_for_image) + 1}: {r.get('defect_type', 'unknown')}",
                                "clean_image": {
                                    "filename": source_name,
                                    "url": clean_url
                                },
                                "output_image": {
                                    "filename": blended_file.name if blended_file else (output_file.name if output_file else "output.png"),
                                    "url": f"/api/output/image/{(blended_file or output_file or raw_file).relative_to(BASE_DIR)}"
                                } if (blended_file or output_file or raw_file) else None,
                                "raw_patch_image": {
                                    "filename": raw_file.name if raw_file else "raw.png",
                                    "url": f"/api/output/image/{raw_file.relative_to(BASE_DIR)}"
                                } if raw_file else None,
                                "mask_image": {
                                    "filename": mask_file.name if mask_file else "mask.png",
                                    "url": f"/api/output/image/{mask_file.relative_to(BASE_DIR)}"
                                } if mask_file else None
                            })
                        elif mask_file and (output_file or blended_file):
                            defects_for_image.append({
                                "defect_type": r.get("defect_type", "unknown"),
                                "defect_label": f"Defect {len(defects_for_image) + 1}: {r.get('defect_type', 'unknown')}",
                                "clean_image": {
                                    "filename": source_name,
                                    "url": clean_url
                                },
                                "mask_image": {
                                    "filename": mask_file.name,
                                    "url": f"/api/output/image/{mask_file.relative_to(BASE_DIR)}"
                                },
                                "output_image": {
                                    "filename": output_file.name if output_file else blended_file.name,
                                    "url": f"/api/output/image/{(output_file or blended_file).relative_to(BASE_DIR)}"
                                }
                            })

            if defects_for_image:
                all_grouped_results.append({
                    "source_index": img_idx,
                    "source_name": source_name,
                    "defect_count": len(defects_for_image),
                    "defects": defects_for_image
                })

        orchestrator.cleanup()
        sys.stdout = old_stdout
        sys.stderr = old_stderr

        elapsed = time.time() - generation_jobs[job_id]['started_at']

        with jobs_lock:
            generation_jobs[job_id]['status'] = 'completed'
            generation_jobs[job_id]['step'] = 5
            generation_jobs[job_id]['step_text'] = 'Generation complete'
            generation_jobs[job_id]['elapsed_time'] = elapsed
            generation_jobs[job_id]['result'] = {
                'success': True,
                'grouped_results': all_grouped_results,
                'output_images': all_output_images,
                'experiment_id': params.get('object_class', 'defectdiffu'),
                'product_type': params.get('object_class', 'defectdiffu'),
                'elapsed_time': elapsed,
                'results_summary': [{"type": img['defect_type'], "success": True} for img in all_output_images],
                'output_dir': str(params['gen_output_dir'].relative_to(BASE_DIR))
            }
    except Exception as e:
        traceback.print_exc()
        with jobs_lock:
            generation_jobs[job_id]['status'] = 'error'
            generation_jobs[job_id]['error'] = str(e)
            generation_jobs[job_id]['step_text'] = f'Error: {str(e)}'
    finally:
        sys.stdout = old_stdout
        sys.stderr = old_stderr


# =============================================================================
# Helpers
# =============================================================================

def allowed_file(filename, exts={".png", ".jpg", ".jpeg", ".bmp", ".tif", ".tiff"}):
    return Path(filename).suffix.lower() in exts


def get_checkpoints():
    """Discover trained checkpoints: object_class -> [defect_types]."""
    mapping = {}
    if not CHECKPOINT_DIR.exists():
        return mapping
    for obj_dir in CHECKPOINT_DIR.iterdir():
        if not obj_dir.is_dir():
            continue
        if obj_dir.name == 'checkpoints_old':
            continue
        defects = []
        for defect_dir in obj_dir.iterdir():
            if not defect_dir.is_dir():
                continue
            ckpt = defect_dir / "checkpoints" / "checkpoint_final.pt"
            ckpt_alt = defect_dir / "checkpoint_final.pt"
            if ckpt.exists() or ckpt_alt.exists():
                defects.append(defect_dir.name)
        if defects:
            mapping[obj_dir.name] = defects
    return mapping


def generate_mask_from_rect(image_path, rect, output_mask_path):
    """Generate a binary mask from rectangle coordinates."""
    img = Image.open(image_path).convert("RGB")
    w, h = img.size
    mask = Image.new("L", (w, h), 0)
    draw = ImageDraw.Draw(mask)

    x1 = max(0, int(rect["x"]))
    y1 = max(0, int(rect["y"]))
    x2 = min(w, int(rect["x"] + rect["width"]))
    y2 = min(h, int(rect["y"] + rect["height"]))

    draw.rectangle([x1, y1, x2, y2], fill=255)
    mask.save(output_mask_path)
    return {"x1": x1, "y1": y1, "x2": x2, "y2": y2}


def stream_training_logs(process):
    """Read training subprocess output in background thread."""
    global training_state
    try:
        for line in iter(process.stdout.readline, ""):
            line_stripped = line.rstrip()
            if line_stripped:
                with training_lock:
                    training_state["log_lines"].append(line_stripped)
                    training_state["log_lines"] = training_state["log_lines"][-2000:]

        process.stdout.close()
        process.wait()

        with training_lock:
            training_state["running"] = False
            training_state["completed"] = (process.returncode == 0)

            if process.returncode == 0:
                object_class = training_state.get("object_class")
                if object_class:
                    sample_img = None
                    good_dir = TRAINING_DATA_DIR / object_class / "test" / "good"
                    if good_dir.exists():
                        for f in sorted(good_dir.iterdir()):
                            if f.is_file() and allowed_file(f.name):
                                sample_img = f
                                break
                    if not sample_img:
                        prod_dir = TRAINING_DATA_DIR / object_class
                        if prod_dir.exists():
                            for f in sorted(prod_dir.rglob("*")):
                                if f.is_file() and allowed_file(f.name) and "mask" not in f.name.lower():
                                    sample_img = f
                                    break
                    if sample_img:
                        save_to_trained_gallery(str(sample_img), object_class)
                    else:
                        print(f"[Gallery Warning] No clean sample image found for '{object_class}'.")
            else:
                training_state["error"] = f"Training exited with code {process.returncode}"

    except Exception as e:
        with training_lock:
            training_state["running"] = False
            training_state["error"] = f"Log thread crashed: {str(e)}"


def save_to_trained_gallery(clean_image_path: str, product_name: str):
    """Copies 1 clean image to the gallery named after the product for DINOv2 matching."""
    try:
        for existing_file in TRAINED_GALLERY_DIR.iterdir():
            if existing_file.is_file() and existing_file.stem == product_name:
                print(f"[Gallery Info] Reference image for '{product_name}' already exists in gallery. Skipping copy.")
                return True

        src = Path(clean_image_path)
        if not src.exists():
            print(f"[Gallery Error] Clean image not found at: {clean_image_path}")
            return False

        ext = src.suffix.lower()
        if ext not in ['.png', '.jpg', '.jpeg', '.bmp', '.tif', '.tiff']:
            ext = '.png'

        dst = TRAINED_GALLERY_DIR / f"{product_name}{ext}"
        shutil.copy2(src, dst)
        print(f"[Gallery Success] Added reference image to gallery: {dst}")
        return True
    except Exception as e:
        print(f"[Gallery Error] Failed to save image to gallery: {e}")
        return False


# =============================================================================
# Routes - Pages
# =============================================================================

@app.route("/")
def index():
    return render_template("index.html")


# =============================================================================
# Routes - DefectDiffu Config
# =============================================================================

@app.route("/api/defectdiffu/config", methods=["GET"])
def get_defectdiffu_config():
    """Get saved DefectDiffu model paths."""
    try:
        with open(DEFECTDIFFU_CONFIG_PATH, "r") as f:
            config = json.load(f)
        return jsonify({"success": True, "config": config})
    except Exception as e:
        return jsonify({"success": False, "error": str(e)})


@app.route("/api/defectdiffu/config", methods=["POST"])
def save_defectdiffu_config():
    """Save DefectDiffu model paths."""
    data = request.get_json()
    config = {
        "ckpt_path": data.get("ckpt_path", "").strip(),
        "vae_path": data.get("vae_path", "").strip(),
        "vlm_model": data.get("vlm_model", "gemma3:12b").strip()
    }
    try:
        with open(DEFECTDIFFU_CONFIG_PATH, "w") as f:
            json.dump(config, f, indent=2)
        return jsonify({"success": True, "message": "DefectDiffu config saved."})
    except Exception as e:
        return jsonify({"success": False, "error": str(e)})


# =============================================================================
# Routes - Training Data Upload & Mask (DefectFill)
# =============================================================================

@app.route("/api/training/good-images/count", methods=["GET"])
def get_good_image_count():
    """Get count of existing good images for a product."""
    object_class = request.args.get("object_class", "").strip()
    if not object_class:
        return jsonify({"success": False, "error": "Object class required."})

    good_dir = TRAINING_DATA_DIR / object_class / "test" / "good"
    count = 0
    if good_dir.exists():
        count = len([f for f in good_dir.iterdir() if f.is_file() and allowed_file(f.name)])

    return jsonify({
        "success": True,
        "count": count,
        "path": str(good_dir.relative_to(DEFECTFILL_ROOT))
    })


@app.route("/api/training/upload-good-images", methods=["POST"])
def upload_good_images():
    """Upload good (non-defective) images to temp staging area."""
    object_class = request.form.get("object_class", "").strip()
    if not object_class:
        return jsonify({"success": False, "error": "Product name is required."})

    if "images" not in request.files:
        return jsonify({"success": False, "error": "No images provided."})

    files = request.files.getlist("images")
    target_dir = TEMP_IMAGE_DIR / object_class / "test" / "good"
    if target_dir.exists():
        shutil.rmtree(target_dir)
    target_dir.mkdir(parents=True, exist_ok=True)

    saved = []
    for file in files:
        if file and allowed_file(file.filename):
            safe_name = Path(file.filename).name
            save_path = target_dir / safe_name
            counter = 1
            while save_path.exists():
                stem = Path(file.filename).stem
                suffix = Path(file.filename).suffix
                save_path = target_dir / f"{stem}_{counter:02d}{suffix}"
                counter += 1
            file.save(save_path)

            img = cv2.imread(str(save_path))
            if img is not None:
                resized = cv2.resize(img, (512, 512), interpolation=cv2.INTER_LANCZOS4)
                cv2.imwrite(str(save_path), resized)

            rel = save_path.relative_to(TEMP_IMAGE_DIR).as_posix()
            saved.append({
                "filename": save_path.name,
                "path": rel,
                "url": f"/api/temp/image/{rel}"
            })

    return jsonify({
        "success": True,
        "saved": saved,
        "count": len(saved),
        "total_in_folder": len([f for f in target_dir.iterdir() if f.is_file() and allowed_file(f.name)]),
        "target_dir": str(target_dir.relative_to(DEFECTFILL_ROOT).as_posix())
    })


@app.route("/api/training/good-image/<path:subpath>")
def serve_good_image(subpath):
    """Serve uploaded good images."""
    return send_from_directory(TRAINING_DATA_DIR, subpath)


@app.route("/api/training/upload-images", methods=["POST"])
def upload_training_images():
    """Upload defect images to temp staging area."""
    object_class = request.form.get("object_class", "").strip()
    defect_type = request.form.get("defect_type", "").strip()

    if not object_class or not defect_type:
        return jsonify({"success": False, "error": "Product name and defect type are required."})

    if "images" not in request.files:
        return jsonify({"success": False, "error": "No images provided."})

    files = request.files.getlist("images")
    if len(files) < 1:
        return jsonify({"success": False, "error": f"At least 1 defect images required. Received {len(files)}."})

    target_dir = TEMP_IMAGE_DIR / object_class / "train" / "defective" / defect_type
    if target_dir.exists():
        shutil.rmtree(target_dir)
    target_dir.mkdir(parents=True, exist_ok=True)

    saved = []
    for file in files:
        if file and allowed_file(file.filename):
            filename = file.filename
            safe_name = Path(filename).name
            save_path = target_dir / safe_name
            counter = 1
            while save_path.exists():
                stem = Path(filename).stem
                suffix = Path(filename).suffix
                save_path = target_dir / f"{stem}_{counter:02d}{suffix}"
                counter += 1
            file.save(save_path)
            rel = save_path.relative_to(TEMP_IMAGE_DIR).as_posix()
            saved.append({
                "filename": save_path.name,
                "path": rel,
                "url": f"/api/temp/image/{rel}"
            })

    return jsonify({
        "success": True,
        "saved": saved,
        "count": len(saved),
        "target_dir": str(target_dir.relative_to(DEFECTFILL_ROOT).as_posix())
    })


@app.route("/api/training/image/<path:subpath>")
def serve_training_image(subpath):
    """Serve uploaded training images."""
    return send_from_directory(TRAINING_DATA_DIR, subpath)


@app.route("/api/temp/image/<path:subpath>")
def serve_temp_image(subpath):
    """Serve images from the temp staging directory."""
    return send_from_directory(TEMP_IMAGE_DIR, subpath)


@app.route("/api/training/save-mask", methods=["POST"])
def save_mask():
    """Save user-drawn rectangle mask for a training image."""
    data = request.get_json()
    image_path = data.get("image_path", "")
    rect = data.get("rect", {})

    if not image_path or not rect:
        return jsonify({"success": False, "error": "Missing image_path or rect data."})

    full_path = TRAINING_DATA_DIR / image_path
    if not full_path.exists():
        return jsonify({"success": False, "error": "Image not found."})

    rel_path = Path(image_path)
    new_parts = []
    for part in rel_path.parts:
        if part == "defective":
            new_parts.append("defective_masks")
        else:
            new_parts.append(part)

    mask_path = TRAINING_DATA_DIR / Path(*new_parts)
    mask_path.parent.mkdir(parents=True, exist_ok=True)

    coords = generate_mask_from_rect(full_path, rect, mask_path)

    return jsonify({
        "success": True,
        "mask_path": str(mask_path.relative_to(TRAINING_DATA_DIR)),
        "coords": coords
    })


@app.route("/api/training/align", methods=["POST"])
def align_images():
    """Align all temp images for a product to a reference image, then move to final."""
    if not CV2_AVAILABLE:
        return jsonify({"success": False, "error": "OpenCV not installed. Cannot align images."})

    object_class = request.form.get("object_class", "").strip()
    reference = request.files.get("reference")

    if not object_class or not reference:
        return jsonify({"success": False, "error": "Object class and reference image are required."})

    ref_dir = TEMP_IMAGE_DIR / object_class / "reference"
    ref_dir.mkdir(parents=True, exist_ok=True)
    ref_path = ref_dir / "reference.png"
    reference.save(ref_path)

    obj_temp_dir = TEMP_IMAGE_DIR / object_class
    aligned = []
    errors = []

    for img_path in sorted(obj_temp_dir.rglob("*")):
        if not img_path.is_file():
            continue
        if "reference" in str(img_path.relative_to(obj_temp_dir)).split(os.sep):
            continue
        if not allowed_file(img_path.name):
            continue

        rel = img_path.relative_to(obj_temp_dir)
        final_path = TRAINING_DATA_DIR / object_class / rel
        final_path.parent.mkdir(parents=True, exist_ok=True)

        try:
            align_image_to_reference(str(ref_path), str(img_path), str(final_path))
            img_arr = cv2.imread(str(final_path))
            if img_arr is not None:
                img_arr_resized = cv2.resize(img_arr, (512, 512), interpolation=cv2.INTER_LANCZOS4)
                cv2.imwrite(str(final_path), img_arr_resized)

            aligned.append({
                "filename": final_path.name,
                "path": str(final_path.relative_to(TRAINING_DATA_DIR).as_posix()),
                "url": f"/api/training/image/{final_path.relative_to(TRAINING_DATA_DIR).as_posix()}"
            })
        except Exception as e:
            errors.append({"file": str(rel), "error": str(e)})
            shutil.copy2(str(img_path), str(final_path))
            img_arr = cv2.imread(str(final_path))
            if img_arr is not None:
                img_resized = cv2.resize(img_arr, (512, 512), interpolation=cv2.INTER_LANCZOS4)
                cv2.imwrite(str(final_path), img_resized)

            aligned.append({
                "filename": final_path.name,
                "path": str(final_path.relative_to(TRAINING_DATA_DIR).as_posix()),
                "url": f"/api/training/image/{final_path.relative_to(TRAINING_DATA_DIR).as_posix()}"
            })

    return jsonify({
        "success": True,
        "aligned_count": len(aligned),
        "error_count": len(errors),
        "errors": errors,
        "images": aligned
    })


@app.route("/api/training/skip-align", methods=["POST"])
def skip_align():
    """Copy all temp images directly to final location without alignment."""
    data = request.get_json()
    object_class = data.get("object_class", "").strip()

    if not object_class:
        return jsonify({"success": False, "error": "Object class required."})

    obj_temp_dir = TEMP_IMAGE_DIR / object_class
    moved = []

    if obj_temp_dir.exists():
        for img_path in sorted(obj_temp_dir.rglob("*")):
            if not img_path.is_file():
                continue
            if "reference" in str(img_path.relative_to(obj_temp_dir)).split(os.sep):
                continue

            rel = img_path.relative_to(obj_temp_dir)
            final_path = TRAINING_DATA_DIR / object_class / rel
            final_path.parent.mkdir(parents=True, exist_ok=True)
            shutil.copy2(str(img_path), str(final_path))

            img_arr = cv2.imread(str(final_path))
            if img_arr is not None:
                img_resized = cv2.resize(img_arr, (512, 512), interpolation=cv2.INTER_LANCZOS4)
                cv2.imwrite(str(final_path), img_resized)

            moved.append({
                "filename": final_path.name,
                "path": str(final_path.relative_to(TRAINING_DATA_DIR).as_posix()),
                "url": f"/api/training/image/{final_path.relative_to(TRAINING_DATA_DIR).as_posix()}"
            })

    return jsonify({"success": True, "moved_count": len(moved), "images": moved})


# =============================================================================
# Routes - Domain Instructions
# =============================================================================

@app.route("/api/domain-instructions", methods=["GET"])
def get_domain_instructions():
    """Get saved domain instructions."""
    try:
        with open(DOMAIN_CONFIG_PATH, "r") as f:
            config = json.load(f)
        return jsonify({"success": True, "config": config})
    except Exception as e:
        return jsonify({"success": False, "error": str(e)})


@app.route("/api/domain-instructions", methods=["POST"])
def save_domain_instructions():
    """Save domain-specific instructions for prompts."""
    data = request.get_json()
    object_class = data.get("object_class", "default").strip()
    instructions = data.get("instructions", "").strip()

    try:
        with open(DOMAIN_CONFIG_PATH, "r") as f:
            config = json.load(f)

        if "active" not in config:
            config["active"] = {}

        # ==========================================================
        # NEW APPEND LOGIC STARTS HERE
        # ==========================================================
        existing_instructions = config["active"].get(object_class, "")

        if instructions.strip():
            # If the user wrote something in the textarea, append it to the old instructions
            if existing_instructions.strip():
                # Use double newlines to clearly separate the old and new contexts
                config["active"][object_class] = existing_instructions.strip() + "\n\n" + instructions.strip()
            else:
                # If there were no old instructions, just use the new ones
                config["active"][object_class] = instructions.strip()
        else:
            # If the user left the textarea EMPTY, we do NOT erase old instructions.
            # We simply keep the existing instructions untouched.
            pass
        # ==========================================================
        # NEW APPEND LOGIC ENDS HERE
        # ==========================================================

        with open(DOMAIN_CONFIG_PATH, "w") as f:
            json.dump(config, f, indent=2)

        return jsonify({"success": True, "message": f"Instructions saved for '{object_class}'."})
    except Exception as e:
        return jsonify({"success": False, "error": str(e)})


# =============================================================================
# Routes - Training Execution (DefectFill)
# =============================================================================

@app.route("/api/check-product-exists", methods=["GET"])
def check_product_exists():
    """Check if a product (object_class) already exists in training data."""
    object_class = request.args.get("object_class", "").strip()
    if not object_class:
        return jsonify({"success": False, "error": "Product name is required."})

    product_dir = TRAINING_DATA_DIR / object_class
    exists = product_dir.exists() and product_dir.is_dir()

    existing_defect_types = []
    sample_image_url = None

    if exists:
        defective_dir = product_dir / "defective"
        if defective_dir.exists():
            for defect_dir in sorted(defective_dir.iterdir()):
                if defect_dir.is_dir():
                    existing_defect_types.append(defect_dir.name)
                    if not sample_image_url:
                        for img_file in sorted(defect_dir.iterdir()):
                            if allowed_file(img_file.name):
                                rel_path = img_file.relative_to(BASE_DIR)
                                sample_image_url = f"/api/training/image/{rel_path}"
                                break

    return jsonify({
        "success": True,
        "exists": exists,
        "existing_defect_types": existing_defect_types,
        "sample_image_url": sample_image_url
    })


@app.route("/api/training/start", methods=["POST"])
def start_training():
    """Start few-shot training with train.py."""
    global training_state

    data = request.get_json()
    object_class = data.get("object_class", "").strip()
    defect_type = data.get("defect_type", "").strip()
    lora_rank = data.get("lora_rank", 8)
    lora_alpha = data.get("lora_alpha", 16)
    max_steps = data.get("max_train_steps", 1500)
    batch_size = data.get("batch_size", 2)
    gradient_accum = data.get("gradient_accumulation_steps", 2)
    lambda_defect = data.get("lambda_defect", 0.5)
    lambda_obj = data.get("lambda_obj", 0.2)
    lambda_attn = data.get("lambda_attn", 0.05)
    alpha = data.get("alpha", 0.3)

    if not object_class or not defect_type:
        return jsonify({"success": False, "error": "Object class and defect type are required."})

    with training_lock:
        if training_state["running"]:
            return jsonify({"success": False, "error": "Training is already in progress."})

    output_dir = CHECKPOINT_DIR / object_class / defect_type
    output_dir.mkdir(parents=True, exist_ok=True)

    train_script = TRAIN_PY
    if not train_script.exists():
        train_script = BASE_DIR.parent / "train.py"

    cmd = [
        sys.executable, '-u', str(train_script),
        "--data_dir", str(TRAINING_DATA_DIR),
        "--object_class", object_class,
        "--defect_type", defect_type,
        "--output_dir", str(output_dir),
        "--lora_rank", str(lora_rank),
        "--lora_alpha", str(lora_alpha),
        "--max_train_steps", str(max_steps),
        "--batch_size", str(batch_size),
        "--gradient_accumulation_steps", str(gradient_accum),
        "--lambda_defect", str(lambda_defect),
        "--lambda_obj", str(lambda_obj),
        "--lambda_attn", str(lambda_attn),
        "--alpha", str(alpha),
        "--save_steps", str(max(500, max_steps // 3)),
        "--lr_warmup_steps", str(min(100, max_steps // 10)),
        "--dilate_mask", "False",
        "--seed", str(int(time.time() * 1000) % (2**31))
    ]

    try:
        with training_lock:
            training_state["running"] = True
            training_state["object_class"] = object_class
            training_state["defect_type"] = defect_type
            training_state["output_dir"] = str(output_dir)
            training_state["log_lines"] = []
            training_state["start_time"] = datetime.now().isoformat()
            training_state["completed"] = False
            training_state["error"] = None

        process = subprocess.Popen(
            cmd,
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT,
            bufsize=1,
            universal_newlines=True,
            cwd=str(DEFECTFILL_ROOT)
        )

        with training_lock:
            training_state["process"] = process

        log_thread = threading.Thread(target=stream_training_logs, args=(process,))
        log_thread.daemon = True
        log_thread.start()

        return jsonify({
            "success": True,
            "message": "Training started.",
            "command": " ".join(cmd),
            "output_dir": str(output_dir.relative_to(DEFECTFILL_ROOT))
        })

    except Exception as e:
        with training_lock:
            training_state["running"] = False
            training_state["error"] = str(e)
        return jsonify({"success": False, "error": str(e)})


@app.route("/api/training/status", methods=["GET"])
def training_status():
    """Get current training status and logs."""
    with training_lock:
        state = dict(training_state)
        state["process"] = None
    return jsonify({"success": True, "status": state})


@app.route("/api/training/stop", methods=["POST"])
def stop_training():
    """Stop running training."""
    global training_state
    with training_lock:
        if training_state["process"] and training_state["process"].poll() is None:
            training_state["process"].terminate()
            training_state["running"] = False
            training_state["error"] = "Training stopped by user."
            return jsonify({"success": True, "message": "Training terminated."})
    return jsonify({"success": False, "error": "No training is running."})


# =============================================================================
# Routes - Checkpoints & Generation (DefectFill)
# =============================================================================

@app.route("/api/detect-similar-product", methods=["POST"])
def detect_similar_product():
    """Endpoint to identify if a clean product image has already been trained."""
    if "image" not in request.files:
        return jsonify({"success": False, "error": "No clean product image uploaded."})

    file = request.files["image"]
    temp_path = UPLOAD_FOLDER / f"temp_detect_{uuid.uuid4().hex}{Path(file.filename).suffix}"

    try:
        file.save(temp_path)
        threshold = float(request.form.get("threshold", 0.85))
        result = check_product_training_status(
            target_image_path=str(temp_path),
            gallery_dir=TRAINED_GALLERY_DIR,
            similarity_threshold=threshold
        )
        return jsonify({"success": True, "result": result})
    except Exception as e:
        return jsonify({"success": False, "error": str(e)})
    finally:
        if temp_path.exists():
            os.remove(temp_path)


@app.route("/api/checkpoints", methods=["GET"])
def list_checkpoints():
    """List available trained checkpoints."""
    return jsonify({"success": True, "checkpoints": get_checkpoints()})


@app.route("/api/generate", methods=["POST"])
def generate_defect():
    """Start DefectFill generation in background and return a job_id for polling."""
    object_class = request.form.get("object_class", "").strip()
    product_desc = request.form.get("product_desc", "").strip()
    detection_mode = request.form.get("detection_mode", "dots")

    if not object_class or not product_desc:
        return jsonify({"success": False, "error": "Object class and product description are required."})

    defect_type = request.form.get("defect_type", "").strip() or None
    num_defects = int(request.form.get("num_defects", 3))
    guidance_scale = float(request.form.get("guidance_scale", 8.0))
    num_steps = int(request.form.get("num_steps", 50))
    image_size = int(request.form.get("image_size", 512))
    device = request.form.get("device", "cuda")

    clean_files = request.files.getlist("clean_images")
    if not clean_files:
        return jsonify({"success": False, "error": "Clean images are required."})

    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    allowed_exts = {'.png', '.jpg', '.jpeg', '.bmp', '.tif', '.tiff', '.webp'}
    clean_paths = []

    for f in clean_files:
        if not f or not f.filename:
            continue
        ext = os.path.splitext(f.filename)[1].lower()
        if ext not in allowed_exts:
            continue
        safe_name = os.path.basename(f.filename)
        clean_filename = f"clean_{timestamp}_{safe_name}"
        clean_path = UPLOAD_FOLDER / clean_filename
        f.save(clean_path)
        clean_paths.append(str(clean_path))

    if not clean_paths:
        return jsonify({"success": False, "error": "No valid image files found in selected folder."})

    ckpt_dir = CHECKPOINT_DIR
    valid_types = get_checkpoints().get(object_class, [])

    if defect_type and defect_type not in valid_types:
        return jsonify({
            "success": False,
            "error": f"Defect type '{defect_type}' not found for '{object_class}'. Available: {valid_types}"
        })

    valid_object_classes = list(get_checkpoints().keys())
    valid_defect_types = get_checkpoints()

    gen_output_dir = OUTPUT_DIR / f"gen_{timestamp}"
    gen_output_dir.mkdir(parents=True, exist_ok=True)

    try:
        domain_hint = ""
        if DOMAIN_CONFIG_PATH.exists():
            with open(DOMAIN_CONFIG_PATH, "r") as f:
                dconfig = json.load(f)
            domain_hint = dconfig.get("active", {}).get(object_class, "")

        # If detection_mode == "auto", do nothing - let the VLM infer entirely from the product description and domain context.
        if detection_mode == "lines":
            product_desc += " Focus defect placement on metallic pins, leads, or legs."
        elif detection_mode == "dots":
            product_desc += " Focus defect placement on solder ball array, BGA grid, or rounded pads."
        elif detection_mode == "single_rounded":
            product_desc += " Focus defect placement on a single, isolated rounded feature such as a mounting hole, circular pad, individual via, or isolated dot."

        # if domain_hint:
        #     product_desc += f" Domain context: {domain_hint}"

        job_params = {
            'device': device,
            'gen_output_dir': gen_output_dir,
            'ckpt_dir': ckpt_dir,
            'object_class': object_class,
            'defect_type': defect_type,
            'valid_object_classes': valid_object_classes,
            'valid_defect_types': valid_defect_types,
            'image_size': image_size,
            'num_steps': num_steps,
            'guidance_scale': guidance_scale,
            'product_desc': product_desc,
            'domain_hint': domain_hint,
            'clean_paths': clean_paths,
            'num_defects': num_defects,
        }

        job_id = str(uuid.uuid4())
        with jobs_lock:
            generation_jobs[job_id] = {
                'status': 'starting',
                'step': 0,
                'total_steps': 5,
                'step_text': 'Initializing...',
                'defect_current': 0,
                'defect_total': 0,
                'image_current': 0,
                'image_total': 0,
                'result': None,
                'error': None,
                'has_error_trace': False
            }

        thread = threading.Thread(target=run_generation_job, args=(job_id, job_params, 'defectfill'), daemon=True)
        thread.start()

        return jsonify({'success': True, 'job_id': job_id})

    except Exception as e:
        traceback.print_exc()
        return jsonify({"success": False, "error": str(e)})


# =============================================================================
# Routes - DefectDiffu Generation
# =============================================================================

@app.route("/api/defectdiffu/generate", methods=["POST"])
def generate_defectdiffu():
    """Start DefectDiffu generation in background and return a job_id for polling."""
    product_desc = request.form.get("product_desc", "").strip()
    defect_type = request.form.get("defect_type", "").strip() or None
    num_defects = int(request.form.get("num_defects", 3))
    image_size = int(request.form.get("image_size", 512))
    num_steps = int(request.form.get("num_steps", 50))
    device = request.form.get("device", "cuda")

    if not product_desc:
        return jsonify({"success": False, "error": "Product description is required."})

    # Load model config
    try:
        with open(DEFECTDIFFU_CONFIG_PATH, "r") as f:
            dd_config = json.load(f)
    except Exception as e:
        return jsonify({"success": False, "error": f"Failed to load DefectDiffu config: {e}"})

    ckpt_path = dd_config.get("ckpt_path", "")
    vae_path = dd_config.get("vae_path", "")
    vlm_model = dd_config.get("vlm_model", "gemma3:12b")

    if not ckpt_path or not vae_path:
        return jsonify({"success": False, "error": "DefectDiffu model paths not configured. Please set them in the UI first."})

    if not Path(ckpt_path).exists():
        return jsonify({"success": False, "error": f"Checkpoint not found: {ckpt_path}"})
    if not Path(vae_path).exists():
        return jsonify({"success": False, "error": f"VAE not found: {vae_path}"})

    clean_files = request.files.getlist("clean_images")
    if not clean_files:
        return jsonify({"success": False, "error": "Clean images are required."})

    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    allowed_exts = {'.png', '.jpg', '.jpeg', '.bmp', '.tif', '.tiff', '.webp'}
    clean_paths = []

    for f in clean_files:
        if not f or not f.filename:
            continue
        ext = os.path.splitext(f.filename)[1].lower()
        if ext not in allowed_exts:
            continue
        safe_name = os.path.basename(f.filename)
        clean_filename = f"clean_dd_{timestamp}_{safe_name}"
        clean_path = UPLOAD_FOLDER / clean_filename
        f.save(clean_path)
        clean_paths.append(str(clean_path))

    if not clean_paths:
        return jsonify({"success": False, "error": "No valid image files found."})

    gen_output_dir = OUTPUT_DIR / f"defectdiffu_{timestamp}"
    gen_output_dir.mkdir(parents=True, exist_ok=True)

    try:
        job_params = {
            'device': device,
            'gen_output_dir': gen_output_dir,
            'defectdiffu_ckpt': ckpt_path,
            'vae_path': vae_path,
            'vlm_model': vlm_model,
            'image_size': image_size,
            'num_steps': num_steps,
            'product_desc': product_desc,
            'clean_paths': clean_paths,
            'num_defects': num_defects,
            'defect_type': defect_type,
        }

        job_id = str(uuid.uuid4())
        with jobs_lock:
            generation_jobs[job_id] = {
                'status': 'starting',
                'step': 0,
                'total_steps': 5,
                'step_text': 'Initializing DefectDiffu...',
                'defect_current': 0,
                'defect_total': 0,
                'image_current': 0,
                'image_total': 0,
                'result': None,
                'error': None,
                'has_error_trace': False
            }

        thread = threading.Thread(target=run_generation_job, args=(job_id, job_params, 'defectdiffu'), daemon=True)
        thread.start()

        return jsonify({'success': True, 'job_id': job_id})

    except Exception as e:
        traceback.print_exc()
        return jsonify({"success": False, "error": str(e)})


@app.route('/api/generation-status/<job_id>', methods=['GET'])
def generation_status(job_id):
    with jobs_lock:
        job = generation_jobs.get(job_id)
    if not job:
        return jsonify({'success': False, 'error': 'Job not found'}), 404
    return jsonify({'success': True, 'job': job})


@app.route("/api/output/image/<path:subpath>")
def serve_output_image(subpath):
    """Serve generated output images."""
    return send_from_directory(BASE_DIR, subpath)


# =============================================================================
# Routes - System Info
# =============================================================================

@app.route("/api/system/info", methods=["GET"])
def system_info():
    """Get system information."""
    import torch
    info = {
        "cuda_available": torch.cuda.is_available(),
        "cuda_device_count": torch.cuda.device_count() if torch.cuda.is_available() else 0,
        "cuda_device_name": torch.cuda.get_device_name(0) if torch.cuda.is_available() else None,
        "project_root": str(PROJECT_ROOT),
        "defectfill_root": str(DEFECTFILL_ROOT),
        "defectdiffu_root": str(DEFECTDIFFU_ROOT),
        "checkpoint_dir": str(CHECKPOINT_DIR),
        "training_data_dir": str(TRAINING_DATA_DIR)
    }
    return jsonify({"success": True, "info": info})


# =============================================================================
# Main
# =============================================================================

if __name__ == "__main__":
    print("=" * 60)
    print(" DefectFill & DefectDiffu Unified Web UI")
    print("=" * 60)
    print(f" Project root:     {PROJECT_ROOT}")
    print(f" DefectFill root:  {DEFECTFILL_ROOT}")
    print(f" DefectDiffu root: {DEFECTDIFFU_ROOT}")
    print(f" Checkpoints:      {CHECKPOINT_DIR}")
    print(f" Training data:    {TRAINING_DATA_DIR}")
    print(f" Outputs:          {OUTPUT_DIR}")
    print("-" * 60)
    print(" Open http://127.0.0.1:5000 in your browser")
    print("=" * 60)

    app.run(host="0.0.0.0", port=7860, debug=False, threaded=True)