#!/usr/bin/env python3
"""
Local GPU VLM Worker for OCR Provenance MCP System

Uses Chandra OCR 2 (v0.2.0) for 100% local GPU processing.
NO cloud API calls, NO fallbacks.

Modes:
  - Single image (CLI args): --image_path, --prompt, --method, --max_tokens
  - Batch mode (JSON on stdin): {"images": [...], "method": "hf", "max_tokens": 2048}
  - Classify mode: --image_path --classify (quick classification)

Output: JSON to stdout matching VLMAnalysisResult interface.
ALL logging to stderr (stdout is JSON-RPC protocol).
"""

import argparse
import json
import logging
import re
import signal
import sys
import time
from pathlib import Path
from typing import Any, Optional

# Configure logging FIRST - ALL to stderr, NEVER stdout
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
    stream=sys.stderr,
)
logger = logging.getLogger(__name__)


# =============================================================================
# ERROR CLASSES
# =============================================================================


class LocalVLMError(Exception):
    """Base local VLM error with category for structured error handling."""

    def __init__(self, message: str, category: str, details: Optional[dict] = None):
        super().__init__(message)
        self.category = category
        self.details = details or {}


class VLMFileError(LocalVLMError):
    """File access or validation errors."""

    def __init__(self, message: str, file_path: str):
        super().__init__(message, "VLM_FILE_ERROR", {"file_path": file_path})


class VLMModelError(LocalVLMError):
    """Model loading or inference errors."""

    def __init__(self, message: str, details: Optional[dict] = None):
        super().__init__(message, "VLM_MODEL_ERROR", details)


class VLMGPUError(LocalVLMError):
    """GPU memory or CUDA errors."""

    def __init__(self, message: str, details: Optional[dict] = None):
        super().__init__(message, "VLM_GPU_ERROR", details)


class VLMParseError(LocalVLMError):
    """JSON parsing errors from model output."""

    def __init__(self, message: str, raw_output: str = ""):
        super().__init__(
            message, "VLM_PARSE_ERROR", {"raw_output_preview": raw_output[:500]}
        )


# =============================================================================
# CUDA STARTUP VALIDATION (runs on module load)
# =============================================================================

try:
    from gpu_utils import clear_gpu_memory, startup_gpu
except ImportError:
    from .gpu_utils import clear_gpu_memory, startup_gpu

# Device startup: configure memory pool, validate environment, log device info
startup_gpu("VLM worker")


# =============================================================================
# PROMPTS
# =============================================================================

ANALYSIS_PROMPT = """Analyze this image from a document. Return ONLY a JSON object with these exact fields:
{
  "imageType": "photograph|diagram|chart|table|form|signature|stamp|map|medical|screenshot|other",
  "primarySubject": "brief description of main subject",
  "paragraph1": "4-5 sentences identifying the image and its context",
  "paragraph2": "6-8 sentences describing detailed content",
  "paragraph3": "4-5 sentences about legal or medical relevance if applicable",
  "extractedText": ["list of visible text strings"],
  "dates": ["dates found in format YYYY-MM-DD"],
  "names": ["people and organization names"],
  "numbers": ["amounts, reference numbers, IDs"],
  "confidence": 0.85
}
Return ONLY valid JSON, no other text."""

CLASSIFY_PROMPT = """Classify this image. Return ONLY a JSON object:
{"type": "photograph|diagram|chart|table|form|signature|other", "hasText": true, "textDensity": "none|sparse|dense", "complexity": "simple|moderate|complex", "confidence": 0.85}
Return ONLY valid JSON, no other text."""

# Required fields for validation
REQUIRED_ANALYSIS_FIELDS = {"imageType", "primarySubject", "paragraph1", "confidence"}
REQUIRED_CLASSIFY_FIELDS = {"type", "confidence"}

# Valid image types
VALID_IMAGE_TYPES = frozenset(
    {
        "photograph",
        "diagram",
        "chart",
        "table",
        "form",
        "signature",
        "stamp",
        "map",
        "medical",
        "screenshot",
        "other",
    }
)


# =============================================================================
# SINGLETON INFERENCE MANAGER
# =============================================================================

