import openai
import json
import re
import os

# For the python function to JSON/dict
import inspect
from typing import Any, Callable, Literal, Union, get_args, get_origin

__version__ = "0.2.2"

import requests
from packaging import version

if True : # You can disable it btw
    try:
        response = requests.get("https://pypi.org/pypi/open-taranis/json", timeout=0.1)
        response.raise_for_status()
        latest_version = response.json()["info"]["version"]
        if version.parse(latest_version) > version.parse(__version__):
            print(f'New (stable) version {latest_version} available for open-taranis !\nUpdate via "pip install -U open-taranis"')
    except Exception:
        pass

# ==============================
# 
# ==============================

class utils:
    def _parse_simple_docstring(doc: str | None) -> dict[str, Any]:
        """Parse docstring minimal (description + args)."""
        result = {"description": "", "args": {}}
        if not doc:
            return result
        
        # Extract main description (first paragraph)
        parts = inspect.cleandoc(doc).split('\n\n', 1)
        result["description"] = parts[0].strip()
        
        # Simple args parsing (Google/NumPy style)
        if len(parts) > 1:
            args_section = parts[1].split('Args:')[-1].split('Returns:')[0].split('Raises:')[0]
            lines = [l.strip() for l in args_section.split('\n') if l.strip()]
            
            for line in lines:
                if ':' in line and not line.startswith(' '):
                    # Format: "arg_name: description" or "arg_name (type): description"
                    arg_match = re.match(r'(\w+)\s*(?:\([^)]*\))?\s*:\s*(.+)', line)
                    if arg_match:
                        arg_name, desc = arg_match.groups()
                        result["args"][arg_name] = desc.strip()
        
        return result

    def _python_type_to_schema(py_type: Any) -> dict[str, Any]:
        """Convert Python type to JSON Schema - MINIMAL version."""
        origin = get_origin(py_type)
        args = get_args(py_type)
        
        # Optional: Union[X, None]
        if origin is Union and type(None) in args:
            non_none = [a for a in args if a is not type(None)]
            if len(non_none) == 1:
                schema = utils._python_type_to_schema(non_none[0])
                schema["nullable"] = True
                return schema
        
        # Literal for enums
        if origin is Literal:
            return {"type": "string", "enum": list(args)}
        
        # Basic types
        if py_type in (str, int, float, bool, type(None)):
            type_map = {str: "string", int: "integer", float: "number", bool: "boolean", type(None): "null"}
            return {"type": type_map[py_type]}
        
        # Collections
        if origin in (list,):
            item_schema = {"type": "string"}  # Default
            if args:
                item_schema = utils._python_type_to_schema(args[0])
            return {"type": "array", "items": item_schema}
        
        if origin in (dict,):
            return {"type": "object"}
        
        # Default fallback
        return {"type": "string"}

    def function_to_openai_tool(func: Callable) -> dict[str, Any]:
        """Convert Python function to OpenAI tool format - MINIMAL."""
        sig = inspect.signature(func)
        type_hints = func.__annotations__
        
        # Parse docstring
        doc_info = utils._parse_simple_docstring(func.__doc__ or "")
        
        # Build schema
        properties = {}
        required = []
        
        for param_name, param in sig.parameters.items():
            # Get type annotation
            py_type = type_hints.get(param_name, str)
            schema = utils._python_type_to_schema(py_type)
            
            # Add description from docstring
            if param_name in doc_info["args"]:
                schema["description"] = doc_info["args"][param_name]
            
            # Handle defaults
            if param.default is not inspect.Parameter.empty:
                schema["default"] = param.default
                if param.default is None:
                    schema["nullable"] = True
            else:
                required.append(param_name)
            
            properties[param_name] = schema
        
        return {
            "type": "function",
            "function": {
                "name": func.__name__,
                "description": doc_info["description"],
                "parameters": {
                    "type": "object",
                    "properties": properties,
                    "required": required,
                    "additionalProperties": False
                }
            }
        }

    @staticmethod
    def generate_clients(api_key_env:str, base_url:str, **kwargs):
        def client_builder(api_key:str=None):
            api_key = os.environ.get(api_key_env) if api_key_env else api_key
                
            return openai.OpenAI(api_key=api_key, base_url=base_url, default_headers=kwargs.get("default_headers"))

        return client_builder


class clients:

