"""Chat with local and remote LLMs in the context of the wordslab-notebooks environment"""

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

# %% auto 0
__all__ = ['ChatTurn', 'ChatTurns', 'refresh_notebook_display', 'get_tools_schemas_and_functions', 'ToolExecutionError', 'Tools',
           'Images', 'ModelClient', 'OllamaModelClient', 'OpenRouterModelClient']

# %% ../nbs/02_chat.ipynb 2
from abc import ABC, abstractmethod
from collections.abc import Sequence as SequenceType
from pydantic import BaseModel, ConfigDict
from typing import Callable, Optional, Union, Sequence, Mapping, Any, Literal, Type

from base64 import b64decode, b64encode
from datetime import datetime
import inspect, json, re, time, traceback
from pathlib import Path
from urllib.parse import urlparse
from urllib.request import urlopen

from toolslm.funccall import get_schema, call_func
from IPython.display import display, Markdown, clear_output

from ollama import Client
from ollama._types import Options
from openai import OpenAI

from .env import WordslabEnv

# %% ../nbs/02_chat.ipynb 4
class ChatTurn:
    def __init__(self, refresh_display: Callable = None):
        self.thinking_chunks = []
        self.content_chunks = []
        self.tool_calls = {}
        self.refresh_display = refresh_display

    def append_thinking(self, chunk:str):
        self.thinking_chunks.append(chunk)
        if self.refresh_display:
            self.refresh_display()

    @property
    def thinking(self):
        return "".join(self.thinking_chunks)
    
    def append_content(self, chunk:str):
        self.content_chunks.append(chunk)
        if self.refresh_display:
            self.refresh_display()

    @property
    def content(self):
        return "".join(self.content_chunks)
    
    def append_tool_call(self, tool_name:str, params:dict):
        self.tool_calls[tool_name] = { "params": params }
        if self.refresh_display:
            self.refresh_display()

    def start_tool_call(self, tool_name:str):
        self.tool_calls[tool_name]["start_time"] = time.time()
        if self.refresh_display:
            self.refresh_display()

    def end_tool_call(self, tool_name:str, result: object):
        self.tool_calls[tool_name]["end_time"] = time.time()
        self.tool_calls[tool_name]["result"] = str(result)
        if self.refresh_display:
            self.refresh_display()

    def to_markdown(self, hide_thinking:bool=True, hide_tool_calls:bool=True):
        output = ""
        if len(self.thinking_chunks) > 0:
            if not hide_thinking:
                output += "> [Thinking]\n\n"
                output += "\n".join(f"> {line}" for line in "".join(self.thinking_chunks).splitlines()) + "\n\n"
            else:
                output += f"> [Thinking] ... thought in {sum(s.count(' ') + s.count('\n') for s in self.thinking_chunks)} words\n\n"
        if len(self.content_chunks) > 0:
            output += "".join(self.content_chunks) + "\n\n"
        if len(self.tool_calls.keys()) > 0:
            if not hide_tool_calls:
                for tool_name in self.tool_calls.keys():
                    output += "> [Tool call]\n"
                    tool_call = self.tool_calls[tool_name]
                    if "params" in tool_call:
                        output += f"> - model wants to call `{tool_name}` with parameters `{tool_call["params"]}`\n"
                    if "start_time" in tool_call:
                        output += f"> - agent called {tool_name} at {datetime.fromtimestamp(tool_call["start_time"]).strftime("%H:%M:%S")}\n"
                    if "end_time" in tool_call:
                        result = tool_call["result"]
                        output += f"> - {tool_name} returned `{result if len(result)<=100 else result[:97]+'...'}` in {(tool_call["end_time"]-tool_call["start_time"]):.3f} sec\n"
                    output += "\n"
            else:
                for tool_name in self.tool_calls.keys():
                    tool_call = self.tool_calls[tool_name]
                    if "end_time" in tool_call:
                        result = self.tool_calls[tool_name]["result"]
                        output += f"> [Tool call] ... `{tool_name}` returned `{result if len(result)<=50 else result[:47]+'...'}`\n\n"
                    elif "start_time" in tool_call:                        
                        output += f"> [Tool call] ... agent is calling `{tool_name}`\n\n"
                    elif "params" in tool_call:                        
                        output += f"> [Tool call] ... model wants to call `{tool_name}`\n\n"
        return output
 

