import base64
from typing import get_type_hints, Callable, Dict, Any
from enum import Enum
from pathlib import Path
import json
import re
from toolla.exceptions import InvalidDescriptionException

def extract_json_from_text(text):
    text_without_newlines = text.replace('\n', '')
    start_index = None
    curly_brace_count = 0
    for i, c in enumerate(text_without_newlines):
        if c == '{':
            curly_brace_count += 1
            if not start_index:
                start_index = i
        if c == '}':
            curly_brace_count -= 1
            if curly_brace_count == 0:
                return json.loads(text_without_newlines[start_index:i+1])
    return None

# Generated by Claude Sonnet 3.5
def parse_descriptions(input_string):
    lines = input_string.strip().split('\n')
    result = {}
    colon_found = False
    for i, line in enumerate(lines):
        if i == 0:
            result['fn_description'] = line.strip()
        elif ':' in line:
            colon_found = True
            key, value = line.split(':', 1)
            key = key.strip()
            value = value.strip()
            result[key] = value
    
    if not colon_found:
        raise InvalidDescriptionException
    
    return result

def build_openai_tool_schema(f: Callable):
    """Build a JSON schema for an OpenAI tool function."""
    descriptions = parse_descriptions(f.__doc__)
    hints = get_type_hints(f)
    schema = {
        "type": "function",
        "function": {
            "name": f.__name__,
            "description": descriptions['fn_description'],
            "parameters": {
                "type": "object",
                "properties": {},
                "required": []
            },
        }
    }

    for param, hint in hints.items():
        if param == "return":
            continue
        if isinstance(hint, type) and issubclass(hint, Enum):
            schema["function"]["parameters"]["properties"][param] = {
                "type": "string",
                "enum": [e.name for e in hint],
                "description": descriptions[param]
            }
            schema["function"]["parameters"]["required"].append(param)
            continue

        var_type = "number" if hint is int or hint is float else "string"
        schema["function"]["parameters"]["properties"][param] = {
            "type": var_type,
            "description": descriptions[param]
        }
        schema["function"]["parameters"]["required"].append(param)
    return schema

# TODO add support for enums
def build_claude_tool_schema(f: Callable):
    """Build a JSON schema for a tool function."""
    descriptions = parse_descriptions(f.__doc__)
    hints = get_type_hints(f)
    schema = {
        "name": f.__name__,
        "description": descriptions['fn_description'],
        "input_schema": {
            "type": "object",
            "properties": {},
            "required": []
        },
    }

    for param, hint in hints.items():
        if param == "return":
            continue
        if isinstance(hint, type) and issubclass(hint, Enum):
            schema["input_schema"]["properties"][param] = {
                "type": "string",
                "enum": [e for e in hint.__members__],
                "description": descriptions[param]
            }
            schema["input_schema"]["required"].append(param)
            continue

        var_type = "number" if hint is int or hint is float else "string"
        schema["input_schema"]["properties"][param] = {
            "type": var_type,
            "description": descriptions[param]
        }
        schema["input_schema"]["required"].append(param)
    return schema

def load_file_base64(file_path: Path) -> str:
    with open(file_path, "rb") as file:
        file_content = file.read()
        return base64.b64encode(file_content).decode('utf-8')

def get_image_mime_type(fpath: Path) -> str:
    suffix = fpath.suffix
    mime_types = {
        '.jpg': 'image/jpeg',
        '.jpeg': 'image/jpeg',
        '.png': 'image/png',
        '.gif': 'image/gif',
        '.webp': 'image/webp',
    }
    return mime_types.get(suffix)