# ==============================
# The clients with their URL
# ==============================

    @staticmethod
    def generic(api_key:str, base_url:str) -> openai.OpenAI:
        """
        Use `clients.generic_request` for call
        """
        return openai.OpenAI(api_key=api_key, base_url=base_url)
    
    veniceai = utils.generate_clients('VENICE_API_KEY', "https://api.venice.ai/api/v1")
    
    deepseek = utils.generate_clients('DEEPSEEK_API_KEY',"https://api.deepseek.com")

    xai = utils.generate_clients('XAI_API_KEY', "https://api.x.ai/v1")

    groq = utils.generate_clients('GROQ_API_KEY', "https://api.groq.com/openai/v1")

    huggingface = utils.generate_clients('HUGGINGFACE_API_KEY', "https://router.huggingface.co/v1")
    
    ollama = utils.generate_clients(None, "http://localhost:11434/v1")

    openrouter = utils.generate_clients('OPENROUTER_API_KEY', "https://openrouter.ai/api/v1")

    kimi_code = utils.generate_clients("KIMI_CODE_API_KEY", "https://api.kimi.com/coding/v1",
        default_headers={"User-Agent": "RooCode/3.30.3","HTTP-Referer": "https://github.com/RooVetGit/Roo-Cline","X-Title": "Roo Code"}
    )
    
# ==============================
# Customers for calls with their specifications"
#
# Like "include_venice_system_prompt" for venice.ai or custom app for openrouter
# ==============================

    @staticmethod
    def generic_request(client: openai.OpenAI, messages: list[dict], model:str="defaut", temperature:float=0.7, max_tokens:int=4096, tools:list[dict]=None, **kwargs) -> openai.Stream:
        base_params = {
            "model": model,
            "messages": messages,
            "temperature": temperature,
            "max_tokens": max_tokens,
            "max_completion_tokens": kwargs.get("max_completion_tokens", 4096),
            "stream": kwargs.get("stream", True),
        }
        
        tool_params = {}
        if tools :
            tool_params = {
                "tools": tools,
                "tool_choice": kwargs.get("tool_choice", "auto")
            }
        
        params = {**base_params, **tool_params}
        
        return client.chat.completions.create(**params)

    @staticmethod
    def veniceai_request(client: openai.OpenAI, messages: list[dict], 
                        model:str="venice-uncensored", temperature:float=0.7, max_tokens:int=4096, tools:list[dict]=None, 
                        include_venice_system_prompt:bool=False, 
                        enable_web_search:bool=False,
                        enable_web_citations:bool=False,
                        disable_thinking:bool=False,
                        **kwargs) -> openai.Stream:
        base_params = {
            "model": model,
            "messages": messages,
            "temperature": temperature,
            "max_tokens": max_tokens,
            "max_completion_tokens": kwargs.get("max_completion_tokens", 4096),
            "stream": kwargs.get("stream", True),
        }
        
        tool_params = {}
        if tools :
            tool_params = {
                "tools": tools,
                "tool_choice": kwargs.get("tool_choice", "auto")
            }
        
        venice_params = {
            "extra_body": {
                "venice_parameters": {
                    "include_venice_system_prompt" : include_venice_system_prompt,
                    "enable_web_search" : "on" if enable_web_search else "off",
                    "enable_web_citations" : enable_web_citations,
                    "disable_thinking" : disable_thinking
                }
            }
        }
        
        params = {**base_params, **tool_params, **venice_params}
        
        return client.chat.completions.create(**params)

    @staticmethod
    def openrouter_request(client: openai.OpenAI, messages: list[dict], model:str="nvidia/nemotron-nano-9b-v2:free", temperature:float=0.7, max_tokens:int=4096, tools:list[dict]=None, **kwargs) -> openai.Stream:
        base_params = {
            "model": model,
            "messages": messages,
            "temperature": temperature,
            "max_tokens": max_tokens,
            "max_completion_tokens": kwargs.get("max_completion_tokens", 4096),
            "stream": kwargs.get("stream", True),
        }
        
        tool_params = {}
        if tools :
            tool_params = {
                "tools": tools,
                "tool_choice": kwargs.get("tool_choice", "auto")
            }
        
        params = {**base_params, **tool_params}
        
        return client.chat.completions.create(
            **params,
            extra_headers={
                "HTTP-Referer": "https://zanomega.com/open-taranis/",
                "X-Title": "open-taranis"
            }
        )