class ChatTurns:
    def __init__(self, refresh_display:Callable = None):
        self.chat_turns = []
        self.refresh_display = refresh_display(self.chat_turns)

    def new_turn(self):
        new_turn = ChatTurn(self.refresh_display)
        self.chat_turns.append(new_turn)
        return new_turn

    
def refresh_notebook_display(hide_thinking:bool=True, hide_tool_calls:bool=True):
    def _refresh_turns(chat_turns: ChatTurns):
        def _refresh():
            clear_output(wait=True)
            output = ""
            for turn in chat_turns:
                output += turn.to_markdown(hide_thinking, hide_tool_calls)
            display(Markdown(output))
        return _refresh
    return _refresh_turns

# %% ../nbs/02_chat.ipynb 14
def _get_function_schema(func: Callable, responsesAPIFormat: bool = False):
    "Get a json schema for a python function defined with comments for all parameters"
    if responsesAPIFormat:        
        return {'type': 'function', **get_schema(func, pname='parameters')}
    else:
        return {'type': 'function', 'function': get_schema(func, pname='parameters')}

def get_tools_schemas_and_functions(funcs: Sequence[Callable], responsesAPIFormat: bool = False):
    """Get a dictionary of json schemas and callable functions which can be used for native tool calling."""
    return {func.__name__: (_get_function_schema(func, responsesAPIFormat), func) for func in funcs}

# %% ../nbs/02_chat.ipynb 17
class ToolExecutionError(Exception):
    """Raised when a tool cannot be executed safely."""

class Tools:
    """"Execute tools implemented as python functions with Large Language Models.
    The python functions must be fully documented:
    - type annotations are mandatory on all parameters and on the return type
    - a docstring after the function definition is mandatory
    - a descriptive comment after each parameter and the return type is also mandatory
    - the expected format is: one parameter by line, a traditional python comment at the end of the line
    """   
    def __init__(self, python_functions:Sequence[Callable], responsesAPIFormat:bool=False):
        self.schemas_and_functions = get_tools_schemas_and_functions(python_functions, responsesAPIFormat=responsesAPIFormat)

    def has_tool(self, tool_name:str):
        return tool_name in self.schemas_and_functions
    
    def get_schemas(self):
        return [t[0] for t in self.schemas_and_functions.values()]

    def get_schema(self, tool_name:str):
        return self.schemas_and_functions[tool_name][0]

    def get_functions(self):
        return [t[1] for t in self.schemas_and_functions.values()]

    def get_function(self, tool_name:str):
        return self.schemas_and_functions[tool_name][1]

    def call(self, tool_name:str, tool_arguments_dict:Mapping[str,Any]):
        # 1. Resolve the tool safely
        try:
            self.get_function(tool_name)
        except Exception as e:
            raise ToolExecutionError(f"Tool '{tool_name}' does not exist or could not be resolved.")
            
        # 2. Execute the tool with runtime protection
        try:
            return call_func(tool_name, tool_arguments_dict, self.get_functions(), raise_on_err=True)
        except Exception as e:
            tb = traceback.format_exc()
            raise ToolExecutionError(f"Tool '{tool_name}' raised an exception: {e}")

