# Familiar - Self-Hosted AI Companion Platform
# Copyright (c) 2026 George Scott Foley
# ORCID: 0009-0006-4957-0540
# Email: Georgescottfoley@proton.me
# Licensed under the MIT License - see LICENSE file for details

"""
Vision LLM integration for Smart Intake.

Uses Ollama vision models (qwen2.5-vl, llava, llama3.2-vision, glm-edge-v)
via the /api/chat endpoint with base64-encoded images. No new dependencies —
uses httpx (core dep) to call local Ollama.

Falls back to Tesseract OCR if no vision model is available.
Auto-installs OCR deps if missing. Auto-pulls a vision model if none found.
"""

import base64
import logging
import subprocess
from pathlib import Path
from typing import Optional

logger = logging.getLogger(__name__)

# Vision-capable models in preference order (smallest last for auto-pull)
VISION_MODELS = [
    "qwen2.5-vl",
    "qwen2.5-vl:7b",
    "llava",
    "llava:13b",
    "llama3.2-vision",
    "llama3.2-vision:11b",
    "glm-edge-v",
    "moondream",
]

# Model to auto-pull if none found (small, fast)
_AUTO_PULL_MODEL = "moondream"

IMAGE_EXTENSIONS = {
    ".png",
    ".jpg",
    ".jpeg",
    ".tiff",
    ".tif",
    ".bmp",
    ".gif",
    ".webp",
}

_cached_vision_model: Optional[str] = None
_ollama_status: Optional[str] = None  # "running", "stopped", "missing"


def _get_ollama_base() -> str:
    """Get Ollama API base URL."""
    import os

    return os.environ.get("OLLAMA_HOST", "http://localhost:11434")


def _check_ollama_status() -> str:
    """Check if Ollama is installed and running. Caches result."""
    global _ollama_status
    if _ollama_status:
        return _ollama_status

    try:
        import httpx

        resp = httpx.get(f"{_get_ollama_base()}/api/tags", timeout=3.0)
        if resp.status_code == 200:
            _ollama_status = "running"
            return "running"
    except Exception:
        pass

    # Check if ollama binary exists
    import shutil

    if shutil.which("ollama"):
        _ollama_status = "stopped"
        return "stopped"

    _ollama_status = "missing"
    return "missing"


def _auto_pull_vision_model() -> Optional[str]:
    """Attempt to auto-pull a small vision model via Ollama."""
    status = _check_ollama_status()
    if status != "running":
        return None

    logger.info("No vision model found — auto-pulling %s...", _AUTO_PULL_MODEL)
    try:
        result = subprocess.run(
            ["ollama", "pull", _AUTO_PULL_MODEL],
            capture_output=True,
            text=True,
            timeout=600,  # 10 min for model download
        )
        if result.returncode == 0:
            logger.info("Auto-pulled vision model: %s", _AUTO_PULL_MODEL)
            return _AUTO_PULL_MODEL
        logger.warning("Failed to auto-pull %s: %s", _AUTO_PULL_MODEL, result.stderr[:200])
    except subprocess.TimeoutExpired:
        logger.warning("Auto-pull timed out for %s", _AUTO_PULL_MODEL)
    except Exception as e:
        logger.warning("Auto-pull failed: %s", e)
    return None


def detect_vision_model(auto_pull: bool = True) -> Optional[str]:
    """Detect first available vision model from Ollama. Auto-pulls if none found."""
    global _cached_vision_model
    if _cached_vision_model:
        return _cached_vision_model

    status = _check_ollama_status()
    if status != "running":
        return None

    try:
        import httpx

        resp = httpx.get(f"{_get_ollama_base()}/api/tags", timeout=5.0)
        if resp.status_code != 200:
            return None
        models = resp.json().get("models", [])
        available = {m["name"].split(":")[0] for m in models}
        available_full = {m["name"] for m in models}

        for candidate in VISION_MODELS:
            if candidate in available or candidate in available_full:
                _cached_vision_model = candidate
                logger.info("Vision model detected: %s", candidate)
                return candidate
    except Exception as e:
        logger.debug("Ollama vision model detection failed: %s", e)

    # No vision model found — try auto-pull
    if auto_pull:
        pulled = _auto_pull_vision_model()
        if pulled:
            _cached_vision_model = pulled
            return pulled

    return None


def clear_vision_cache():
    """Clear cached vision model and Ollama status (for testing)."""
    global _cached_vision_model, _ollama_status
    _cached_vision_model = None
    _ollama_status = None


def encode_image_base64(filepath: Path) -> str:
    """Read image file and return base64-encoded string."""
    with open(filepath, "rb") as f:
        return base64.b64encode(f.read()).decode("utf-8")


