from typing import Any, Dict, Union, get_args, get_origin, Optional

from phi.utils.log import logger


def get_json_type_for_py_type(arg: str) -> str:
    """
    Get the JSON schema type for a given type.
    :param arg: The type to get the JSON schema type for.
    :return: The JSON schema type.
    """
    # logger.info(f"Getting JSON type for: {arg}")
    if arg in ("int", "float", "complex", "Decimal"):
        return "number"
    elif arg in ("str", "string"):
        return "string"
    elif arg in ("bool", "boolean"):
        return "boolean"
    elif arg in ("NoneType", "None"):
        return "null"
    elif arg in ("list", "tuple", "set", "frozenset"):
        return "array"
    elif arg in ("dict", "mapping"):
        return "object"

    # If the type is not recognized, return "object"
    return "object"


def get_json_schema_for_arg(t: Any) -> Optional[Dict[str, Any]]:
    # logger.info(f"Getting JSON schema for arg: {t}")
    type_args = get_args(t)
    # logger.info(f"Type args: {type_args}")
    type_origin = get_origin(t)
    # logger.info(f"Type origin: {type_origin}")

    if type_origin is not None:
        if type_origin in (list, tuple, set, frozenset):
            json_schema_for_items = get_json_schema_for_arg(type_args[0]) if type_args else {"type": "string"}
            return {"type": "array", "items": json_schema_for_items}
        elif type_origin is dict:
            # Handle both key and value types for dictionaries
            key_schema = get_json_schema_for_arg(type_args[0]) if type_args else {"type": "string"}
            value_schema = get_json_schema_for_arg(type_args[1]) if len(type_args) > 1 else {"type": "string"}
            return {"type": "object", "propertyNames": key_schema, "additionalProperties": value_schema}
        elif type_origin is Union:
            types = []
            for arg in type_args:
                if arg is not type(None):
                    try:
                        schema = get_json_schema_for_arg(arg)
                        if schema:
                            types.append(schema)
                    except Exception:
                        continue
            return {"anyOf": types} if types else None

    return {"type": get_json_type_for_py_type(t.__name__)}


def get_json_schema(
    type_hints: Dict[str, Any], param_descriptions: Optional[Dict[str, str]] = None, strict: bool = False
) -> Dict[str, Any]:
    json_schema: Dict[str, Any] = {
        "type": "object",
        "properties": {},
    }
    if strict:
        json_schema["additionalProperties"] = False

    for k, v in type_hints.items():
        # logger.info(f"Parsing arg: {k} | {v}")
        if k == "return":
            continue

        try:
            # Check if type is Optional (Union with NoneType)
            type_origin = get_origin(v)
            type_args = get_args(v)
            is_optional = type_origin is Union and len(type_args) == 2 and any(arg is type(None) for arg in type_args)

            # Get the actual type if it's Optional
            if is_optional:
                v = next(arg for arg in type_args if arg is not type(None))

            # Handle cases with no type hint
            if v:
                arg_json_schema = get_json_schema_for_arg(v)
            else:
                arg_json_schema = {}

            if arg_json_schema is not None:
                if is_optional:
                    # Handle null type for optional fields
                    if isinstance(arg_json_schema["type"], list):
                        arg_json_schema["type"].append("null")
                    else:
                        arg_json_schema["type"] = [arg_json_schema["type"], "null"]

                # Add description
                if param_descriptions and k in param_descriptions and param_descriptions[k]:
                    arg_json_schema["description"] = param_descriptions[k]

                json_schema["properties"][k] = arg_json_schema

            else:
                logger.warning(f"Could not parse argument {k} of type {v}")
        except Exception as e:
            logger.error(f"Error processing argument {k}: {str(e)}")
            continue

    return json_schema