# %% ../nbs/02_chat.ipynb 26
class Images:
    def __init__(self, image:Union[str,Path,bytes] = None):
        self.images = []
        self.urls = []
        
        if isinstance(image, bytes):
            self.add_bytes(image)
        elif isinstance(image, Path):
            self.add_file_path(image)
        elif isinstance(image, str):
            if Images.is_web_url(image):
                self.add_web_url(image)
            else:
                self.add_file_path(image)
        else:
            raise ValueError(f'Image must be a string, Path or bytes object')

    def add_bytes(self, image_bytes: bytes):
        self.images.append(b64encode(image_bytes).decode())
        self.urls.append(False)
    
    def add_file_path(self, image_path: Union[str, Path]):
        if isinstance(image_path, str):
            image_path = Path(image_path)
        if isinstance(image_path, Path):
            if image_path.is_file():
                self.add_bytes(image_path.read_bytes())
            else:
                raise ValueError(f'Invalid image file path: {image_path}')
        else:
            raise ValueError(f'Image file path must be a string or a Path object')

    def is_web_url(image_url: str):
        try:
            parsed = urlparse(image_url)
            return parsed.scheme in ("http", "https") and bool(parsed.netloc)
        except Exception:
            return False
    
    def add_web_url(self, image_url: str): 
        if Images.is_web_url(image_url):
            self.images.append(image_url)
            self.urls.append(True)
        else:
            raise ValueError(f'Invalid image web url: {image_url}')

    def get_base64_data_or_url(self):
        return zip(self.images,self.urls)

    def get_base64_data(self):
        for i in range(len(self.images)):
            if not self.urls[i]:
                yield self.images[i]
            else:
                with urlopen(self.images[i], timeout=30) as response:
                    yield b64encode(response.read()).decode()

# %% ../nbs/02_chat.ipynb 33
class ModelClient(ABC):
    def __init__(
        self,
        model: str,        
        context_size: Optional[int] = None,
        base_url: Optional[str] = None,
        api_key: Optional[str] = None,
    ):
        self.model = model
        self.base_url = base_url
        self.api_key = api_key
        self.context_size = context_size

    @abstractmethod
    def __call__(
        self,
        messages: Sequence[Mapping[str, Any]] = None,
        system_prompt: str = None,
        user_prompt: str = None,
        user_images: Images = None,
        assistant_prefill: str = None,
        tools: Tools = None,             
        max_tool_calls:int = 20,
        think: Union[bool, Literal["xhigh", "high", "medium", "low", "minimal", "none"], int, None] = None,
        web_search: bool = False,
        output_model: Type[BaseModel] = None,
        max_new_tokens: Optional[int] = None,
        seed: Optional[int] = None,
        temperature: Optional[float] = None,
        top_k: Optional[int] = None,
        top_p: Optional[float] = None,
        min_p: Optional[float] = None, 
        refresh_display: Optional[Callable] = refresh_notebook_display()
    ) -> None:
        """
        Execute a model call and return the model response.
        """
        raise NotImplementedError

    def _check_call_parameters_type(self, user_images, tools, output_model):
        if user_images and not isinstance(user_images, Images):
            raise TypeError("Argument user_images must be of type wordslab_notebooks_lib.chat.Images. Create an images object with the syntax: Images('file.jpg') or Images('https://example.com/file.jpg') or Images(image_bytes).")
        if tools and not isinstance(tools, Tools):
            raise TypeError("Argument tools must be of type wordslab_notebooks_lib.chat.Tools. Create a tools object with the syntax: Tools([func1, func2, func3]), where the parameters are documented python functions.")
        if output_model and not (isinstance(output_model, type) and issubclass(output_model, BaseModel)):
            raise TypeError("Argument output_model must be a pydantic model - a subclass of pydantic.BaseModel.")

    def _manage_messages(self, messages, system_prompt, user_prompt, images_value, assistant_prefill):
        if system_prompt:
            system_message = {'role': 'system', 'content': system_prompt}
            if len(messages) == 0:                
                messages.append(system_message) 
            elif messages[0]['role'] == 'system':
                messages[0] = system_message
            else:
                messages.insert(0, system_message)
        if user_prompt or images_value:
            ollama_format = False
            if images_value:
                ollama_format = isinstance(images_value[0], str)
            user_message = {'role': 'user'}  
            if (not images_value or ollama_format):
                if user_prompt:
                    user_message['content'] = user_prompt
                if images_value:
                    user_message['images'] = images_value
            else:
                content_list = []
                if user_prompt:  
                    content_list.append({"type": "text", "text": user_prompt})
                for image_data,is_url in images_value:
                    if is_url:
                        content_list.append({"type": "image_url", "image_url": { "url": image_data }})
                    else:
                        content_list.append({"type": "image_url", "image_url": { "url":f"data:image/jpeg;base64,{image_data}" }})
                user_message['content'] = content_list
            messages.append(user_message)
        if assistant_prefill:
            messages.append({'role': 'assistant', 'content': assistant_prefill}) 