# ==============================
# Functions for the streaming
# ==============================

@staticmethod
def handle_streaming(stream: openai.Stream):
    """
    return :
    - token : str or None
    - tool : list
    - tool_bool : bool
    """
    tool_calls = []
    accumulated_tool_calls = {}
    arg_chunks = {}  # Per tool_call index: list of argument chunks

    is_thinking = False

    # Process each chunk
    for chunk in stream:
        # Skip if no choices
        if not chunk.choices:
            continue

        delta = chunk.choices[0].delta
        if delta is None:
            continue
        
        # Handle reasoning streaming
        if getattr(delta, 'reasoning_content', None):
            if is_thinking :
                yield delta.reasoning_content, [], False
            else :
                yield "<think>\n", [], False
                yield delta.reasoning_content, [], False
                is_thinking = True

        # Handle content streaming
        if delta.content :
            if is_thinking :
                yield "</think>\n", [], False
                yield delta.content, [], False
                is_thinking = False
            else :
                yield delta.content, [], False

        # Handle tool calls in delta
        if delta.tool_calls:
            for tool_call in delta.tool_calls:
                index = tool_call.index
                if index not in accumulated_tool_calls:
                    accumulated_tool_calls[index] = {
                        "id": tool_call.id,
                        "function": {"name": "", "arguments": ""},
                        "type": tool_call.type,
                        "arg_chunks": []  # New: list for arguments
                    }
                    arg_chunks[index] = []
                if tool_call.function:
                    if tool_call.function.name:
                        accumulated_tool_calls[index]["function"]["name"] += tool_call.function.name
                    if tool_call.function.arguments:
                        # Append to list instead of +=
                        arg_chunks[index].append(tool_call.function.arguments)

    # Stream finished - check if we have accumulated tool calls
    # Finalize arguments for each tool call
    for idx in accumulated_tool_calls:
        call = accumulated_tool_calls[idx]
        # Join arg chunks
        joined_args = ''.join(arg_chunks.get(idx, []))
        if joined_args:
            # Try to parse the full joined string
            try:
                parsed_args = json.loads(joined_args)
                call["function"]["arguments"] = json.dumps(parsed_args)
            except json.JSONDecodeError:
                # Fallback: attempt to extract valid JSON substring
                # Look for balanced braces starting from end
                start = joined_args.rfind('{')
                if start != -1:
                    potential_json = joined_args[start:]
                    try:
                        parsed_args = json.loads(potential_json)
                        call["function"]["arguments"] = json.dumps(parsed_args)
                    except json.JSONDecodeError:
                        # Last resort: use raw joined as string
                        call["function"]["arguments"] = joined_args
                else:
                    call["function"]["arguments"] = joined_args

    if accumulated_tool_calls:
        tool_calls = [
            {
                "id": call["id"],
                "function": call["function"],
                "type": call["type"]
            }
            for call in accumulated_tool_calls.values()
        ]
    yield "", tool_calls, len(tool_calls) > 0

@staticmethod
def handle_tool_call(tool_call:dict) -> tuple[str, str, dict, str] :
    """
    Return :
    - function id : str
    - function name : str
    - arguments : dict
    - error_message : str 
    """
    fid = tool_call.get("id", "")
    fname = tool_call.get("function", {}).get("name", "")
    raw_args = tool_call.get("function", {}).get("arguments", "{}")
    
    try:
        cleaned = re.sub(r'(?<=\d)_(?=\d)', '', raw_args)
        args = json.loads(cleaned)
    except json.JSONDecodeError as e:
        return fid, fname, {}, str(e)

    return fid, fname, args, ""

# Utility for multiple functions, code by Kimi k2 thinking
@staticmethod
def functions_to_tools(funcs: list[Callable]) -> list[dict[str, Any]]:
    return [utils.function_to_openai_tool(f) for f in funcs]

# ==============================
# Functions to simplify the messages roles
# ==============================

@staticmethod
def create_assistant_response(content:str, tool_calls:list[dict]=None, reasoning_content:str=None) -> dict[str, str]:
    """
    Creates an assistant message, optionally with tool calls.
    
    Args:
        content (str): Textual content of the response
        tool_calls (list[dict], optional): List of tool calls
        
    Returns:
        dict: Message formatted for the API
    """
    r = {"role": "assistant","content": content}
    if tool_calls : 
        r.update({"tool_calls":tool_calls})
    if reasoning_content :
        r.update({"reasoning_content":reasoning_content})
    return r

