"""Simple utilities for working with Google's Gemini API"""

# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/00_gem.ipynb.

# %% auto 0
__all__ = ['upload_file', 'gem']

# %% ../nbs/00_gem.ipynb 5
import os, time
from pathlib import Path
from fastcore.all import *
from google import genai
from google.genai import types
from functools import partial
from fastprogress import progress_bar

# %% ../nbs/00_gem.ipynb 9
def _client():
    "Get Gemini client context manager"
    return genai.Client()

# %% ../nbs/00_gem.ipynb 12
def upload_file(pth):
    if not Path(pth).exists(): raise ValueError(f"File {pth} does not exist.")
    with _client() as c:
        f = c.files.upload(file=pth)
        time.sleep(2)
        for i in progress_bar(range(30)):
            try:
                f = c.files.get(name=f.name)
                if f.state == 'ACTIVE': return f
                elif f.state == 'FAILED': raise Exception(f'File processing for {pth} failed.')
                time.sleep(10)
            except: pass # because the gemini file thing is jank
        raise Exception(f'Timeout processing {pth}')

# %% ../nbs/00_gem.ipynb 15
def _is_url(s):
    "Check if string is a URL"
    if not isinstance(s, str): return False
    return (s.startswith('http://') or 
            s.startswith('https://') or 
            s.startswith('www.') or 
            'youtube.com' in s or 
            'youtu.be' in s)

def _make_part(o):
    "Convert object to Gemini Part"
    if isinstance(o, types.File):
        return types.Part.from_uri(file_uri=o.uri, mime_type=o.mime_type)
    if isinstance(o, (str, Path)):
        p = Path(o)
        if p.exists():
            if p.suffix.lower() == '.mp4':
                f = upload_file(o)
                return types.Part.from_uri(file_uri=f.uri, mime_type=f.mime_type)
            mime_map = {'.pdf': 'application/pdf', 
                        '.png': 'image/png', 
                        '.jpg': 'image/jpeg', 
                        '.jpeg': 'image/jpeg', 
                        '.gif': 'image/gif'}
            mime = mime_map.get(p.suffix.lower(), 'application/octet-stream')
            return types.Part.from_bytes(mime_type=mime, data=p.read_bytes())
        elif _is_url(o): return types.Part.from_uri(file_uri=o, mime_type='video/*')
        else: raise ValueError(f"Could not parse file or url: {o}")
    return None

# %% ../nbs/00_gem.ipynb 18
def gem(prompt, # Text prompt
        o=None, # Optional file/URL attachment or list of attachments
        model='gemini-2.5-flash',
        thinking=-1,
        search=False):
    "Generate content with Gemini"
    parts = [types.Part.from_text(text=prompt)]
    # Handle single attachment or list of attachments
    attachments = o if isinstance(o, list) else [o] if o else []
    for attachment in attachments:
        if part := _make_part(attachment): parts.insert(0, part)
    
    contents = types.Content(role='user', parts=parts) if attachments else prompt    
    config_dict = {
        'thinking_config': types.ThinkingConfig(thinking_budget=thinking),
        'response_mime_type': 'text/plain'
    }
    # Adjust media_resolution for videos for more tokens
    if any(p.file_data and p.file_data.mime_type.startswith('video') for p in parts):
        config_dict['media_resolution'] = 'MEDIA_RESOLUTION_LOW'
    config_dict['tools'] = []
    if search: config_dict['tools'].append(types.Tool(google_search=types.GoogleSearch()))
    cfg = types.GenerateContentConfig(**config_dict)
    with _client() as client:
        resp = client.models.generate_content(model=model, contents=contents, config=cfg)
    return resp.text