# %% ../nbs/02_chat.ipynb 35
_whitespace_pattern = re.compile(r"\s+")

def _messages_words(messages):
    return sum([len(_whitespace_pattern.findall(message["content"][0]["text"])) if isinstance(message["content"], list) else len(_whitespace_pattern.findall(message["content"])) for message in messages if message["role"] in {"system", "user", "assistant", "tool"}])

class OllamaModelClient(ModelClient):
    def __init__(
        self,
        model: str,
        context_size: int = 32768, # This is the default value for the ollama server in wordslab-notebooks
        base_url: str = "http://localhost:11434",
        api_key: Optional[str] = None,  # If not provided, the optional key will be pulled from WordslabEnv
    ):
        super().__init__(model, context_size, base_url, api_key)

        # Initialize API client
        if not api_key:
            env = WordslabEnv()
            api_key = env.cloud_ollama_api_key
        if api_key:
            self.has_api_key = True
            headers = {'Authorization': 'Bearer ' + api_key}
        else:
            self.has_api_key = False
            headers = {}
        self.client = Client(host=self.base_url, headers=headers)
        
        # Load model in memory with the right context length
        print(f"ollama: loading model {self.model} with context size {self.context_size} ... ", end="");
        self.client.chat(model=self.model, messages=[{'role': 'user', 'content': 'say yes'}], options=Options(num_ctx=self.context_size, num_predict=1))
        print(f"ok");

    def __call__(
        self,
        messages: Sequence[Mapping[str, Any]] = None,
        system_prompt: str = None,
        user_prompt: str = None,
        user_images: Images = None,
        assistant_prefill: str = None,
        tools: Tools = None,             
        max_tool_calls:int = 20,
        think: Union[bool, Literal["low", "medium", "high"], None] = None,
        web_search: bool = False,
        output_model: Type[BaseModel] = None,
        max_new_tokens: Optional[int] = None,
        seed: Optional[int] = None,
        temperature: Optional[float] = None,
        top_k: Optional[int] = None,
        top_p: Optional[float] = None,
        min_p: Optional[float] = None, 
        refresh_display: Optional[Callable] = refresh_notebook_display()
    ) -> None:
        # Reset response field
        self.response = None
        
        # Check parameters type
        super()._check_call_parameters_type(user_images, tools, output_model)

        # Images
        images_value = None
        if user_images:
            images_value = list(user_images.get_base64_data())

        # Messages
        if messages is None:
            messages = []
        super()._manage_messages(messages, system_prompt, user_prompt, images_value, assistant_prefill)        

        # Web search
        if web_search:
            if self.has_api_key:
                tools = Tools([self.client.web_search, self.client.web_fetch] + (tools if tools else []))
            else:
                raise RuntimeError("Web search requires an api_key in OllamaModelClient constructor.")
                
        chat_turns = ChatTurns(refresh_display)
        
        for i in range(max_tool_calls):
            # Immediate user feedback
            print(f"ollama: processing {_messages_words(messages)} words with `{self.model}` ...")
            
            # Observable conversation turn
            chat_turn = chat_turns.new_turn()
            if assistant_prefill and len(chat_turns.chat_turns) == 1:
                chat_turn.append_content(assistant_prefill)
            
            stream = self.client.chat(
                model = self.model,
                messages = messages,
                tools = tools.get_schemas() if tools else None,
                stream = True,
                think = think,
                format = output_model.model_json_schema() if output_model else None,
                options = Options(num_ctx = self.context_size, num_predict = max_new_tokens, seed = seed,
                                  temperature = temperature, top_k=top_k, top_p=top_p, min_p=min_p)
            )
        
            # Streaming: accumulate the partial fields
            tool_calls = []        
            for chunk in stream:
                if chunk.message.thinking:
                    chat_turn.append_thinking(chunk.message.thinking)                
                if chunk.message.content:
                    chat_turn.append_content(chunk.message.content)
                if chunk.message.tool_calls:
                    tool_calls.extend(chunk.message.tool_calls)
                    for tc in chunk.message.tool_calls:
                        chat_turn.append_tool_call(tc.function.name, tc.function.arguments)
            
            # append accumulated fields to the messages
            if chat_turn.thinking or chat_turn.content or tool_calls:
                messages.append({'role': 'assistant', 'thinking': chat_turn.thinking, 'content': chat_turn.content, 'tool_calls': tool_calls})
        
            # end the loop if there is no more tool calls
            if not tool_calls: 
                # Set response field
                if 'content' in messages[-1]:
                    self.response = messages[-1]['content']
                if output_model:
                    self.response = output_model.model_validate_json(self.response)
                return 
                
            # execute tool calls  
            else:    
                for tc in tool_calls:
                    if tools.has_tool(tc.function.name):
                        chat_turn.start_tool_call(tc.function.name)
                        result = tools.call(tc.function.name, tc.function.arguments)
                        chat_turn.end_tool_call(tc.function.name, result)
                    else:
                        result = 'Unknown tool'
            
                    # append tool call result to the messages 
                    messages.append({'role': 'tool', 'tool_name': tc.function.name, 'content': str(result)})
    
            # continue the loop after tool calls

