vision_agent.py

"""
vision_agent.py — Vision Agent (Fallback Only) — OCR Text Extraction Only

Supports 3 backends:
  1. LMStudio (Local) — OpenAI Vision API with base64-encoded images
  2. ZeroGPU (HF Spaces) — Moondream2 via transformers + GPU
  3. Transformers (CPU) — direct Moondream2 loading (slow fallback)

Prompt Strategy:
  Pure OCR — Extract ALL visible text from the banner image.
  No classification or offer parsing is done here.
  The extracted text is written back to the email's body_text field
  so the Classifier Agent (Qwen) can handle classification in Phase 3.

Only invoked when an email's body text is insufficient (< 50 chars)
and the HTML contains banner image URLs. Downloads images from URLs
and uses vision model to extract text from promotional banners.
"""

import gc
import base64
import logging
import requests
import io
from typing import Dict, List, Optional
from PIL import Image

from database import BankingDatabase
from config import (
    BANNER_DOWNLOAD_TIMEOUT, BANNERS_DIR,
    MOONDREAM_HF_MODEL_ID, MOONDREAM_HF_REVISION,
    LMSTUDIO_BASE_URL, LMSTUDIO_VISION_MODEL,
    BACKEND_LMSTUDIO, BACKEND_ZEROGPU, BACKEND_GGUF,
    DEFAULT_BACKEND, MODELS_DIR
)

logger = logging.getLogger(__name__)

# ─── Try importing optional dependencies ─────────────────────────────────────
try:
    import spaces
    HAS_SPACES = True
except ImportError:
    HAS_SPACES = False

try:
    from openai import OpenAI
    HAS_OPENAI = True
except ImportError:
    HAS_OPENAI = False


# ══════════════════════════════════════════════════════════════════════════════
# VISION PROMPT TEMPLATES — Multi-Pass Strategy
# ══════════════════════════════════════════════════════════════════════════════

# OCR Extraction — get ALL text from the image using RECAT Strategy
VISION_OCR_PROMPT = """**Role**: You are a highly accurate Optical Character Recognition (OCR) assistant.
**Expertise**: You specialize in analyzing Indian banking and financial images, including transaction alerts, e-statements, and promotional banners.
**Context**: You are analyzing an image extracted from a banking email. The image may contain critical transaction details (money sent/received) OR promotional offers. 
**Action**: Extract ALL text you can see in the image exactly as written. Pay extremely close attention to:
1. Transaction details: Amounts (Rs./₹/INR), "debited from" / "credited to", merchant names, and transaction IDs/references.
2. Account details: Last 4 digits of cards or accounts.
3. Promotional details: Percentages (%), flat discounts, promo codes.
4. Dates and timestamps.
**Tone**: Be purely objective, accurate, and precise. Do not summarize or interpret; simply transcribe the text clearly and preserve the visual hierarchy."""

# LMStudio vision system prompt (OCR-focused)
VISION_SYSTEM_PROMPT = """You are a highly accurate OCR assistant specializing in Indian banking.
Extract all visible text accurately without adding commentary. Pay extremely close attention to transaction amounts, currency symbols (Rs, INR, ₹), merchant names, card numbers, percentages, and dates."""