# Cache globally - Chandra loads model into VRAM, must not reload per-call
_inference_manager: Any = None
_inference_method: Optional[str] = None


def get_inference_manager(method: str = "hf") -> Any:
    """
    Get or create the singleton InferenceManager.

    The manager is cached globally because model loading is expensive (~10.6GB VRAM).
    If the method changes, the old manager is discarded and a new one created.
    """
    global _inference_manager, _inference_method

    if _inference_manager is not None and _inference_method == method:
        return _inference_manager

    # If method changed, clear old manager
    if _inference_manager is not None:
        logger.info("Method changed from %s to %s, recreating manager", _inference_method, method)
        _inference_manager = None
        clear_gpu_memory()

    logger.info("Loading Chandra InferenceManager (method=%s)...", method)
    start = time.time()

    try:
        import inspect

        import torch
        from chandra.model import InferenceManager

        # Check if InferenceManager accepts torch_dtype for BF16
        init_params = inspect.signature(InferenceManager.__init__).parameters
        use_bf16 = False
        if "torch_dtype" in init_params and torch.cuda.is_available():
            try:
                use_bf16 = torch.cuda.is_bf16_supported()
            except Exception:
                use_bf16 = False

        if use_bf16:
            logger.info("Chandra InferenceManager supports torch_dtype, using bfloat16")
            _inference_manager = InferenceManager(method=method, torch_dtype=torch.bfloat16)
        else:
            dtype_reason = "FP32" if "torch_dtype" in init_params else "torch_dtype not accepted"
            logger.info("Chandra InferenceManager using %s", dtype_reason)
            _inference_manager = InferenceManager(method=method)

        _inference_method = method
        elapsed = time.time() - start

        # Log post-load device memory state
        if torch.cuda.is_available():
            vram_used = torch.cuda.memory_allocated(0) / (1024**3)
            logger.info("Chandra InferenceManager loaded in %.1fs (VRAM used: %.1f GB)", elapsed, vram_used)
        else:
            logger.info("Chandra InferenceManager loaded in %.1fs (non-CUDA device)", elapsed)
        return _inference_manager
    except (VLMGPUError, VLMModelError):
        raise
    except Exception as e:
        raise VLMModelError(
            f"Failed to load Chandra InferenceManager: {e}",
            {"method": method, "error": str(e)},
        )


# =============================================================================
# JSON PARSING (robust, with fallback)
# =============================================================================


def parse_vlm_json(raw: str, required_fields: set[str]) -> dict[str, Any]:
    """
    Robustly parse JSON from VLM model output.

    Strategy:
      1. Try json.loads on raw output directly
      2. Strip markdown code fences and retry
      3. Search for JSON block with regex
      4. Validate required fields
      5. FALLBACK: if all fails, return minimal structure with confidence 0.3

    Args:
        raw: Raw model output string
        required_fields: Set of field names that must be present

    Returns:
        Parsed dict with at least the required fields
    """
    if not raw or not raw.strip():
        logger.warning("Empty VLM output, returning fallback")
        return _make_fallback(raw, required_fields)

    # Step 1: Try direct parse
    try:
        result = json.loads(raw)
        if isinstance(result, dict) and required_fields.issubset(result.keys()):
            return result
        if isinstance(result, dict):
            logger.warning("Parsed JSON missing fields: %s", required_fields - result.keys())
            return _fill_defaults(result, required_fields)
    except json.JSONDecodeError:
        pass

    # Step 2: Strip markdown fences
    clean = re.sub(r"```(?:json)?\n?", "", raw).strip()
    clean = re.sub(r"\n?```$", "", clean).strip()
    try:
        result = json.loads(clean)
        if isinstance(result, dict):
            return _fill_defaults(result, required_fields)
    except json.JSONDecodeError:
        pass

    # Step 3: Regex search for JSON object
    match = re.search(r"\{[\s\S]*\}", clean)
    if match:
        try:
            result = json.loads(match.group())
            if isinstance(result, dict):
                return _fill_defaults(result, required_fields)
        except json.JSONDecodeError:
            logger.warning("Regex-extracted JSON failed to parse")

    # Step 4: All attempts failed - fallback
    logger.warning(
        "All JSON parse attempts failed (raw length=%d, preview=%s)",
        len(raw),
        raw[:200],
    )
    return _make_fallback(raw, required_fields)


