"""Components that are not part of the core functionality of the library but in general make its use better"""

# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/01_utils.ipynb.

# %% auto 0
__all__ = ['all_model_types', 'thinking_models', 'imagen_models', 'vertex_models', 'models', 'pricings', 'audio_token_pricings',
           'usage', 'get_pricing', 'get_repr', 'det_repr', 'contents', 'response_md_repr']

# %% ../nbs/01_utils.ipynb 4
import os
import base64

from google import genai
from google.genai import types

from fastcore.all import *

# %% ../nbs/01_utils.ipynb 10
all_model_types = {
    "gemini-2.0-flash": "llm-vertex#gemini-2.0-flash",
    "gemini-2.0-flash-lite": "llm#gemini-2.0-flash-lite",
    "gemini-2.5-pro-exp-03-25": "llm-thinking#gemini-2.5-pro",
    "gemini-2.0-flash-exp": "llm-imagen#gemini-2.0-flash",
    "gemini-2.0-flash-exp-image-generation": "llm-imagen#gemini-2.0-flash",
    "gemini-2.0-flash-001": "llm-vertex#gemini-2.0-flash",
    "gemini-2.0-pro-exp-02-05": "llm#gemini-2.0-pro",
    "gemini-1.5-flash": "llm-vertex#gemini-1.5-flash",
    "gemini-1.5-pro": "llm-vertex#gemini-1.5-pro",
    "gemini-1.5-pro-002": "llm-vertex#gemini-1.5-pro",
    "gemini-1.5-flash-8b": "llm#gemini-1.5-flash-8b",
    "gemini-2.0-flash-thinking-exp-01-21": "llm-thinking#gemini-2.0-flash-thinking",
    "imagen-3.0-generate-002": "imagen#imagen-3.0"
}

thinking_models = [m for m in all_model_types if "thinking" in all_model_types[m]]

imagen_models = [m for m in all_model_types if "imagen" in all_model_types[m]]

vertex_models = [m for m in all_model_types if "vertex" in all_model_types[m]]

models = [m for m in all_model_types if "llm" in all_model_types[m]]

models

# %% ../nbs/01_utils.ipynb 14
def usage(inp=0,     # Number of input tokens (excluding cached)
          out=0,     # Number of output tokens
          cached=0): # Number of cached tokens
    """A quicker and simpler constructor for the Usage Metadata model"""
    return types.GenerateContentResponseUsageMetadata(cached_content_token_count=cached, 
                                                      candidates_token_count=out, 
                                                      prompt_token_count=inp + cached, 
                                                      total_token_count=inp + out + cached)

# %% ../nbs/01_utils.ipynb 18
@patch(as_prop=True)
def cached(self: types.GenerateContentResponseUsageMetadata): 
    return self.cached_content_token_count or 0

@patch(as_prop=True)
def inp(self: types.GenerateContentResponseUsageMetadata): 
    return (self.prompt_token_count - self.cached) or 0

@patch(as_prop=True)
def out(self: types.GenerateContentResponseUsageMetadata): 
    return self.candidates_token_count or 0

@patch(as_prop=True)
def total(self: types.GenerateContentResponseUsageMetadata): 
    return self.total_token_count or self.prompt_token_count + self.candidates_token_count

# %% ../nbs/01_utils.ipynb 21
@patch
def __repr__(self: types.GenerateContentResponseUsageMetadata):
    return f"Cached: {self.cached}; In: {self.inp}; Out: {self.out}; Total: {self.total}"

@patch
def _repr_markdown_(self: types.GenerateContentResponseUsageMetadata):
    return self.__repr__()

# %% ../nbs/01_utils.ipynb 24
@patch
def __add__(self: types.GenerateContentResponseUsageMetadata, other):
    cached = getattr(self, "cached", 0) + getattr(other, "cached", 0)
    return usage(self.inp + other.inp, self.out + other.out, cached)

# %% ../nbs/01_utils.ipynb 27
# $/1M input (non cached) tokens, $/1M output tokens, $/1M cached input tokens, 

