"""
Vision MCP tool - visual Q&A via Groq Vision API (llama-3.2-11b-vision-preview).

Cloud path  : image -> Groq Llama 3.2 Vision -> structured answer (FREE, 14400 req/day)
Local path  : image -> Ollama qwen2-vl -> answer (dev only)

Registered in mcp_server.py as 'granite_vision.qa'.
"""
import os
import base64
import requests
from typing import Optional


def qa(image_path: str, question: str, model_path: Optional[str] = None) -> dict:
    """Answer a question about a document image using Groq Vision or local Ollama."""

    print(f"[vision.qa] image={image_path} question={question[:80]}")

    # CLOUD MODE: Groq Vision API
    if os.environ.get("USE_HF_CLOUD", "0") in ("1", "true", "True"):
        groq_key = os.environ.get("GROQ_API_KEY")
        groq_model = os.environ.get("GROQ_VISION_MODEL", "meta-llama/llama-4-scout-17b-16e-instruct")

        if groq_key:
            print(f"[granite_vision.qa] CLOUD MODE: {groq_model} via Groq")
            try:
                from groq import Groq

                client = Groq(api_key=groq_key)

                with open(image_path, "rb") as f:
                    img_b64 = base64.b64encode(f.read()).decode("utf-8")

                ext = os.path.splitext(image_path)[-1].lower().lstrip(".")
                mime = "image/jpeg" if ext in ("jpg", "jpeg") else f"image/{ext}"

                response = client.chat.completions.create(
                    model=groq_model,
                    messages=[
                        {
                            "role": "user",
                            "content": [
                                {
                                    "type": "text",
                                    "text": (
                                        "You are analyzing a document image. "
                                        "Look carefully at the layout, tables, numbers, and text. "
                                        "Answer the following question based only on what you see:\n\n"
                                        + question
                                    )
                                },
                                {
                                    "type": "image_url",
                                    "image_url": {
                                        "url": f"data:{mime};base64,{img_b64}"
                                    }
                                }
                            ]
                        }
                    ],
                    max_tokens=512
                )

                answer = response.choices[0].message.content.strip()
                print(f"[granite_vision.qa] Groq answer ({len(answer)} chars): {answer[:200]}")
                return {"answer": answer, "model": f"{groq_model} (Groq Vision)"}

            except Exception as e:
                print(f"[granite_vision.qa] Groq failed: {e}")
                return {
                    "answer": f"Groq Vision error: {e}",
                    "model": groq_model,
                    "error": str(e)
                }
        else:
            print("[granite_vision.qa] GROQ_API_KEY not set in HF secrets.")
            return {
                "answer": "GROQ_API_KEY is not configured. Please add it in HuggingFace Space secrets.",
                "model": "none",
                "error": "missing_groq_key"
            }

    # LOCAL MODE (dev only): Ollama
    local_model = os.environ.get("GRANITE_MODEL_PATH", "qwen2-vl:2b")
    print(f"[granite_vision.qa] LOCAL MODE: sending image to Ollama ({local_model})")

    try:
        with open(image_path, "rb") as f:
            img_b64 = base64.b64encode(f.read()).decode("utf-8")

        base_url = os.environ.get("OLLAMA_BASE_URL", "http://localhost:11434").rstrip("/")
        url = f"{base_url}/api/chat"
        payload = {
            "model": local_model,
            "messages": [
                {
                    "role": "user",
                    "content": f"Carefully analyze this document image. {question}",
                    "images": [img_b64]
                }
            ],
            "stream": False,
            "keep_alive": 60,
            "options": {"num_ctx": 1024, "temperature": 0.1}
        }
        r = requests.post(url, json=payload, timeout=300)
        if r.status_code == 200:
            answer = r.json().get("message", {}).get("content", "").strip()
            print(f"[granite_vision.qa] Ollama answer ({len(answer)} chars): {answer[:200]}")
            return {"answer": answer, "model": f"{local_model} (Ollama)"}
        else:
            msg = f"Ollama HTTP {r.status_code}: {r.text}"
            print(f"[granite_vision.qa] {msg}")
            return {"answer": msg, "model": f"{local_model} (Ollama)", "error": "http_error"}

    except Exception as e:
        msg = f"[ERROR] Ollama Q&A failed: {e}"
        print(f"[granite_vision.qa] {msg}")
        return {"answer": msg, "model": f"{local_model} (Ollama)", "error": str(e)}