def _fill_defaults(parsed: dict[str, Any], required_fields: set[str]) -> dict[str, Any]:
    """Fill missing required fields with defaults."""
    defaults: dict[str, Any] = {
        "imageType": "other",
        "primarySubject": "",
        "paragraph1": "",
        "paragraph2": "",
        "paragraph3": "",
        "extractedText": [],
        "dates": [],
        "names": [],
        "numbers": [],
        "confidence": 0.5,
        "type": "other",
        "hasText": False,
        "textDensity": "unknown",
        "complexity": "moderate",
    }
    for field in required_fields:
        if field not in parsed:
            parsed[field] = defaults.get(field, "")
    return parsed


def _make_fallback(raw: str, required_fields: set[str]) -> dict[str, Any]:
    """Create a fallback result when JSON parsing completely fails."""
    text = (raw or "").strip()
    result: dict[str, Any] = {
        "imageType": "other",
        "primarySubject": text[:200] if text else "Unable to analyze image",
        "paragraph1": text if text else "The model output could not be parsed as structured JSON.",
        "paragraph2": "",
        "paragraph3": "",
        "extractedText": [],
        "dates": [],
        "names": [],
        "numbers": [],
        "confidence": 0.3,
        "type": "other",
        "hasText": False,
        "textDensity": "unknown",
        "complexity": "moderate",
    }
    return _fill_defaults(result, required_fields)


# =============================================================================
# IMAGE VALIDATION
# =============================================================================

SUPPORTED_EXTENSIONS = frozenset({".png", ".jpg", ".jpeg", ".tiff", ".tif", ".bmp", ".webp", ".gif"})


def validate_image_path(image_path: str) -> Path:
    """
    Validate that an image path exists and is a supported format.

    Args:
        image_path: Path to the image file

    Returns:
        Resolved Path object

    Raises:
        VLMFileError: If file doesn't exist or format unsupported
    """
    path = Path(image_path).resolve()

    if not path.exists():
        raise VLMFileError(f"Image file not found: {path}", str(path))

    if not path.is_file():
        raise VLMFileError(f"Path is not a file: {path}", str(path))

    if path.suffix.lower() not in SUPPORTED_EXTENSIONS:
        raise VLMFileError(
            f"Unsupported image format: {path.suffix}. Supported: {sorted(SUPPORTED_EXTENSIONS)}",
            str(path),
        )

    # Check file is not empty
    if path.stat().st_size == 0:
        raise VLMFileError(f"Image file is empty: {path}", str(path))

    return path


# =============================================================================
# SINGLE IMAGE ANALYSIS
# =============================================================================