def extract_with_vision(
    filepath: Path,
    prompt: Optional[str] = None,
    model: Optional[str] = None,
    timeout: float = 60.0,
) -> Optional[str]:
    """
    Extract text/data from an image using a vision LLM via Ollama.

    Returns:
        Extracted text, or None if vision is unavailable.
    """
    vision_model = model or detect_vision_model()
    if not vision_model:
        return None

    if not filepath.exists():
        logger.warning("Image file not found: %s", filepath)
        return None

    if filepath.suffix.lower() not in IMAGE_EXTENSIONS:
        logger.warning("Not an image file: %s", filepath.suffix)
        return None

    if not prompt:
        prompt = (
            "Extract ALL text visible in this image. Preserve the original layout, "
            "formatting, and structure as closely as possible. Include headers, "
            "labels, numbers, dates, addresses, and any handwritten text. "
            "If this is a form or receipt, identify each field and its value."
        )

    image_b64 = encode_image_base64(filepath)

    try:
        import httpx

        resp = httpx.post(
            f"{_get_ollama_base()}/api/chat",
            json={
                "model": vision_model,
                "messages": [
                    {
                        "role": "user",
                        "content": prompt,
                        "images": [image_b64],
                    }
                ],
                "stream": False,
            },
            timeout=timeout,
        )
        if resp.status_code != 200:
            logger.warning("Vision API error %d: %s", resp.status_code, resp.text[:200])
            return None

        result = resp.json()
        content = result.get("message", {}).get("content", "")
        return content.strip() if content else None

    except Exception as e:
        logger.warning("Vision extraction failed: %s", e)
        return None


def _ensure_ocr() -> bool:
    """Auto-install OCR dependencies if missing. Returns True if available."""
    try:
        import pytesseract  # noqa: F401
        from PIL import Image  # noqa: F401

        return True
    except ImportError:
        pass

    # Try auto-install via familiar's dep system
    try:
        from familiar.core.deps import OCR_PACKAGES, ensure_packages

        ok, _ = ensure_packages(OCR_PACKAGES)
        return ok
    except ImportError:
        return False


def extract_with_ocr(filepath: Path) -> Optional[str]:
    """Fallback: extract text from image using Tesseract OCR. Auto-installs deps."""
    if not _ensure_ocr():
        logger.debug("OCR unavailable — pytesseract/Pillow could not be installed")
        return None

    try:
        import pytesseract
        from PIL import Image

        image = Image.open(str(filepath))
        if image.mode != "L":
            image = image.convert("L")
        text = pytesseract.image_to_string(image)
        return text.strip() if text.strip() else None
    except Exception as e:
        # Check for missing tesseract system package specifically
        err_str = str(e).lower()
        if "tesseract" in err_str and ("not found" in err_str or "not installed" in err_str):
            logger.warning(
                "Tesseract system package missing. Install with: sudo apt install tesseract-ocr"
            )
        else:
            logger.warning("OCR extraction failed: %s", e)
        return None


def _get_extraction_diagnostics() -> str:
    """Build a helpful diagnostic message when all extraction methods fail."""
    parts = []

    status = _check_ollama_status()
    if status == "missing":
        parts.append("Ollama is not installed (needed for vision AI)")
    elif status == "stopped":
        parts.append("Ollama is installed but not running — start it with: ollama serve")
    else:
        # Ollama running but no vision model
        parts.append(f"No vision model installed — pull one with: ollama pull {_AUTO_PULL_MODEL}")

    if not _ensure_ocr():
        parts.append(
            "OCR not available — install with: sudo apt install tesseract-ocr && pip install pytesseract"
        )

    if not parts:
        parts.append("Both vision and OCR are available but could not read this image")

    return ". ".join(parts) + "."


def extract_text_from_image(
    filepath: Path,
    prompt: Optional[str] = None,
    prefer_ocr: bool = False,
) -> str:
    """
    Tiered image text extraction: vision LLM first, OCR fallback.

    Returns:
        Extracted text, or actionable error message explaining what's needed.
    """
    path = Path(filepath)
    if not path.exists():
        return f"File not found: {filepath}"

    if prefer_ocr:
        ocr_result = extract_with_ocr(path)
        if ocr_result:
            return ocr_result
        vision_result = extract_with_vision(path, prompt=prompt)
        if vision_result:
            return vision_result
    else:
        # Default: vision first (better quality), OCR fallback
        vision_result = extract_with_vision(path, prompt=prompt)
        if vision_result:
            return vision_result
        ocr_result = extract_with_ocr(path)
        if ocr_result:
            return ocr_result

    # Both failed — give actionable diagnostics
    return f"Could not extract text from image. {_get_extraction_diagnostics()}"
