import os
import time
from pathlib import Path

from dotenv import load_dotenv
from google import genai
from google.genai import errors, types


# ============================================================
# Config
# ============================================================

load_dotenv()

MODELS = [
    "gemini-3.6-flash",
    "gemini-3.5-flash",
    "gemini-3.5-flash-lite",
    "gemini-flash-lite-latest",
]
MAX_QUESTION_LENGTH = 300

SYSTEM_PROMPT = (
    "You answer questions about a single photo. "
    "Base your answer only on what is visible in the image. "
    "If the answer can't be seen in the photo, say so. "
    "Answer in one to three short sentences."
)

GENERATION_CONFIG = types.GenerateContentConfig(
    system_instruction=SYSTEM_PROMPT,
    temperature=0.2,
    thinking_config=types.ThinkingConfig(thinking_level="low"),
)

HTTP_OPTIONS = types.HttpOptions(
    timeout=10_000,
    retry_options=types.HttpRetryOptions(attempts=1),
)

client = genai.Client(http_options=HTTP_OPTIONS) if os.getenv("GEMINI_API_KEY") else None


# ============================================================
# Q&A
# ============================================================

def ask_about_image(image_path: str, question: str) -> str:
    question = question.strip()
    if not question:
        return "Type a question about the image first."
    if len(question) > MAX_QUESTION_LENGTH:
        return f"Please keep the question under {MAX_QUESTION_LENGTH} characters."
    if client is None:
        return "The Q&A model isn't configured: GEMINI_API_KEY is missing."

    image_part = types.Part.from_bytes(
        data=Path(image_path).read_bytes(),
        mime_type="image/jpeg",
    )

    for model in MODELS:
        start = time.perf_counter()
        try:
            response = client.models.generate_content(
                model=model,
                contents=[image_part, question],
                config=GENERATION_CONFIG,
            )
            print(f"[qa] {model}: OK in {time.perf_counter() - start:.1f}s")
            return response.text.strip()
        except errors.ServerError as e:
            print(f"[qa] {model}: {e.code} after {time.perf_counter() - start:.1f}s, trying next")
            continue
        except errors.ClientError as e:
            print(f"[qa] {model}: {e.code} after {time.perf_counter() - start:.1f}s")
            if e.code in (404, 429):
                continue
            return f"The request was rejected ({e.code}). Please try a different question."
        except Exception as e:
            print(f"[qa] {model}: {type(e).__name__} after {time.perf_counter() - start:.1f}s, trying next")
            continue

    return "The model is busy right now. Please try again in a moment."


# ============================================================
# Manual test
# ============================================================

if __name__ == "__main__":
    import sys

    path = "data/Flicker8k_Dataset/1000268201_693b08cb0e.jpg"
    question = " ".join(sys.argv[1:]) or "What is the girl wearing?"
    print("Q:", question)
    start = time.perf_counter()
    answer = ask_about_image(path, question)
    print("A:", answer)
    print(f"Total: {time.perf_counter() - start:.1f}s")