@staticmethod
def create_function_response(id:str, result:str, name:str) -> dict[str, str, str]:
    if not id or not name:
        raise ValueError("id and name are required")
    return {"role": "tool", "content": json.dumps(result), "tool_call_id": id, "name": name}

@staticmethod
def create_system_prompt(content:str) -> dict[str, str] :
    return {"role":"system", "content":content}

@staticmethod
def create_user_prompt(content:str) -> dict[str, str] :
    return {"role":"user", "content":content}

# ==============================
# Agents base
# ==============================

class agent_base:
    def __init__(self, is_thinking_enabled:bool=False, yield_thinking:bool=False):

        self._system_prompt = [create_system_prompt("")]
        self.messages = []
        self.tools = []

        self.meta = {
            "create_stream":True,
            "is_thinking_enabled":is_thinking_enabled,
            "yield_thinking":yield_thinking,
        }
    
    def create_stream(self):
        """
        # TO IMPLEMENT

        like this :
        ```python
        return clients.Your_request(
            client=clients.Your,
            messages=self._system_prompt+self.messages, # Need to be keep !
            model="Yout model",
            tools=self.tools
        )
        ```
        but with your customisation

        Yout can define your client in `__init__()`
        """

        if self.meta["create_stream"]:
            raise "You MUST define 'agent_base.create_stream()'"   

    def manage_user_prompt(self, prompt):
        """
        # TO IMPLEMENT if needed
        """

        return prompt

    def manage_assistant_response(self, response):
        """
        # TO IMPLEMENT if needed
        """

        return response

    def manage_messages_in_reply(self):
        """
        Function to manage message history, executed at each step (after agent response or tool call)


        ```
        """
        pass

    def manage_messages_after_reply(self):
        """
        Message history management function, executed after each reply

        Ex:
        - Compress messages
        - Reduce to the last X
        - And more...

        ---

        Example to always store only the 2 lasts turns (without tools !) :
        ```python
        self.messages = self.messages[-4:]
        ```
        """
        pass

    def execute_tools(self, fname, args):
        raise NotImplementedError("Subclasses must implement execute_tools()")
    
    def manage_token_yield(self, token, is_thinking=None):
        """
        # TO IMPLEMENT if needed for custom front !
        """
        return token

    def __call__(self, prompt):

        run = True

        self.messages.append(create_user_prompt(
            self.manage_user_prompt(prompt)
        ))
        
        while run :
            is_thinking=False # Reset at each request
            response = ""
            reasoning = ""
            for token, tool_calls, run in handle_streaming(self.create_stream()) :
                if token :

                    if "<think>" in token or is_thinking :
                        is_thinking=True

                        if "</think>" in token :
                            is_thinking=False

                        if self.meta["is_thinking_enabled"] :
                            reasoning += token
                        else :
                            response += token
                        
                        if self.meta["yield_thinking"]:
                            yield self.manage_token_yield(token, is_thinking)
                    
                    else :
                        yield self.manage_token_yield(token, is_thinking)
                        response += token

            if run:

                if self.meta["is_thinking_enabled"] :
                    self.messages.append(create_assistant_response(
                        self.manage_assistant_response(response), tool_calls,
                        reasoning_content=reasoning)
                    )
                else :
                    self.messages.append(create_assistant_response(
                        self.manage_assistant_response(response), tool_calls)
                    )

                for tool_call in tool_calls :
                    fid, fname, args, _ = handle_tool_call(tool_call)

                    result = self.execute_tools(fname, args)

                    self.messages.append(create_function_response(
                        id=fid, result=result, name=fname
                    ))
            
            self.manage_messages_in_reply()

        reasoning = re.sub(r'<think>\n?|</think>\n', '', reasoning).strip()
            
        if self.meta["is_thinking_enabled"] :
            self.messages.append(create_assistant_response(
                self.manage_assistant_response(response), tool_calls,
                reasoning_content=reasoning)
            )
        else :
            self.messages.append(create_assistant_response(
                self.manage_assistant_response(response), tool_calls)
            )

        self.manage_messages_after_reply()