def analyze_image(
    image_path: str,
    prompt: Optional[str] = None,
    method: str = "hf",
    max_tokens: int = 2048,
) -> dict[str, Any]:
    """
    Analyze a single image using Chandra VLM.

    Args:
        image_path: Path to the image file
        prompt: Custom prompt (defaults to ANALYSIS_PROMPT)
        method: Inference method ("hf" or "vllm")
        max_tokens: Maximum output tokens

    Returns:
        VLMAnalysisResult dict matching the TypeScript interface
    """
    path = validate_image_path(image_path)
    use_prompt = prompt or ANALYSIS_PROMPT
    start_time = time.time()

    logger.info("Analyzing image: %s (method=%s, max_tokens=%d)", path.name, method, max_tokens)

    try:
        from PIL import Image
        from chandra.model.schema import BatchInputItem

        img = Image.open(path)
        img.load()  # Force load to catch corrupted images early
        logger.info("Image loaded: %dx%d, mode=%s", img.width, img.height, img.mode)

    except Exception as e:
        if isinstance(e, LocalVLMError):
            raise
        raise VLMFileError(f"Failed to open image: {e}", str(path))

    try:
        manager = get_inference_manager(method)
        item = BatchInputItem(image=img, prompt=use_prompt)
        results = manager.generate([item], max_output_tokens=max_tokens)

        if not results or len(results) == 0:
            raise VLMModelError("Chandra returned empty results", {"image_path": str(path)})

        result = results[0]

        if result.error:
            raise VLMModelError(
                f"Chandra reported error processing image: {path.name}",
                {"image_path": str(path), "raw": result.raw[:500] if result.raw else ""},
            )

        # Parse the raw output as JSON
        raw_output = result.raw or ""
        analysis = parse_vlm_json(raw_output, REQUIRED_ANALYSIS_FIELDS)

        # Normalize imageType to valid values
        if analysis.get("imageType") not in VALID_IMAGE_TYPES:
            analysis["imageType"] = "other"

        # Ensure list fields are lists
        for field in ("extractedText", "dates", "names", "numbers"):
            if not isinstance(analysis.get(field), list):
                analysis[field] = []

        # Ensure confidence is a float in [0, 1]
        try:
            analysis["confidence"] = max(0.0, min(1.0, float(analysis.get("confidence", 0.5))))
        except (TypeError, ValueError):
            analysis["confidence"] = 0.5

        # Build description from paragraphs
        paragraphs = [
            analysis.get("paragraph1", ""),
            analysis.get("paragraph2", ""),
            analysis.get("paragraph3", ""),
        ]
        description = "\n\n".join(p for p in paragraphs if p)

        elapsed_ms = int((time.time() - start_time) * 1000)
        tokens_used = result.token_count if result.token_count else 0

        vlm_result: dict[str, Any] = {
            "description": description,
            "analysis": {
                "imageType": analysis["imageType"],
                "primarySubject": analysis.get("primarySubject", ""),
                "paragraph1": analysis.get("paragraph1", ""),
                "paragraph2": analysis.get("paragraph2", ""),
                "paragraph3": analysis.get("paragraph3", ""),
                "extractedText": analysis.get("extractedText", []),
                "dates": analysis.get("dates", []),
                "names": analysis.get("names", []),
                "numbers": analysis.get("numbers", []),
                "confidence": analysis["confidence"],
            },
            "model": "chandra",
            "processingTimeMs": elapsed_ms,
            "tokensUsed": tokens_used,
        }

        logger.info(
            "Analysis complete: type=%s, confidence=%.2f, tokens=%d, time=%dms",
            analysis["imageType"],
            analysis["confidence"],
            tokens_used,
            elapsed_ms,
        )

        return vlm_result

    except LocalVLMError:
        raise
    except Exception as e:
        error_str = str(e)
        if "CUDA" in error_str or "out of memory" in error_str.lower():
            clear_gpu_memory()
            raise VLMGPUError(
                f"GPU error during analysis: {e}",
                {"image_path": str(path), "error": error_str},
            )
        raise VLMModelError(
            f"Chandra inference error: {e}",
            {"image_path": str(path), "error": error_str},
        )
    finally:
        img.close()  # F-2: free PIL image memory after inference


# =============================================================================
# CLASSIFY IMAGE
# =============================================================================


