"""The main aspects of content generation with Gemini (like streaming, generation parameters, input manipulation etc.)"""

# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/02_core.ipynb.

# %% auto 0
__all__ = ['is_youtube_url', 'mk_part', 'mk_parts', 'mk_content', 'mk_contents', 'mk_msg', 'mk_msgs', 'goog_doc', 'prep_tool',
           'f_result', 'f_results', 'mk_fres_content', 'prep_tools']

# %% ../nbs/02_core.ipynb 4
from .utils import *

import os

import mimetypes
import inspect

from urllib.parse import urlparse, parse_qs
from functools import wraps

from google import genai
from google.genai import types

import PIL

from fastcore import imghdr
from fastcore.all import *
from fastcore.docments import *

from toolslm.funccall import call_func



# %% ../nbs/02_core.ipynb 8
def is_youtube_url(url: str) -> bool:
    """Check if the given URL is a valid YouTube video URL using urllib for parsing."""
    parsed = urlparse(url)
    if parsed.scheme not in ('http', 'https'): return False
    host = parsed.netloc.lower()
    
    # Standard YouTube URL (e.g., https://www.youtube.com/watch?v=VIDEO_ID)
    if host in ('www.youtube.com', 'youtube.com', 'm.youtube.com'):
        if parsed.path == '/watch':
            query_params = parse_qs(parsed.query)
            if 'v' in query_params:
                video_id = query_params['v'][0]
                return len(video_id) == 11
            return False
        # Embedded YouTube URL (e.g., https://www.youtube.com/embed/VIDEO_ID)
        elif parsed.path.startswith('/embed/'):
            video_id = parsed.path.split('/embed/')[1]
            return len(video_id) == 11

    # Shortened YouTube URL (e.g., https://youtu.be/VIDEO_ID)
    elif host == 'youtu.be':
        video_id = parsed.path.lstrip('/')
        return len(video_id) == 11
    return False

# %% ../nbs/02_core.ipynb 10
def mk_part(inp: Union[str, Path, types.Part, types.File, PIL.Image.Image], c: genai.Client|None=None):
    "Turns an input fragment into a multimedia `Part` to be sent to a Gemini model"
    api_client = c or genai.Client(api_key=os.environ["GEMINI_API_KEY"])
    if isinstance(inp, types.Part): return inp
    if isinstance(inp, types.File): return types.Part(file_data={"file_uri": inp.uri, "mime_type": inp.mime_type})
    if isinstance(inp, PIL.Image.Image): return types.Part.from_bytes(data=inp.tobytes(), mime_type=inp.get_format_mimetype())
    if isinstance(inp, bytes):
        mt = mimetypes.types_map["." + imghdr.what(None, h=inp)]
        return types.Part.from_bytes(data=inp, mime_type=mt)
    p_inp = Path(inp)
    if p_inp.exists():
        mt = mimetypes.guess_type(p_inp)[0]
        if mt.split("/")[0] == "image": return types.Part.from_bytes(data=p_inp.read_bytes(), mime_type=mt)
        file = api_client.files.upload(file=p_inp)
        return mk_part(file, c)
    if is_youtube_url(inp): return types.Part.from_uri(file_uri=inp, mime_type='video/*')
    return types.Part.from_text(text=inp)
        

# %% ../nbs/02_core.ipynb 19
def mk_parts(inps, c=None):
    return list(L(inps).map(mk_part, c=c)) if inps else [""]

# %% ../nbs/02_core.ipynb 23
def mk_content(content, role='user', cli=None):
    if isinstance(content, types.Content): return content.model_copy(update={'role': content.role or role})
    if isinstance(content, dict): mk_content(types.Content.model_construct(types.ContentDict(**content)), role)
    c = cli or genai.Client(api_key=os.environ['GEMINI_API_KEY'])
    return types.Content(role=role, parts=mk_parts(content, c=c))

def _is_msg(item):
    if isinstance(item, (types.Content, list)): return True
    if isinstance(item, dict):
        try:
            types.ContentDict(**item)
            return True
        except: return False

