import inspect
from inspect import signature
from typing import Dict, List, Union, get_args, get_origin, get_type_hints

from pydantic import BaseModel


def convert_function_call_to_json(func, *args, **kwargs):
    input_data = dict(zip(signature(func).parameters, args))
    input_data.update(kwargs)
    for key, value in input_data.items():
        if isinstance(value, BaseModel):
            input_data[key] = value.model_dump()
        elif isinstance(value, list) and all(isinstance(v, BaseModel) for v in value):
            input_data[key] = [v.model_dump() for v in value]
    return input_data


def type_to_json_schema(type_hint):
    """Convert a Python type to a JSON schema."""
    if type_hint == int:
        return {"type": "integer"}
    elif type_hint == str:
        return {"type": "string"}
    elif type_hint == float:
        return {"type": "number"}
    elif type_hint == bool:
        return {"type": "boolean"}
    elif inspect.isclass(type_hint) and issubclass(type_hint, BaseModel):
        return type_hint.model_json_schema()
    elif hasattr(type_hint, "__origin__"):
        # Handling generic types (List, Dict, etc.)
        if get_origin(type_hint) == Union and type(None) in get_args(type_hint):
            return type_to_json_schema(get_args(type_hint)[0])  # for Optional
        elif type_hint.__origin__ in (List, list):
            item_type = get_args(type_hint)[0]
            if inspect.isclass(item_type):
                return {"type": "array", "items": type_to_json_schema(item_type)}
            if get_origin(item_type) == List:
                return {"type": "array", "items": type_to_json_schema(item_type)}
            elif inspect.isclass(item_type) and issubclass(item_type, BaseModel):
                return {"type": "array", "items": item_type.schema()}
            else:
                return {"type": "array", "items": type_to_json_schema(item_type)}
        elif get_origin(type_hint) == Dict:
            return {"type": "object"}
        else:
            raise NotImplementedError(
                f"Type {type_hint} is not supported for JSON schema generation.{get_origin(type_hint)}"
            )
    else:
        return {"type": "object"}


def get_input_schema(func):
    """Generate input JSON schema from a function signature."""
    sig = inspect.signature(func)
    type_hints = get_type_hints(func)

    schema = {"type": "object", "properties": {}, "required": []}

    for param in sig.parameters.values():
        param_type = type_hints.get(
            param.name, str
        )  # Default to string if type is not specified
        schema["properties"][param.name] = type_to_json_schema(param_type)
        if param.default is inspect.Parameter.empty:
            schema["required"].append(param.name)

    if not schema["properties"]:
        return {
            "type": "object",
            "description": "No input parameters",
            "properties": {},
            "additionalProperties": False,
        }

    return schema


def get_output_schema(func):
    """Generate output JSON schema from a function signature."""
    return_type = get_type_hints(func).get("return", None)
    if return_type:
        schema = type_to_json_schema(return_type)
    else:
        schema = {"type": "unknown"}

    return schema
