"""Agent-friendly wrapper around and LLM instance"""

# AUTOGENERATED! DO NOT EDIT! File to edit: ../../../nbs/04_pipelines.agent.model.ipynb.

# %% auto 0
__all__ = ['AgentModel']

# %% ../../../nbs/04_pipelines.agent.model.ipynb 3
import re

from enum import Enum

from smolagents import ChatMessage, Model, get_clean_message_list, tool_role_conversions
from smolagents.models import get_tool_call_from_text, remove_stop_sequences
import onprem

# %% ../../../nbs/04_pipelines.agent.model.ipynb 4
class AgentModel(Model):
    """
    A smolagents Model implementation that wraps an onprem LLM instance.
    
    This adapter allows onprem LLM instances to be used with smolagents Agents.
    
    Parameters:
        llm (LLM): An instance of onprem.llm.base.LLM
        **kwargs: Additional keyword arguments to pass to the parent Model class
    """
    
    def __init__(
        self,
        llm: onprem.LLM,
        **kwargs
    ):
        self.llm = llm
        self.maxlength=8192
        self.model_id = self.llm.model_name
        super().__init__(
            flatten_messages_as_text=True,
            **kwargs
        )
        
    def generate(
        self,
        messages,
        stop_sequences=None,
        response_format=None, 
        tools_to_call_from=None,
        **kwargs
        ):
        """
        Process the input messages and return the model's response.
        
        Parameters:
            messages: A list of message dictionaries to be processed.
            stop_sequences: A list of strings that will stop generation if encountered.
            response_format: The response format to use in the model's response.
            tools_to_call_from: A list of tools that the model can use.
            **kwargs: Additional keyword arguments to pass to the LLM.
            
        Returns:
            ChatMessage: A chat message object containing the model's response.
        """
        # Convert smolagents messages to a format that onprem LLM can use
        messages = self.clean(messages)
        
        # Call the LLM with the processed messages
        #import litellm as api
        #result = api.completion(
                #model='ollama_chat/llama3.1:8b',
                #messages=messages,
                #max_tokens=self.maxlength,
                #stream=False,
                #stop=stop_sequences,
            #)
        #data = result["choices"][0]
        #text = data.get("text", data.get("message", data.get("delta")))
        #text = text if isinstance(text, str) else text.get("content")
        #response = text

        response = self.llm.prompt(
            messages,
            stop=stop_sequences or [],
            max_tokens = self.maxlength,
            **kwargs
        )
        #print(f'RESPONSE: {response}')

        # Remove stop sequences from LLM output
        if stop_sequences is not None:
            response = remove_stop_sequences(response, stop_sequences)
        
        # Create and return a ChatMessage with the response
        message =  ChatMessage(role="assistant", content=response)

        # Extract first tool action
        if tools_to_call_from:
            # Try different patterns to extract JSON
            extracted_json = None
            
            # Helper function to extract balanced JSON
            def extract_balanced_json(text, start_pos):
                brace_count = 0
                in_string = False
                escape_next = False
                
                for i, char in enumerate(text[start_pos:], start_pos):
                    if escape_next:
                        escape_next = False
                        continue
                    
                    if char == '\\' and in_string:
                        escape_next = True
                        continue
                    
                    if char == '"' and not escape_next:
                        in_string = not in_string
                        continue
                    
                    if not in_string:
                        if char == '{':
                            brace_count += 1
                        elif char == '}':
                            brace_count -= 1
                            if brace_count == 0:
                                return text[start_pos:i+1]
                
                return None
            
            # Pattern 1: Action: {...}
            match = re.search(r"Action:\s*(\{)", response, flags=re.DOTALL)
            if match:
                start_pos = match.start(1)
                extracted_json = extract_balanced_json(response, start_pos)
            else:
                # Pattern 2: Called Tool: 'name' with arguments: {...}
                match = re.search(r"Called Tool: '([^']+)' with arguments:\s*(\{)", response, flags=re.DOTALL)
                if match:
                    tool_name = match.group(1)
                    start_pos = match.start(2)
                    args_json = extract_balanced_json(response, start_pos)
                    if args_json:
                        extracted_json = f'{{"name": "{tool_name}", "arguments": {args_json}}}'
                else:
                    # Pattern 3: Just the raw Called Tool format without proper JSON structure
                    match = re.search(r"Called Tool: '([^']+)' with arguments:\s*(.+?)(?=\n|$)", response, flags=re.DOTALL)
                    if match:
                        tool_name, args_str = match.groups()
                        # Try to fix the arguments format
                        args_str = args_str.strip()
                        if args_str.startswith('{') and args_str.endswith('}'):
                            extracted_json = f'{{"name": "{tool_name}", "arguments": {args_str}}}'
            
            if extracted_json:
                #print(f'EXTRACTED JSON: {repr(extracted_json)}')
                # Fix JSON formatting: replace single quotes around keys/values but preserve apostrophes
                import json
                try:
                    # Try to parse as-is first
                    json.loads(extracted_json)
                except json.JSONDecodeError:
                    # Try a different approach - use ast.literal_eval style fixing
                    import ast
                    try:
                        # Replace single quotes with double quotes more carefully
                        # First fix obvious key patterns
                        extracted_json = re.sub(r"'(\w+)':", r'"\1":', extracted_json)
                        
                        # Try to parse with Python's literal_eval which is more forgiving
                        parsed = ast.literal_eval(extracted_json)
                        # Convert back to JSON string
                        import json
                        extracted_json = json.dumps(parsed)
                    except (ValueError, SyntaxError):
                        # Fallback: manually escape problematic characters
                        # Fix single quotes to double quotes for JSON structure
                        lines = extracted_json.split('\n')
                        fixed_lines = []
                        for line in lines:
                            # Don't mess with the overall JSON structure
                            if ':' in line and not line.strip().startswith('"'):
                                # This looks like a key-value pair, escape the value part
                                key_part, value_part = line.split(':', 1)
                                value_part = value_part.strip()
                                # If it starts with a quote, escape internal quotes and newlines
                                if value_part.startswith("'") and value_part.endswith("'"):
                                    value_content = value_part[1:-1]  # Remove outer quotes
                                    value_content = value_content.replace('\\', '\\\\').replace('"', '\\"').replace('\n', '\\n')
                                    line = f'{key_part}: "{value_content}"'
                                elif value_part.startswith('"') and value_part.endswith('"'):
                                    value_content = value_part[1:-1]  # Remove outer quotes
                                    value_content = value_content.replace('\\', '\\\\').replace('"', '\\"').replace('\n', '\\n')
                                    line = f'{key_part}: "{value_content}"'
                            fixed_lines.append(line)
                        extracted_json = '\n'.join(fixed_lines)
                    #print(f'FIXED JSON: {repr(extracted_json)}')
                
                message.tool_calls = [
                    get_tool_call_from_text(
                        extracted_json, self.tool_name_key, self.tool_arguments_key
                    )
                ]

        return message


    def parameters(self, maxlength):
        """
        Set LLM inference parameters.

        Args:
            maxlength: maximum sequence length
        """

        self.maxlength = maxlength


    def clean(self, messages):
        """
        Gets a clean message list.

        Args:
            messages: input messages

        Returns:
            clean messages
        """

        # Get clean message list
        messages = get_clean_message_list(messages, role_conversions=tool_role_conversions, flatten_messages_as_text=self.flatten_messages_as_text)

        # Ensure all roles are strings and not enums for compability across LLM frameworks
        for message in messages:
            if "role" in message:
                message["role"] = message["role"].value if isinstance(message["role"], Enum) else message["role"]

        return messages