# %% ../nbs/02_chat.ipynb 63
class OpenRouterModelClient(ModelClient):
    def __init__(
        self,
        model: str,
        context_size: Optional[int] = None, # For OpenRouter this parameter is ignored, we inherit the remote model config
        base_url: str = "https://openrouter.ai/api/v1",
        api_key: Optional[str] = None, # If not provided, the mandatory key will be pulled from WordslabEnv
    ):
        super().__init__(model, context_size, base_url, api_key)

        # Initialize API client
        if not api_key:
            env = WordslabEnv()
            api_key = env.cloud_openrouter_api_key
        self.client = OpenAI(base_url=base_url, api_key=api_key)
        
        # Check connection
        print(f"openrouter: testing model {self.model} ... ", end="");
        self.client.chat.completions.create(model=self.model, messages=[{'role': 'user', 'content': 'say yes'}], max_tokens=16)
        print(f"ok");

    def __call__(
        self,
        messages: Sequence[Mapping[str, Any]] = None,
        system_prompt: str = None,
        user_prompt: str = None,
        user_images: Images = None,
        assistant_prefill: str = None,
        tools: Tools = None,             
        max_tool_calls:int = 20,
        think: Union[bool, Literal["xhigh", "high", "medium", "low", "minimal", "none"], int, None] = None,
        web_search: bool = False,
        output_model: Type[BaseModel] = None,
        max_new_tokens: Optional[int] = None,
        seed: Optional[int] = None,
        temperature: Optional[float] = None,
        top_k: Optional[int] = None,
        top_p: Optional[float] = None,
        min_p: Optional[float] = None, 
        refresh_display: Optional[Callable] = refresh_notebook_display()
    ) -> None:
        # Reset response field
        self.response = None
        self.annotations = None
        
        # Check parameters type
        super()._check_call_parameters_type(user_images, tools, output_model)

        # Images
        images_value = None
        if user_images:
            images_value = list(user_images.get_base64_data_or_url())
        
        # Messages
        if messages is None:
            messages = []
        super()._manage_messages(messages, system_prompt, user_prompt, images_value, assistant_prefill)

        # Map "think" → reasoning > enabled / reasoning > effort / resoning > max_tokens
        reasoning = None
        if think is True:
            reasoning = {"reasoning": {"enabled": True}} # reasoning on/off
        elif think in ("low", "medium", "high"):
            reasoning = {"reasoning": {"effort": think}} # "xhigh", "high", "medium", "low", "minimal" or "none" (OpenAI-style)
        elif isinstance(think, int):
            reasoning = {"reasoning": {"max_tokens": think}} # specified token budget for extended thinking (Anthropic-style)

        # Structured outputs
        response_format = None
        if output_model:
            schema = output_model.model_json_schema()
            schema["additionalProperties"] = False
            response_format = {
                    "type": "json_schema",
                    "json_schema": {
                        "name": output_model.__name__,
                        "schema": schema,
                        "strict": True,
                    },
                }
        
        chat_turns = ChatTurns(refresh_display)
        
        for i in range(max_tool_calls):
            # Immediate user feedback
            print(f"openrouter: processing {_messages_words(messages)} words with `{self.model + (":online" if web_search else "")}` ...")
            if web_search:
                print("-> web search activated")
            
            # Observable conversation turn
            chat_turn = chat_turns.new_turn()
            if assistant_prefill and len(chat_turns.chat_turns) == 1:
                chat_turn.append_content(assistant_prefill)            
            
            stream = self.client.chat.completions.create(
                model = self.model + (":online" if web_search else ""),
                messages = messages,
                tools = tools.get_schemas() if tools else None,
                stream = True,
                extra_body = reasoning,
                response_format = response_format,
                max_tokens = max_new_tokens,
                seed = seed,
                temperature = temperature,
                top_p = top_p,
            )
        
            # Streaming: accumulate the partial fields
            tool_calls = {} 
            annotations = None
            for chunk in stream:
                delta = chunk.choices[0].delta
                if hasattr(delta, "reasoning") and delta.reasoning:
                    chat_turn.append_thinking(delta.reasoning)                
                if hasattr(delta, "content") and delta.content:
                    chat_turn.append_content(delta.content)
                # web search results - standardized by openrouter
                if hasattr(delta, "annotations") and delta.annotations:
                    annotations = delta.annotations
                if hasattr(delta, "tool_calls") and delta.tool_calls:
                    for tool_call in delta.tool_calls:
                        idx = tool_call.index
                        # First tool call chunk
                        if idx not in tool_calls:
                            tool_calls[idx] = {
                                "id": tool_call.id,  # only present in first chunk
                                "name": tool_call.function.name,
                                "arguments": ""
                            }       
                        # Append streamed argument fragments
                        tool_calls[idx]["arguments"] += (
                            tool_call.function.arguments or ""
                        )
                        
            # We need to wait the end of the stream to make sure the tool calls are complete
            for tc in tool_calls.values():
                chat_turn.append_tool_call(tc["name"], tc["arguments"])
            
            # Append accumulated fields to the messages
            if chat_turn.content or tool_calls:
                messages.append({
                    "role": "assistant",
                    "content": chat_turn.content,
                    "tool_calls": [
                        {
                            "id": tc["id"],
                            "type": "function",
                            "function": {
                                "name": tc["name"],
                                "arguments": tc["arguments"]
                            }
                        }
                        for tc in tool_calls.values()
                    ]
                })
                
            # end the loop if there is no more tool calls
            if not tool_calls: 
                # Set response field
                if 'content' in messages[-1]:
                    self.response = messages[-1]['content']
                if annotations:
                    self.annotations = annotations
                if output_model:
                    self.response = output_model.model_validate_json(self.response)
                return      
                
            # execute tool calls  
            else:    
                for tc in tool_calls.values():
                    if tools.has_tool(tc["name"]):
                        chat_turn.start_tool_call(tc["name"])
                        result = tools.call(tc["name"], json.loads(tc["arguments"]))
                        chat_turn.end_tool_call(tc["name"], result)
                    else:
                        result = 'Unknown tool'
            
                    # append tool call result to the messages 
                    messages.append({"role": "tool", "tool_call_id": tc["id"], "content": str(result)})
    
            # continue the loop after tool calls