def mk_contents(inps, cli=None):
    if not (is_listy(inps) and any(_is_msg(o) for o in inps)): return [mk_content(inps, cli=cli)]
    return [mk_content(o, ('user', 'model')[i % 2], cli) for i, o in enumerate(inps)]

# %% ../nbs/02_core.ipynb 31
def mk_msg(content: list | str | types.Content, role:str='user', *args, api='genai', **kw):
    """Create a `Content` object from the actual content (GenAI's equivalent of a Message)"""
    c = kw.get('client', genai.Client(api_key=os.environ['GEMINI_API_KEY']))
    return mk_content(content, role, cli=c)


def mk_msgs(msgs: list | str, *args, api:str="openai", **kw) -> list:
    "Create a list of messages compatible with the GenAI sdk"
    if isinstance(msgs, str): msgs = [msgs]
    return [mk_msg(o, ('user', 'model')[i % 2], *args, api=api, **kw) for i, o in enumerate(msgs)]

# %% ../nbs/02_core.ipynb 38
@patch(as_prop=True)
def use(self: genai.models.Models): return getattr(self, "_u", usage())

@patch(as_prop=True)
def cost(self: genai.models.Models): return getattr(self, "_cost", 0)


@patch(as_prop=True)
def use(self: genai.Client): return self.models.use

@patch(as_prop=True)
def cost(self: genai.Client): return self.models.cost

# %% ../nbs/02_core.ipynb 39
@patch
def _r(self: genai.models.Models, r):
    """Process a complete model result, storing cost and usage on the `Models` instance."""
    self.result = r
    self.result_content = [nested_idx(r, "candidates", 0, "content")]
    self._u = self.use + getattr(r, "usage_metadata", usage())
    self._cost = self.cost + r.cost
    for func in getattr(self, 'post_cbs', []): func(r)
    return r

@patch(as_prop=True)
def result(self: genai.Client): return nested_idx(self, "models", "result")

@patch(as_prop=True)
def result_content(self: genai.Client): return nested_idx(self, "models", "result_content")

# %% ../nbs/02_core.ipynb 45
@patch(as_prop=True)
def _parts(self: types.GenerateContentResponse): return nested_idx(self, "candidates", 0, "content", "parts") or []
    

@patch
def _stream(self: genai.models.Models, s):
    all_parts = []
    for r in s:
        all_parts.extend(r._parts)
        yield r.text
    r.candidates[0].content.parts = all_parts
    self._r(r)

# %% ../nbs/02_core.ipynb 47
@patch
def _gen(self:genai.models.Models, contents, model:str, config=None, stream:bool=False):
    gen_f = self.generate_content_stream if stream else self.generate_content
    r = gen_f(model=model, contents=contents, config=config)
    return self._stream(r) if stream else self._r(r)

# %% ../nbs/02_core.ipynb 53
@patch
def _genconf(self: genai.models.Models, **kw):
    """Builds a GenerateContentConfigDict from call parameters"""
    config= {k: v for k, v in kw.items() if k in types.GenerateContentConfigDict.__annotations__}
    if _sp := kw.get("sp", False) or kw.get('system_instruction', False) or getattr(self, 'sp', False):
        config['system_instruction'] = _sp 
    if _temp := kw.get("temp", False) or kw.get('temperature', False) or getattr(self, 'temp', False):
        config['temperature'] = _temp
    if maxtok := kw.get("maxtok", False): config['max_output_tokens'] = maxtok
    if stop := kw.get("stop", False): config['stop_sequences'] = [stop] if isinstance(stop, str) else stop

    if tools:= kw.get("tools", False):
        config['tools'] = tools
        tc = config.get('tool_config', dict())
        fcc = tc.get('function_calling_config', dict())
        fcc['mode'] = kw.get("tool_mode", 'AUTO')
        tc['function_calling_config'] = fcc
        config['tool_config']= tc
        
    if model := kw.get('model', None) in imagen_models and not getattr(self, "text_only", False):
        config['response_modalities'] = kw.get('response_modalities', ['Text', 'Image'])
        
    return config