def classify_image(
    image_path: str,
    method: str = "hf",
    max_tokens: int = 256,
) -> dict[str, Any]:
    """
    Quick classification of an image using Chandra VLM.

    Args:
        image_path: Path to the image file
        method: Inference method ("hf" or "vllm")
        max_tokens: Maximum output tokens (classification is small)

    Returns:
        Classification dict with type, hasText, textDensity, complexity, confidence
    """
    path = validate_image_path(image_path)
    start_time = time.time()

    logger.info("Classifying image: %s (method=%s)", path.name, method)

    try:
        from PIL import Image
        from chandra.model.schema import BatchInputItem

        img = Image.open(path)
        img.load()

    except Exception as e:
        if isinstance(e, LocalVLMError):
            raise
        raise VLMFileError(f"Failed to open image: {e}", str(path))

    try:
        manager = get_inference_manager(method)
        item = BatchInputItem(image=img, prompt=CLASSIFY_PROMPT)
        results = manager.generate([item], max_output_tokens=max_tokens)

        if not results or len(results) == 0:
            raise VLMModelError("Chandra returned empty results for classification")

        result = results[0]

        if result.error:
            raise VLMModelError(
                f"Chandra reported error classifying image: {path.name}",
                {"image_path": str(path)},
            )

        raw_output = result.raw or ""
        classification = parse_vlm_json(raw_output, REQUIRED_CLASSIFY_FIELDS)

        # Normalize type field
        valid_types = {"photograph", "diagram", "chart", "table", "form", "signature", "other"}
        if classification.get("type") not in valid_types:
            classification["type"] = "other"

        # Normalize boolean
        classification["hasText"] = bool(classification.get("hasText", False))

        # Normalize textDensity
        if classification.get("textDensity") not in ("none", "sparse", "dense"):
            classification["textDensity"] = "unknown"

        # Normalize complexity
        if classification.get("complexity") not in ("simple", "moderate", "complex"):
            classification["complexity"] = "moderate"

        # Normalize confidence
        try:
            classification["confidence"] = max(0.0, min(1.0, float(classification.get("confidence", 0.5))))
        except (TypeError, ValueError):
            classification["confidence"] = 0.5

        elapsed_ms = int((time.time() - start_time) * 1000)

        logger.info(
            "Classification complete: type=%s, hasText=%s, complexity=%s, confidence=%.2f, time=%dms",
            classification["type"],
            classification["hasText"],
            classification["complexity"],
            classification["confidence"],
            elapsed_ms,
        )

        return {
            "type": classification["type"],
            "hasText": classification["hasText"],
            "textDensity": classification["textDensity"],
            "complexity": classification["complexity"],
            "confidence": classification["confidence"],
            "model": "chandra",
            "processingTimeMs": elapsed_ms,
        }

    except LocalVLMError:
        raise
    except Exception as e:
        error_str = str(e)
        if "CUDA" in error_str or "out of memory" in error_str.lower():
            clear_gpu_memory()
            raise VLMGPUError(f"GPU error during classification: {e}")
        raise VLMModelError(f"Chandra classification error: {e}")
    finally:
        img.close()  # F-2: free PIL image memory after inference


# =============================================================================
# BATCH PROCESSING
# =============================================================================