pricings = {
    'gemini-2.0-flash': [0.1, 0.4, 0.025],
    'gemini-2.0-flash-lite': [0.075, 0.3, 0.01875],
    'gemini-1.5-flash_short': [0.075, 0.3, 0.01875],
    'gemini-1.5-flash_long': [0.15, 0.6, 0.0375], 
    'gemini-1.5-flash-8b_short': [0.0375, 0.15, 0.01],
    'gemini-1.5-flash-8b_long': [0.075, 0.3, 0.02],
    'gemini-1.5-pro_short': [1.25, 5., 0.3125],   
    'gemini-1.5-pro_long': [2.5, 10., 0.625],
 }


audio_token_pricings = {
    'gemini-2.0-flash': [0.7, 0.4, 0.175],
}

def get_pricing(model, prompt_tokens):
    if "exp" in model: return [0, 0, 0]
    suff = "_long" if prompt_tokens > 128_000 else "_short"
    m = all_model_types.get(model, "#").split("#")[-1]
    m += suff if "1.5" in m else ""
    return pricings.get(m, [0, 0, 0])


# %% ../nbs/01_utils.ipynb 30
@patch(as_prop=True)
def cost(self: types.GenerateContentResponse):
    ip, op, cp = get_pricing(self.model_version, self.usage_metadata.prompt_token_count)
    return ((self.usage_metadata.inp * ip) + (self.usage_metadata.out * op) + (self.usage_metadata.cached * cp)) / 1e6

# %% ../nbs/01_utils.ipynb 32
@patch(as_prop=True)
def cost(self: types.GenerateImagesResponse): return 0.03 * len(self.generated_images)

# %% ../nbs/01_utils.ipynb 41
def get_repr(m, lab=""):
    """Recurisvely fetch the markdown representation of genai.types fields, wrapping lists into `<details>` blocks"""
    if hasattr(m, '_repr_markdown_'): return m._repr_markdown_()
    if is_listy(m): return "\n".join([f"<details open='true'><summary>{lab}[{i}]</summary>{get_repr(li)}</details>" for i, li in enumerate(m)])
    if isinstance(m, dict): return "<ul>" + "\n".join([f"<li><b>{i}</b>: {get_repr(li, i)}</li>" for i, li in m.items()]) + "</ul>"
    if isinstance(m, bytes): return m[:10] + b'...'
    return str(m)


# %% ../nbs/01_utils.ipynb 45
def det_repr(m): return "<ul>" + "".join(f"<li><code>{d}</code>: {get_repr(getattr(m, d), d)}</li>" for d in m.model_fields_set) + "</ul>"

# %% ../nbs/01_utils.ipynb 47
@patch
def _repr_markdown_(self: genai._common.BaseModel):
    return det_repr(self)

# %% ../nbs/01_utils.ipynb 51
def contents(r: genai.types.GenerateContentResponse | genai.types.GenerateImagesResponse):
    """Returns a dictionary with the contents of a Gemini model response"""
    cts = {'text': '', 'images': []}
    for part in nested_idx(r, 'candidates', 0, 'content', 'parts') or []:
        if part.text is not None: cts['text'] += part.text
        if part.inline_data is not None:
            cts['images'].append(types.Image(image_bytes=part.inline_data.data, mime_type=part.inline_data.mime_type))
    for im in nested_idx(r, 'generated_images') or []:
        cts['images'].append(im.image)
    return cts

# %% ../nbs/01_utils.ipynb 53
def response_md_repr(resp: types.GenerateContentResponse | types.GenerateImagesResponse):
    c = None
    cts = contents(resp)
    if cts['images'] or cts['text']:
        c = ''
        for img in cts['images']:
            b64 = base64.b64encode(img.image_bytes).decode("utf-8")
            c += f'<div style="width: 200px; height: auto;"><img src="data:{img.mime_type};base64,{b64}" /></div>'
        c += cts['text'].replace("\n", "\n\n")
    elif getattr(resp, "function_calls", False):
        calls = (f"<code>{call.name}({', '.join([f'{a}={v}' for a, v in call.args.items()])})</code>" for call in resp.function_calls)
        calls_repr = '\n'.join(f'<li>{c}</li>' for c in calls)
        c = f"<ul>{calls_repr}</ul>"
    dets = det_repr(resp)
    return f"""{c}\n<details>{dets}</details>"""
    

@patch
def _repr_markdown_(self: types.GenerateContentResponse):
    return response_md_repr(self)

@patch
def _repr_markdown_(self: types.GenerateImagesResponse):
    return response_md_repr(self)
