import os
import json
import io
import logging
from PIL import Image
from google import genai
from google.genai import types
from ..config_loader import get_config

log = logging.getLogger("ainventory")

def get_best_models():
    # Priority list based on aliases and future-gen models found in user logs
    return [
        "gemini-flash-latest",           # Points to the best stable Flash version
        "gemini-3.1-flash-lite-preview", # Cutting edge 3.1 Lite
        "gemini-flash-lite-latest",      # Best Lite alias
        "gemini-pro-latest"              # Pro fallback alias
    ]

def extract(image_bytes: bytes, prompt: str):
    config = get_config()
    api_key = config.get("ai", {}).get("gemini_api_key")
    if not api_key:
        log.error("CRITICAL: gemini_api_key is MISSING in configuration!")
        return None
    
    # Log partial key for safety debug
    key_hint = f"{api_key[:4]}...{api_key[-4:]}" if len(api_key) > 8 else "too short"
    log.info(f"🔑 Using API Key: {key_hint}")

    try:
        # Switching to v1beta which is required for many experimental/new models
        client = genai.Client(api_key=api_key, http_options={'api_version': 'v1beta'})
        
        # DEBUG: List allowed models for this key
        log.debug("🔍 Checking available models for your key...")
        try:
            for m in client.models.list():
                log.debug(f"   - Found: {m.name}")
        except Exception as list_e:
            log.warning(f"   ⚠️ Could not list models: {list_e}")

        models_to_try = get_best_models()

        # Try models in order
        for model_name in models_to_try:
            try:
                log.info(f"🚀 AI Launching (v2 SDK): {model_name}...")
                
                # In the new SDK, we pass a list of parts (text string and image bytes)
                response = client.models.generate_content(
                    model=model_name,
                    contents=[
                        prompt,
                        types.Part.from_bytes(data=image_bytes, mime_type="image/jpeg")
                    ]
                )

                if not response or not response.text:
                    log.warning(f"⚠️ Gemini {model_name} returned NO text.")
                    continue

                text = response.text.strip()
                log.info(f"✅ AI Response Received ({len(text)} bytes)")
                log.debug(f"FULL RESPONSE:\n{text}")
                
                # Extract JSON block
                if "```json" in text:
                    text = text.split("```json")[1].split("```")[0].strip()
                elif "```" in text:
                    text = text.split("```")[1].strip()
                     
                return json.loads(text)
                
            except Exception as e:
                log.error(f"❌ Gemini {model_name} failed: {e}")
                continue
                
    except Exception as outer_e:
        log.error(f"❌ Gemini Client Init failed: {outer_e}")
        
    return None