"""Use Gemini Vision to identify meal components for better USDA queries.

Gemini is used as a nutrition metadata assistant, not as the segmentation
metric model. The local YOLO model still provides area fractions; Gemini helps
turn coarse classes like "meat" into better food queries like "grilled chicken
breast" or "beef teriyaki".

Set the API key before running:
    $env:GEMINI_API_KEY = "your-api-key"

Example:
    python scripts/gemini_meal_analyzer.py --image path/to/meal.jpg --segments data/nutrition/segments.json
"""

from __future__ import annotations

import argparse
import base64
import json
import mimetypes
import os
from pathlib import Path
from typing import Any
from urllib.error import HTTPError, URLError
from urllib.parse import urlencode
from urllib.request import Request, urlopen


PROJECT_ROOT = Path(__file__).resolve().parents[1]
DEFAULT_SEGMENTS = PROJECT_ROOT / "data" / "nutrition" / "sample_segments.json"
DEFAULT_OUTPUT = PROJECT_ROOT / "data" / "nutrition" / "gemini_meal_analysis.json"
DEFAULT_MODEL = "gemini-2.5-flash"
GEMINI_ENDPOINT = "https://generativelanguage.googleapis.com/v1beta/models/{model}:generateContent"


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--image", type=Path, required=True)
    parser.add_argument("--segments", type=Path, default=DEFAULT_SEGMENTS)
    parser.add_argument("--output", type=Path, default=DEFAULT_OUTPUT)
    parser.add_argument("--model", default=DEFAULT_MODEL)
    parser.add_argument("--api-key", default=os.getenv("GEMINI_API_KEY"))
    return parser.parse_args()


def load_segments(path: Path) -> Any:
    if not path.exists():
        return []
    return json.loads(path.read_text(encoding="utf-8"))


def image_part(path: Path) -> dict[str, Any]:
    mime_type, _ = mimetypes.guess_type(path)
    if mime_type is None:
        mime_type = "image/jpeg"

    return {
        "inline_data": {
            "mime_type": mime_type,
            "data": base64.b64encode(path.read_bytes()).decode("ascii"),
        }
    }


def build_prompt(segments: Any) -> str:
    return f"""Analyze this meal image for a macronutrient estimation pipeline.

The segmentation model produced these coarse area fractions:
{json.dumps(segments, indent=2)}

Return only valid JSON with this schema:
{{
  "meal_summary": "short description",
  "components": [
    {{
      "class_name": "meat|rice|vegetables",
      "likely_food": "specific food name",
      "fdc_query": "short USDA FoodData Central search query",
      "confidence": 0.0,
      "notes": "short note"
    }}
  ],
  "portion_assumptions": {{
    "total_plate_grams": 500,
    "rationale": "short rationale"
  }}
}}

Prefer simple USDA-friendly search terms over restaurant brand names. If the
image is ambiguous, choose a conservative generic query.
"""


def call_gemini(api_key: str, model: str, image_path: Path, segments: Any) -> dict[str, Any]:
    endpoint = GEMINI_ENDPOINT.format(model=model)
    url = f"{endpoint}?{urlencode({'key': api_key})}"
    payload = {
        "contents": [
            {
                "parts": [
                    {"text": build_prompt(segments)},
                    image_part(image_path),
                ]
            }
        ],
        "generationConfig": {
            "temperature": 0.2,
            "response_mime_type": "application/json",
        },
    }

    request = Request(
        url,
        data=json.dumps(payload).encode("utf-8"),
        headers={"Content-Type": "application/json"},
        method="POST",
    )

    try:
        with urlopen(request, timeout=60) as response:
            return json.loads(response.read().decode("utf-8"))
    except HTTPError as exc:
        detail = exc.read().decode("utf-8", errors="replace")
        raise RuntimeError(
            f"Gemini API error {exc.code}. "
            "Check that GEMINI_API_KEY is valid, billing/API access is enabled, "
            f"and model '{model}' is available. Response: {detail}"
        ) from exc
    except URLError as exc:
        raise RuntimeError(
            "Gemini connection error. Check network/proxy access to "
            f"{GEMINI_ENDPOINT.format(model=model)}. Detail: {exc.reason}"
        ) from exc


def extract_json_text(response: dict[str, Any]) -> dict[str, Any]:
    candidates = response.get("candidates", [])
    if not candidates:
        raise RuntimeError("Gemini response did not include candidates.")

    parts = candidates[0].get("content", {}).get("parts", [])
    text = "".join(part.get("text", "") for part in parts).strip()
    if not text:
        raise RuntimeError("Gemini response did not include text.")

    return json.loads(text)


def main() -> None:
    args = parse_args()
    if not args.api_key:
        raise SystemExit("Missing Gemini API key. Set GEMINI_API_KEY or pass --api-key.")

    segments = load_segments(args.segments)
    try:
        raw_response = call_gemini(args.api_key, args.model, args.image, segments)
        analysis = extract_json_text(raw_response)
    except Exception as exc:
        raise SystemExit(f"[ERROR] {exc}") from None

    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_text(json.dumps(analysis, indent=2), encoding="utf-8")
    print(f"Wrote Gemini meal analysis: {args.output}")


if __name__ == "__main__":
    main()