# %% ../nbs/02_core.ipynb 64
def _googlify_docs(fdoc:str,                  # Docstring of a function
                    argdescs: dict|None=None, # Dict of arg:docment of the arguments of the function
                    retd: str|None=None       # Return docoment of the function
                   )-> str: # The function docstring following Google style guide
    """Turns a function docment and docstring into a docstrings that function """
    if argdescs: fdoc += "\n\nArgs:\n"  + "\n".join([f"    {p}: {desc}" for p, desc in argdescs.items()])
    if retd: fdoc += f"\n\nReturns:\n    {retd}"
    return fdoc

def goog_doc(f:callable # A docment style function
            )->str:     # Google style docstring
    """Builds the docstring for a docment style function following Google style guide"""
    fdoc = f.__doc__
    args = {par: doc for par, doc in docments(f, returns=False).items() if doc is not None}
    retd = docments(f, full=True, returns=True)['return']['docment']
    return _googlify_docs(fdoc, args, retd)
    

# %% ../nbs/02_core.ipynb 67
def _geminify(f: callable) -> callable:
    """Makes a function suitable to be turned into a function declaration: 
    infers argument types from default values and removes the values from the signature"""
    docs = docments(f, full=True)
    new_params = [inspect.Parameter(name=n,
                                    kind=inspect.Parameter.POSITIONAL_OR_KEYWORD,
                                    annotation=i.anno) for n, i in docs.items() if n != 'return']

        
    @wraps(f)
    def wrapper(*args, **kwargs):
        return f(*args, **kwargs)
    
    wrapper.__signature__ = inspect.Signature(new_params, return_annotation=docs['return']['anno'])
    wrapper.__annotations__ = {n: i['anno'] for n, i in docs.items() if n != 'return'}
    return wrapper


def prep_tool(f:callable, # The function to be passed to the LLM
             as_decl:bool=False,  # Return an enriched genai.types.FunctionDeclaration?
             googlify_docstring:bool=True): # Use docments to rewrite the docstring following Google Style Guide? 
    """Optimizes a dunction for function calling with the Gemini api. Best suited for docments style functions."""
    docs = goog_doc(f) if googlify_docstring else f.__doc__
    _f = _geminify(f)
    _f.__doc__ = docs
    if not as_decl: return _f
    f_decl = types.FunctionDeclaration.from_callable_with_api_option(callable=_f, api_option='GEMINI_API')
    for par, desc in docments(f, returns=False).items():
        if desc: f_decl.parameters.properties[par].description = desc
    required_params = [p for p, d in docments(f, full=True, returns=False).items() if d['default'] == inspect._empty]
    if hasattr(f_decl.parameters, "required"): f_decl.parameters.required = required_params
    return f_decl

# %% ../nbs/02_core.ipynb 72
def f_result(fname, fargs, ns=None):
    try: return {"result": call_func(fname, fargs, ns or globals())}
    except Exception as e: return {'error': str(e)}

def f_results(fcalls, ns=None):
    return [{"name": c.name, "response": f_result(c.name, c.args, ns)} for c in fcalls]

def mk_fres_content(fres):
    return types.Content(role='tool', parts=[types.Part.from_function_response(**d) for d in fres])

# %% ../nbs/02_core.ipynb 75
@patch
def _call_tools(self: genai.models.Models, r):
    if r.function_calls:
        self.result_content.append(mk_fres_content(f_results(r.function_calls, ns=getattr(self, "_tools", None))))
   

# %% ../nbs/02_core.ipynb 77
def prep_tools(tools, toolify_everything=False):
    funcs = [prep_tool(f, as_decl=toolify_everything) for f in tools if inspect.isfunction(f) or inspect.ismethod(f)]
    if toolify_everything: funcs = [types.Tool(function_declarations=[f]) for f in funcs]
    tools_ = [t for t in tools if isinstance(t, types.Tool)]
    class_tools = [types.Tool(function_declarations=[prep_tool(f, as_decl=True)]) for f in tools if inspect.isclass(f)]
    return funcs + tools_ + class_tools
