"""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', 'upload_file_async', 'gem', 'gem_async']

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

# %% ../nbs/00_gem.ipynb 10
async def _to_thread(func, /, *args, **kwargs):
    "Run blocking helper on a background thread"
    return await asyncio.to_thread(func, *args, **kwargs)


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

# %% ../nbs/00_gem.ipynb 15
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 16
async def upload_file_async(pth):
    "Async wrapper around upload_file"
    return await _to_thread(upload_file, pth)


# %% ../nbs/00_gem.ipynb 20
_VIDEO_EXTS = {'.mp4', '.mpeg', '.mov', '.avi', '.flv', '.mpg', '.webm', '.wmv', '.3gpp'}
_AUDIO_EXTS = {'.wav', '.mp3', '.aiff', '.aac', '.ogg', '.flac', '.m4a'}
_INLINE_MIME_MAP = {
    '.pdf': 'application/pdf',
    '.png': 'image/png',
    '.jpg': 'image/jpeg',
    '.jpeg': 'image/jpeg',
    '.gif': 'image/gif',
    '.txt': 'text/plain',
    '.vtt': 'text/plain',
    '.md': 'text/markdown',
    '.json': 'text/plain',  # Gemini doesn't accept application/json
    '.yaml': 'text/yaml',
    '.yml': 'text/yaml',
    '.toml': 'text/plain',
    '.ipynb': 'text/plain'
}

def _normalize_attachments(o):
    "Ensure attachments are always a list"
    if isinstance(o, list): return o
    return [o] if o else []

def _is_media_file(p: Path):
    return p.suffix.lower() in (_VIDEO_EXTS | _AUDIO_EXTS)

def _mime_for_path(p: Path):
    mime = _INLINE_MIME_MAP.get(p.suffix.lower())
    if mime: return mime
    guessed_mime, _ = mimetypes.guess_type(str(p))
    if guessed_mime is None:
        raise ValueError(f"Cannot determine MIME type for file: {p}. Unsupported extension: {p.suffix}")
    return guessed_mime


# %% ../nbs/00_gem.ipynb 22
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 (sync)"
    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 _is_media_file(p):
                f = upload_file(o)
                return types.Part.from_uri(file_uri=f.uri, mime_type=f.mime_type)
            return types.Part.from_bytes(mime_type=_mime_for_path(p), 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 23
async def _make_part_async(o):
    "Async wrapper around _make_part"
    return await _to_thread(_make_part, o)


# %% ../nbs/00_gem.ipynb 28
def _parts_from_attachments(prompt, attachments):
    "Build Parts list from normalized attachments"
    parts = [types.Part.from_text(text=prompt)]
    for attachment in attachments:
        if part := _make_part(attachment):
            parts.insert(0, part)
    return parts

#| export
async def _parts_from_attachments_async(prompt, attachments):
    "Async Parts builder"
    parts = [types.Part.from_text(text=prompt)]
    for attachment in attachments:
        part = await _make_part_async(attachment)
        if part:
            parts.insert(0, part)
    return parts


# %% ../nbs/00_gem.ipynb 30
def _content_payload(prompt, attachments, parts):
    "Return prompt or Content depending on attachments"
    return types.Content(role='user', parts=parts) if attachments else prompt


# %% ../nbs/00_gem.ipynb 32
def _build_config(thinking, search, parts):
    config_dict = {
        'thinking_config': types.ThinkingConfig(thinking_budget=thinking),
        'response_mime_type': 'text/plain'
    }
    if any(getattr(p, 'file_data', None) and getattr(p.file_data, 'mime_type', '').startswith('video')
           for p in parts):
        config_dict['media_resolution'] = 'MEDIA_RESOLUTION_LOW'
    tools = []
    if search:
        tools.append(types.Tool(google_search=types.GoogleSearch()))
    if tools:
        config_dict['tools'] = tools
    return types.GenerateContentConfig(**config_dict)


# %% ../nbs/00_gem.ipynb 34
def _chunk_text(chunk):
    "Best-effort way to pull text out of a streaming chunk"
    if hasattr(chunk, 'text') and chunk.text:
        return chunk.text
    cand = getattr(chunk, 'candidates', None)
    if cand:
        candidate = cand[0]
        content = getattr(candidate, 'content', None)
        if content and getattr(content, 'parts', None):
            part = content.parts[-1]
            if hasattr(part, 'text') and part.text:
                return part.text
    return str(chunk)

#| export
async def _stream_generate_content(model, contents, cfg):
    loop = asyncio.get_running_loop()
    queue = asyncio.Queue()

    def _run_stream():
        try:
            with _client() as client:
                for chunk in client.models.generate_content_stream(model=model, contents=contents, config=cfg):
                    fut = asyncio.run_coroutine_threadsafe(queue.put(chunk), loop)
                    fut.result()
        except Exception as e:
            asyncio.run_coroutine_threadsafe(queue.put(e), loop).result()
        finally:
            asyncio.run_coroutine_threadsafe(queue.put(None), loop).result()

    threading.Thread(target=_run_stream, daemon=True).start()

    while True:
        item = await queue.get()
        if item is None:
            break
        if isinstance(item, Exception):
            raise item
        yield _chunk_text(item)



# %% ../nbs/00_gem.ipynb 35
def _generate_content(model, contents, cfg):
    with _client() as client:
        return client.models.generate_content(model=model, contents=contents, config=cfg)


# %% ../nbs/00_gem.ipynb 36
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"
    attachments = _normalize_attachments(o)
    parts = _parts_from_attachments(prompt, attachments)
    contents = _content_payload(prompt, attachments, parts)
    cfg = _build_config(thinking, search, parts)
    resp = _generate_content(model, contents, cfg)
    return resp.text

# %% ../nbs/00_gem.ipynb 37
async def gem_async(prompt,
                    o=None,
                    model='gemini-2.5-flash',
                    thinking=-1,
                    search=False,
                    stream=False):
    "Async wrapper around gem using background threads. Set stream=True for an async iterator."
    attachments = _normalize_attachments(o)
    parts = await _parts_from_attachments_async(prompt, attachments)
    contents = _content_payload(prompt, attachments, parts)
    cfg = _build_config(thinking, search, parts)
    if stream:
        return _stream_generate_content(model, contents, cfg)
    resp = await _to_thread(_generate_content, model, contents, cfg)
    return resp.text