class VisionAgent:
    """
    Vision agent for extracting text from promotional banners (OCR only).
    
    FALLBACK ONLY: Only used when email body text lacks useful information
    and the email HTML contains image URLs to promotional banners.
    
    This agent ONLY extracts raw text from images. It does NOT classify
    or categorize the content. The extracted text is appended to the
    email's body_text field so the Classifier Agent (Qwen) can process it.
    """

    def __init__(self):
        self.model = None
        self.tokenizer = None
        self._loaded = False
        self._backend = DEFAULT_BACKEND
        self._openai_client = None
        self._resolved_vision_model = LMSTUDIO_VISION_MODEL  # Will be updated to actual LMStudio model ID

    # ─── Backend Management ──────────────────────────────────────────────────

    def set_backend(self, backend_name: str):
        """Switch inference backend. Unloads any previously loaded model."""
        if backend_name == self._backend and self._loaded:
            return
        self.unload_model()
        self._backend = backend_name
        logger.info(f"Vision backend set to: {backend_name}")

    def get_backend(self) -> str:
        return self._backend

    # ─── Model Loading ───────────────────────────────────────────────────────

    def load_model(self, progress_callback=None):
        """Load vision model based on current backend."""
        if self._loaded:
            return True

        def log(msg):
            logger.info(msg)
            if progress_callback:
                progress_callback(msg)

        if self._backend == BACKEND_LMSTUDIO:
            return self._load_lmstudio(log)
        elif self._backend == BACKEND_ZEROGPU:
            return self._load_zerogpu(log)
        else:
            return self._load_llama_cpp(log)

    def _load_lmstudio(self, log) -> bool:
        """Initialize OpenAI client for LMStudio vision model."""
        if not HAS_OPENAI:
            log("[ERROR] `openai` package not installed for LMStudio vision")
            return False
        try:
            self._openai_client = OpenAI(
                base_url=LMSTUDIO_BASE_URL,
                api_key="lm-studio",
            )
            models = self._openai_client.models.list()
            model_ids = [m.id for m in models.data]
            
            # Check if vision model is loaded in LMStudio
            # Find the actual full model ID that contains our search term
            matched_model = None
            for m in model_ids:
                if LMSTUDIO_VISION_MODEL.lower() in m.lower():
                    matched_model = m
                    break
            
            if not matched_model and model_ids:
                log(f"[WARN] Vision model containing '{LMSTUDIO_VISION_MODEL}' not loaded in LMStudio.")
                log(f"       Loaded models: {', '.join(model_ids)}")
                log("       Please load the vision model in LMStudio to enable banner extraction.")
                self._loaded = False
                return False

            # Store the resolved full model ID for API calls
            self._resolved_vision_model = matched_model
            self._loaded = True
            log(f"[OK] Vision: Connected to LMStudio")
            log(f"  Model: {self._resolved_vision_model}")
            return True
        except Exception as e:
            log(f"[ERROR] LMStudio vision connection failed: {e}")
            self._loaded = False
            return False

    def _load_zerogpu(self, log) -> bool:
        """Load Moondream2 for ZeroGPU."""
        log("Loading Moondream2 vision model (ZeroGPU)...")
        try:
            from transformers import AutoModelForCausalLM, AutoTokenizer
            import torch

            log(f"  Model: {MOONDREAM_HF_MODEL_ID} (rev: {MOONDREAM_HF_REVISION})")

            self.model = AutoModelForCausalLM.from_pretrained(
                MOONDREAM_HF_MODEL_ID,
                revision=MOONDREAM_HF_REVISION,
                trust_remote_code=True,
                torch_dtype=torch.float16,
                device_map="auto",
            )
            self.tokenizer = AutoTokenizer.from_pretrained(
                MOONDREAM_HF_MODEL_ID,
                revision=MOONDREAM_HF_REVISION,
                trust_remote_code=True,
            )
            self.model.eval()
            self._loaded = True
            log("[OK] Moondream2 loaded on ZeroGPU")
            return True
        except Exception as e:
            log(f"[ERROR] Vision model loading failed: {e}")
            self._loaded = False
            return False

    def _load_llama_cpp(self, log) -> bool:
        """Load Moondream2 natively using llama-cpp-python from local GGUF models."""
        log("Loading Moondream2 vision model (Local GGUF)...")
        try:
            from llama_cpp import Llama
            from llama_cpp.llama_chat_format import Llava15ChatHandler
            import os

            text_model = os.path.join(MODELS_DIR, "moondream2-text-model-q4_k_m.gguf")
            mmproj = os.path.join(MODELS_DIR, "moondream2-mmproj-f16.gguf")

            if not os.path.exists(text_model) or not os.path.exists(mmproj):
                log("[ERROR] Vision GGUF models not found in models/ directory.")
                self._loaded = False
                return False

            log(f"  Text model: {os.path.basename(text_model)}")
            log(f"  Vision proj: {os.path.basename(mmproj)}")

            chat_handler = Llava15ChatHandler(clip_model_path=mmproj)
            self.model = Llama(
                model_path=text_model,
                chat_handler=chat_handler,
                n_ctx=2048,
                verbose=False
            )
            self._loaded = True
            log("[OK] Moondream2 loaded natively via llama.cpp")
            return True
        except Exception as e:
            logger.error(f"Failed to load Moondream2 GGUF: {e}")
            log(f"[ERROR] Vision model loading failed: {e}")
            log("[WARN] Banner extraction will be skipped")
            self._loaded = False
            return False

    def unload_model(self):
        """Unload model to free memory for other agents."""
        if self.model is not None:
            del self.model
            self.model = None
        if self.tokenizer is not None:
            del self.tokenizer
            self.tokenizer = None
        self._openai_client = None
        self._loaded = False
        gc.collect()
        logger.info("Vision model unloaded, memory freed")

    # ─── Image Utilities ─────────────────────────────────────────────────────

    def _download_image(self, url: str) -> Optional[Image.Image]:
        """Download an image from a URL. Returns PIL Image or None."""
        try:
            headers = {
                "User-Agent": "Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36"
            }
            response = requests.get(
                url, headers=headers, timeout=BANNER_DOWNLOAD_TIMEOUT, stream=True
            )
            response.raise_for_status()

            # Check content type
            content_type = response.headers.get("Content-Type", "")
            if not content_type.startswith("image/"):
                logger.debug(f"Skipping non-image URL: {url}")
                return None

            # Check size (skip tiny tracking pixels and huge files)
            content_length = int(response.headers.get("Content-Length", 0))
            if content_length > 0 and content_length < 1000:
                logger.debug(f"Skipping tiny image (likely tracker): {url}")
                return None
            if content_length > 10_000_000:
                logger.debug(f"Skipping oversized image: {url}")
                return None

            img = Image.open(io.BytesIO(response.content))
            img = img.convert("RGB")

            # Skip very small images (tracking pixels)
            if img.width < 50 or img.height < 50:
                logger.debug(f"Skipping tiny image {img.width}x{img.height}: {url}")
                return None

            return img

        except Exception as e:
            logger.debug(f"Failed to download image from {url}: {e}")
            return None

    def _image_to_base64(self, image: Image.Image) -> str:
        """Convert PIL Image to base64 data URI for LMStudio vision API."""
        buf = io.BytesIO()
        image.save(buf, format="PNG")
        buf.seek(0)
        b64 = base64.b64encode(buf.read()).decode("utf-8")
        return f"data:image/png;base64,{b64}"

    # ─── Vision Inference ────────────────────────────────────────────────────

    def extract_text_from_image(self, image: Image.Image) -> str:
        """
        OCR text extraction from a banner image.
        Returns raw extracted text without any classification.
        """
        if not self._loaded:
            return ""

        if self._backend == BACKEND_LMSTUDIO:
            return self._extract_lmstudio_ocr(image)
        elif self._backend == BACKEND_ZEROGPU:
            return self._extract_moondream(image)
        else:
            return self._extract_llama_cpp(image)

    def _extract_lmstudio_ocr(self, image: Image.Image) -> str:
        """
        OCR-only extraction via LMStudio's OpenAI Vision API.
        Extracts all visible text from the image without classification.
        """
        if not self._openai_client:
            return ""

        image_b64 = self._image_to_base64(image)

        try:
            ocr_response = self._openai_client.chat.completions.create(
                model=self._resolved_vision_model,
                messages=[
                    {"role": "system", "content": VISION_SYSTEM_PROMPT},
                    {"role": "user", "content": [
                        {"type": "text", "text": VISION_OCR_PROMPT},
                        {"type": "image_url", "image_url": {"url": image_b64}},
                    ]},
                ],
                max_tokens=500,
                temperature=0.1,
            )
            ocr_text = ocr_response.choices[0].message.content.strip()
            logger.info(f"  Vision OCR: {len(ocr_text)} chars extracted")
            return ocr_text if ocr_text and len(ocr_text) >= 10 else ""

        except Exception as e:
            logger.warning(f"LMStudio vision OCR failed: {e}")
            return ""

    def _extract_moondream(self, image: Image.Image) -> str:
        """
        OCR text extraction using Moondream2 (transformers).
        Extracts all visible text without classification.
        """
        if not self._loaded or self.model is None:
            return ""

        try:
            enc_image = self.model.encode_image(image)

            # Pure OCR prompt
            result = self.model.answer_question(
                enc_image,
                VISION_OCR_PROMPT,
                self.tokenizer,
            )
            text = result.strip() if result else ""

            logger.info(f"  Vision OCR: {len(text)} chars extracted via Moondream")
            return text if len(text) >= 10 else ""

        except Exception as e:
            logger.warning(f"Moondream vision OCR failed: {e}")
            return ""

    def _extract_llama_cpp(self, image: Image.Image) -> str:
        """
        OCR text extraction using llama-cpp-python natively.
        """
        if not self._loaded or self.model is None:
            return ""

        try:
            image_b64 = self._image_to_base64(image)
            response = self.model.create_chat_completion(
                messages=[
                    {"role": "system", "content": VISION_SYSTEM_PROMPT},
                    {
                        "role": "user",
                        "content": [
                            {"type": "text", "text": VISION_OCR_PROMPT},
                            {"type": "image_url", "image_url": {"url": image_b64}},
                        ],
                    }
                ],
                max_tokens=500,
                temperature=0.1,
            )
            text = response["choices"][0]["message"]["content"].strip()
            logger.info(f"  Vision OCR: {len(text)} chars extracted via llama.cpp")
            return text if len(text) >= 10 else ""
        except Exception as e:
            logger.warning(f"Llama.cpp vision OCR failed: {e}")
            return ""

            # If too short, try a simpler OCR prompt
            if not text or len(text) < 15:
                result2 = self.model.answer_question(
                    enc_image,
                    "Read all text visible in this image. List every word and number you can see.",
                    self.tokenizer,
                )
                if result2 and len(result2.strip()) > len(text):
                    text = result2.strip()

            return text

        except Exception as e:
            logger.warning(f"Moondream2 inference failed: {e}")
            return ""

    # ─── Banner Processing Pipeline ──────────────────────────────────────────

    def process_email_banners(
        self,
        db: BankingDatabase,
        email_id: int,
        image_urls: List[str],
        bank_name: str,
    ) -> Dict:
        """
        Download and process banner images from an email.
        OCR-only: extracts text and appends to email body_text.
        Classification is deferred to the Classifier Agent (Qwen).
        """
        all_extracted_text = []

        for url in image_urls[:5]:  # Limit to 5 images per email
            logger.info(f"  Vision processing banner: {url[:80]}...")

            image = self._download_image(url)
            if not image:
                continue

            text = self.extract_text_from_image(image)
            if text and len(text) > 10:
                all_extracted_text.append(text[:500])
                logger.info(f"  OCR extracted: {text[:100]}...")

        # Append all extracted text back to the email's body_text
        if all_extracted_text:
            combined_text = "\n".join(all_extracted_text)
            db.append_body_text(email_id, combined_text)

        # Mark email as vision-processed
        db.mark_vision_processed(email_id)
        return {"texts_extracted": len(all_extracted_text)}

    def process_pending_emails(
        self,
        db: BankingDatabase,
        email_agent,
        progress_callback=None,
    ) -> Dict[str, int]:
        """
        Process all emails flagged for vision processing.
        OCR-only: extracts text and appends to email body_text.
        Classification is handled by the Classifier Agent in Phase 3.
        """
        results = {"processed": 0, "texts_extracted": 0, "skipped": 0}

        def log(msg):
            logger.info(msg)
            if progress_callback:
                progress_callback(msg)

        # Get emails flagged for vision
        pending = db.get_emails_needing_vision()

        if not pending:
            log("Vision: No emails need vision processing")
            return results

        log(f"Vision: {len(pending)} emails need OCR processing...")

        # Load model on-demand
        if not self.load_model(progress_callback):
            log("Vision model not available, skipping banner extraction")
            results["skipped"] = len(pending)
            return results

        try:
            for email_data in pending:
                email_id = email_data["id"]
                bank_name = email_data.get("bank_name", "Unknown")

                # Get image URLs from email HTML
                image_urls = email_agent.get_image_urls_for_email(email_id)

                if not image_urls:
                    db.mark_vision_processed(email_id)
                    results["skipped"] += 1
                    continue

                result = self.process_email_banners(db, email_id, image_urls, bank_name)
                results["processed"] += 1
                results["texts_extracted"] += result.get("texts_extracted", 0)

                if result.get("texts_extracted", 0) > 0:
                    log(f"  [{bank_name}] {result['texts_extracted']} texts OCR'd from banners")
                else:
                    log(f"  [{bank_name}] No text found in banners")

        finally:
            # Always unload to free memory for Classifier Agent
            self.unload_model()
            log("Vision model unloaded, memory freed")

        log(f"\nVision processing: {results['processed']} emails, "
            f"{results['texts_extracted']} texts extracted (classification deferred to Qwen)")

        return results