def process_batch(
    images: list[dict[str, Any]],
    method: str = "hf",
    max_tokens: int = 2048,
    batch_size: int = 4,
) -> list[dict[str, Any]]:
    """
    Process multiple images in batches.

    Each item in images must have:
      - image_path: str (required)
      - prompt: str (optional, defaults to ANALYSIS_PROMPT)

    Args:
        images: List of image specs
        method: Inference method
        max_tokens: Maximum output tokens per image
        batch_size: Number of images per GPU batch

    Returns:
        List of VLMAnalysisResult dicts (one per input image)
    """
    if not images:
        return []

    results: list[dict[str, Any]] = []
    total = len(images)

    logger.info("Processing batch of %d images (batch_size=%d, method=%s)", total, batch_size, method)

    # Process in sub-batches for GPU memory management
    for batch_start in range(0, total, batch_size):
        batch_end = min(batch_start + batch_size, total)
        batch_items = images[batch_start:batch_end]
        batch_num = (batch_start // batch_size) + 1
        total_batches = (total + batch_size - 1) // batch_size

        logger.info("Processing sub-batch %d/%d (%d images)", batch_num, total_batches, len(batch_items))

        batch_results = _process_single_batch(batch_items, method, max_tokens)
        results.extend(batch_results)

        # Free GPU memory between sub-batches
        if batch_end < total:
            clear_gpu_memory()

    # Final GPU cleanup
    clear_gpu_memory()

    logger.info("Batch complete: %d/%d images processed", len(results), total)
    return results


def _process_single_batch(
    batch_items: list[dict[str, Any]],
    method: str,
    max_tokens: int,
) -> list[dict[str, Any]]:
    """
    Process a single sub-batch of images through Chandra.

    NOTE: Chandra's generate() with custom prompts returns only 1 result
    regardless of batch size. Therefore we process images individually,
    reusing the cached InferenceManager singleton for efficiency.
    """
    results: list[dict[str, Any]] = []

    for idx, spec in enumerate(batch_items):
        image_path = spec.get("image_path", "")
        prompt = spec.get("prompt")

        try:
            result = analyze_image(
                image_path,
                prompt=prompt,
                method=method,
                max_tokens=max_tokens,
            )
            results.append(result)
            logger.info("Batch item %d/%d complete: %s", idx + 1, len(batch_items), image_path)
        except LocalVLMError as e:
            logger.error("Batch item %d/%d failed: %s: %s", idx + 1, len(batch_items), e.category, e)
            results.append(_make_error_result(image_path, str(e)))
        except Exception as e:
            error_str = str(e)
            if "CUDA" in error_str or "out of memory" in error_str.lower():
                clear_gpu_memory()
                logger.error("GPU OOM on batch item %d/%d: %s", idx + 1, len(batch_items), image_path)
                results.append(_make_error_result(image_path, f"GPU out of memory: {e}"))
            else:
                logger.error("Unexpected error on batch item %d/%d: %s", idx + 1, len(batch_items), e)
                results.append(_make_error_result(image_path, str(e)))

    return results


def _make_error_result(image_path: str, error_msg: str) -> dict[str, Any]:
    """Create an error result entry for a failed image."""
    return {
        "description": "",
        "analysis": {
            "imageType": "other",
            "primarySubject": f"Error: {error_msg}",
            "paragraph1": f"Failed to analyze image: {error_msg}",
            "paragraph2": "",
            "paragraph3": "",
            "extractedText": [],
            "dates": [],
            "names": [],
            "numbers": [],
            "confidence": 0.0,
        },
        "model": "chandra",
        "processingTimeMs": 0,
        "tokensUsed": 0,
        "error": error_msg,
        "image_path": image_path,
    }


# =============================================================================
# CLI INTERFACE
# =============================================================================


def parse_args() -> argparse.Namespace:
    """Parse command-line arguments."""
    parser = argparse.ArgumentParser(
        description="Local GPU VLM Worker using Chandra OCR 2 (v0.2.0)",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  # Single image analysis
  python vlm_worker_local.py --image_path photo.png --method hf

  # Single image with custom prompt
  python vlm_worker_local.py --image_path photo.png --prompt "Describe this chart"

  # Quick classification
  python vlm_worker_local.py --image_path photo.png --classify

  # Batch mode (JSON on stdin)
  echo '{"images": [{"image_path": "a.png"}, {"image_path": "b.png"}]}' | python vlm_worker_local.py --batch
        """,
    )
    # Daemon mode flag
    parser.add_argument("--daemon", action="store_true",
                        help="Run in daemon mode: load model once, accept JSON work items on stdin")

    # CLI mode args
    parser.add_argument("--image_path", type=str, help="Path to image file (single mode)")
    parser.add_argument("--prompt", type=str, help="Custom analysis prompt (overrides default)")
    parser.add_argument("--method", type=str, default="hf", choices=["hf", "vllm"],
                        help="Inference method (default: hf)")
    parser.add_argument("--max_tokens", type=int, default=2048,
                        help="Maximum output tokens (default: 2048)")
    parser.add_argument("--batch_size", type=int, default=4,
                        help="Images per GPU batch (default: 4)")
    parser.add_argument("--classify", action="store_true",
                        help="Quick classification mode instead of full analysis")
    parser.add_argument("--batch", action="store_true",
                        help="Batch mode: read JSON from stdin")
    return parser.parse_args()


def _load_model_for_daemon(method: str = "hf") -> tuple[Any, int]:
    """
    Load the Chandra model once for daemon mode.

    Args:
        method: Inference method ("hf" or "vllm")

    Returns:
        (inference_manager, model_load_ms)

    Raises:
        VLMGPUError, VLMModelError on failure.
    """
    logger.info("Daemon: Loading Chandra model (one-time, method=%s)...", method)
    model_load_start = time.time()

    # This calls get_inference_manager which handles GPU checks and model loading
    manager = get_inference_manager(method)

    model_load_ms = int((time.time() - model_load_start) * 1000)
    logger.info("Daemon: Chandra model loaded in %dms", model_load_ms)

    return manager, model_load_ms


def daemon_loop(manager: Any, model_load_ms: int, default_method: str = "hf") -> None:
    """
    Persistent processing loop - model already loaded in GPU memory.

    Protocol (JSON lines over stdin/stdout):
      Input:  {"action": "analyze", "image_path": "...", "prompt": "optional", "method": "hf", "max_tokens": 2048}
      Output: VLMAnalysisResult as JSON

      Input:  {"action": "classify", "image_path": "...", "method": "hf"}
      Output: Classification result as JSON

      Input:  {"action": "ping"}
      Output: {"status": "ready", "items_processed": N}

      Input:  {"action": "shutdown"}
      Output: (none, clean exit)

    On error: {"error": "...", "category": "...", "details": {...}}

    Reads one JSON line per iteration, processes, writes one JSON line response.
    Exits cleanly on stdin EOF (TypeScript process closed pipe) or shutdown action.
    """
    items_processed = 0

    # Output ready signal so the TypeScript caller knows model is loaded
    _output_json({
        "status": "ready",
        "model_loaded": True,
        "load_time_ms": model_load_ms,
    })

    logger.info("Daemon: Entering main loop, waiting for work items on stdin...")

    for line in sys.stdin:
        line = line.strip()
        if not line:
            continue

        try:
            request = json.loads(line)
        except json.JSONDecodeError as e:
            _output_json({
                "error": f"Invalid JSON: {e}",
                "category": "VLM_USAGE_ERROR",
                "details": {"raw_input": line[:200]},
            })
            continue

        action = request.get("action", "analyze")

        if action == "ping":
            _output_json({
                "status": "ready",
                "items_processed": items_processed,
            })
            continue

        if action == "shutdown":
            logger.info(
                "Daemon: Shutdown requested after %d items processed.", items_processed
            )
            break

        if action == "analyze":
            image_path = request.get("image_path", "")
            prompt = request.get("prompt")
            method = request.get("method", default_method)
            max_tokens = request.get("max_tokens", 2048)

            try:
                result = analyze_image(
                    image_path,
                    prompt=prompt,
                    method=method,
                    max_tokens=max_tokens,
                )
                _output_json(result)
                items_processed += 1
            except LocalVLMError as e:
                logger.error("Daemon VLM error [%s]: %s", e.category, e)
                _output_json({
                    "error": str(e),
                    "category": e.category,
                    "details": e.details,
                })
            except Exception as e:
                logger.exception("Daemon unexpected error: %s", e)
                _output_json({
                    "error": str(e),
                    "category": "VLM_INTERNAL_ERROR",
                    "details": {"type": type(e).__name__},
                })

        elif action == "classify":
            image_path = request.get("image_path", "")
            method = request.get("method", default_method)

            try:
                result = classify_image(
                    image_path,
                    method=method,
                )
                _output_json(result)
                items_processed += 1
            except LocalVLMError as e:
                logger.error("Daemon VLM error [%s]: %s", e.category, e)
                _output_json({
                    "error": str(e),
                    "category": e.category,
                    "details": e.details,
                })
            except Exception as e:
                logger.exception("Daemon unexpected error: %s", e)
                _output_json({
                    "error": str(e),
                    "category": "VLM_INTERNAL_ERROR",
                    "details": {"type": type(e).__name__},
                })

        else:
            _output_json({
                "error": f"Unknown action: {action}",
                "category": "VLM_USAGE_ERROR",
                "details": {"valid_actions": ["analyze", "classify", "ping", "shutdown"]},
            })
            continue

    # Clean exit: stdin EOF or shutdown action
    logger.info(
        "Daemon: Exiting after %d items processed. Freeing GPU memory.", items_processed
    )
    clear_gpu_memory()


def main() -> None:
    """Main entry point."""
    args = parse_args()

    # ---- DAEMON MODE ----
    if args.daemon:
        # Set up signal handlers for clean shutdown
        def _signal_handler(signum: int, _frame: Any) -> None:
            sig_name = signal.Signals(signum).name
            logger.info("Daemon: Received %s, shutting down...", sig_name)
            clear_gpu_memory()
            sys.exit(0)

        signal.signal(signal.SIGTERM, _signal_handler)
        signal.signal(signal.SIGINT, _signal_handler)

        try:
            manager, model_load_ms = _load_model_for_daemon(args.method)
            daemon_loop(manager, model_load_ms, default_method=args.method)
        except LocalVLMError as e:
            logger.error("Daemon startup error [%s]: %s", e.category, e)
            _output_json({
                "error": str(e),
                "category": e.category,
                "details": e.details,
            })
            sys.exit(1)
        except Exception as e:
            logger.exception("Daemon startup unexpected error: %s", e)
            _output_json({
                "error": str(e),
                "category": "VLM_INTERNAL_ERROR",
                "details": {"type": type(e).__name__},
            })
            sys.exit(1)
        finally:
            clear_gpu_memory()
        return

    # ---- CLI MODE ----
    try:
        if args.batch:
            # Batch mode: read JSON from stdin
            logger.info("Reading batch input from stdin...")
            try:
                raw_input = sys.stdin.read()
                if not raw_input.strip():
                    raise VLMFileError("Empty stdin input for batch mode", "stdin")
                batch_spec = json.loads(raw_input)
            except json.JSONDecodeError as e:
                raise VLMParseError(f"Invalid JSON on stdin: {e}", raw_input[:500] if raw_input else "")

            images = batch_spec.get("images", [])
            if not isinstance(images, list) or len(images) == 0:
                raise VLMFileError("Batch input must have non-empty 'images' array", "stdin")

            method = batch_spec.get("method", args.method)
            max_tokens = batch_spec.get("max_tokens", args.max_tokens)
            batch_size = batch_spec.get("batch_size", args.batch_size)

            results = process_batch(images, method=method, max_tokens=max_tokens, batch_size=batch_size)
            _output_json(results)

        elif args.image_path:
            if args.classify:
                # Classification mode
                result = classify_image(args.image_path, method=args.method)
                _output_json(result)
            else:
                # Single image analysis
                result = analyze_image(
                    args.image_path,
                    prompt=args.prompt,
                    method=args.method,
                    max_tokens=args.max_tokens,
                )
                _output_json(result)
        else:
            # No image_path and not batch - check if stdin has data
            if not sys.stdin.isatty():
                logger.info("No --image_path, trying stdin batch mode...")
                raw_input = sys.stdin.read()
                if raw_input.strip():
                    try:
                        batch_spec = json.loads(raw_input)
                        images = batch_spec.get("images", [])
                        if isinstance(images, list) and len(images) > 0:
                            method = batch_spec.get("method", args.method)
                            max_tokens = batch_spec.get("max_tokens", args.max_tokens)
                            batch_size = batch_spec.get("batch_size", args.batch_size)
                            results = process_batch(images, method=method, max_tokens=max_tokens, batch_size=batch_size)
                            _output_json(results)
                            return
                    except json.JSONDecodeError:
                        pass

            print(
                json.dumps({
                    "error": "No input provided. Use --image_path for single mode or --batch for batch mode.",
                    "category": "VLM_USAGE_ERROR",
                }),
                file=sys.stdout,
                flush=True,
            )
            sys.exit(1)

    except LocalVLMError as e:
        logger.error("%s: %s", e.category, e)
        print(
            json.dumps({
                "error": str(e),
                "category": e.category,
                "details": e.details,
            }),
            file=sys.stdout,
            flush=True,
        )
        sys.exit(1)

    except Exception as e:
        logger.error("Unexpected error: %s", e, exc_info=True)
        print(
            json.dumps({
                "error": str(e),
                "category": "VLM_INTERNAL_ERROR",
                "details": {"type": type(e).__name__},
            }),
            file=sys.stdout,
            flush=True,
        )
        sys.exit(1)

    finally:
        clear_gpu_memory()


def _output_json(data: Any) -> None:
    """Write JSON to stdout (the only thing that goes to stdout)."""
    print(json.dumps(data, ensure_ascii=False), file=sys.stdout, flush=True)


if __name__ == "__main__":
    main()
