#|default_exp funccall

#| exports
import inspect
from collections import abc
from fastcore.utils import *
from fastcore.docments import docments
from typing import get_origin, get_args, Dict, List, Optional, Tuple, Union
from types import UnionType
from datetime import date

#|hide
from fastcore.test import *

empty = inspect.Parameter.empty

def silly_sum(
    a:int, # First thing to sum
    b:int=1, # Second thing to sum
    c:list[int]=None, # A pointless argument
) -> int: # The sum of the inputs
    "Adds a + b."
    return a + b

d = docments(silly_sum, full=True)
d

d.a.docment, d['a']['anno']

#| exports
def _types(t:type)->tuple[str,Optional[str]]:
    "Tuple of json schema type name and (if appropriate) array item name."
    if t is empty: raise TypeError('Missing type')
    tmap = {int:"integer", float:"number", str:"string", bool:"boolean", list:"array", dict:"object"}
    tmap.update({k.__name__: v for k, v in tmap.items()})
    if getattr(t, '__origin__', None) in (list,tuple):
        args = getattr(t, '__args__', None)
        item_type = "object" if not args else tmap.get(t.__args__[0].__name__, "object")
        return "array", item_type
    # if t is a string like 'int', directly use the string as the key
    elif isinstance(t, str): return tmap.get(t, "object"), None
    # if t is the type itself and a container
    elif get_origin(t): return tmap.get(get_origin(t).__name__, "object"), None
    # if t is the type itself like int, use the __name__ representation as the key
    else: return tmap.get(t.__name__, "object"), None

_types(list[int]), _types(int), _types('int')

_types(List[int]), _types(Optional[str]), _types(str | None), _types(Tuple[str, int])

class Custom: a: int
_types(list[Custom]), _types(Custom)

#| exports
def _param(name, info):
    "json schema parameter given `name` and `info` from docments full dict."
    paramt,itemt = _types(info.anno)
    pschema = dict(type=paramt, description=info.docment or "")
    if itemt: pschema["items"] = {"type": itemt}
    if info.default is not empty: pschema["default"] = info.default
    return pschema

n,o = first(d.items())
print(n,'//', o)
_param(n, o)

custom_types = {Path, date}

def _handle_type(t, defs):
    "Handle a single type, creating nested schemas if necessary"
    if t is NoneType: return {'type': 'null'}
    if t in custom_types: return {'type':'string', 'format':t.__name__}
    if t in (dict, list, tuple, set): return {'type': _types(t)[0]}
    if isinstance(t, type) and not issubclass(t, (int, float, str, bool)) or inspect.isfunction(t):
        defs[t.__name__] = _get_nested_schema(t)
        return {'$ref': f'#/$defs/{t.__name__}'}
    return {'type': _types(t)[0]}

_handle_type(int, None), _handle_type(Path, None)

def _is_container(t):
    "Check if type is a container (list, dict, tuple, set, Union)"
    origin = get_origin(t)
    return origin in (list, dict, tuple, set, Union) if origin else False

def _is_parameterized(t):
    "Check if type has arguments (e.g. list[int] vs list, dict[str, int] vs dict)"
    return _is_container(t) and (get_args(t) != ())

assert _is_parameterized(list[int]) == True
assert _is_parameterized(int) == False
assert _is_container(list[int]) == True
assert _is_container(dict[str, int]) == True
assert _is_container(int) == False

def _example_new_unioin(opt_tup: str | None):
    pass

d = docments(_example_new_unioin, full=True)
anno1 = first(d.items())[1].anno
(anno1, get_origin(anno1), get_args(anno1))

def _example_old_union(opt_tup: Union[str, type(None)] =None):
    pass

d = docments(_example_old_union, full=True)
anno2 = first(d.items())[1].anno
(anno2, get_origin(anno2), get_args(anno2))

def _handle_container(origin, args, defs):
    "Handle container types like dict, list, tuple, set, and Union"
    if origin is Union or origin is UnionType:
        return {"anyOf": [_handle_type(arg, defs) for arg in args]}
    if origin is dict:
        value_type = args[1].__args__[0] if hasattr(args[1], '__args__') else args[1]
        return {
            'type': 'object',
            'additionalProperties': (
                {'type': 'array', 'items': _handle_type(value_type, defs)}
                if hasattr(args[1], '__origin__') else _handle_type(args[1], defs)
            )
        }
    elif origin in (list, tuple, set):
        schema = {'type': 'array', 'items': _handle_type(args[0], defs)}
        if origin is set:
            schema['uniqueItems'] = True
        return schema
    return None

def _process_property(name, obj, props, req, defs):
    "Process a single property of the schema"
    p = _param(name, obj)
    props[name] = p
    if obj.default is empty: req[name] = True

    if _is_container(obj.anno) and _is_parameterized(obj.anno):
        p.update(_handle_container(get_origin(obj.anno), get_args(obj.anno), defs))        
    else:
        # Non-container type or container without arguments
        p.update(_handle_type(obj.anno, defs))

def _get_nested_schema(obj):
    "Generate nested JSON schema for a class or function"
    d = docments(obj, full=True)
    props, req, defs = {}, {}, {}

    for n, o in d.items():
        if n != 'return' and n != 'self':
            _process_property(n, o, props, req, defs)

    tkw = {}
    if isinstance(obj, type): tkw['title']=obj.__name__
    schema = dict(type='object', properties=props, **tkw)
    if req: schema['required'] = list(req)
    if defs: schema['$defs'] = defs
    return schema

# Test primitive types
defs = {}
assert _handle_type(int, defs) == {'type': 'integer'}
assert _handle_type(str, defs) == {'type': 'string'}
assert _handle_type(bool, defs) == {'type': 'boolean'}
assert _handle_type(float, defs) == {'type': 'number'}

# Test custom class
class TestClass:
    def __init__(self, x: int, y: int): store_attr()

result = _handle_type(TestClass, defs)
assert result == {'$ref': '#/$defs/TestClass'}
assert 'TestClass' in defs
assert defs['TestClass']['type'] == 'object'
assert 'properties' in defs['TestClass']

# Test primitive types in containers
assert _handle_container(list, (int,), defs) == {'type': 'array', 'items': {'type': 'integer'}}
assert _handle_container(tuple, (str,), defs) == {'type': 'array', 'items': {'type': 'string'}}
assert _handle_container(set, (str,), defs) == {'type': 'array', 'items': {'type': 'string'}, 'uniqueItems': True}
assert _handle_container(dict, (str,bool), defs) == {'type': 'object', 'additionalProperties': {'type': 'boolean'}}

result = _handle_container(list, (TestClass,), defs)
assert result == {'type': 'array', 'items': {'$ref': '#/$defs/TestClass'}}
assert 'TestClass' in defs

# Test complex nested structure
ComplexType = dict[str, list[TestClass]]
result = _handle_container(dict, (str, list[TestClass]), defs)
assert result == {
    'type': 'object',
    'additionalProperties': {
        'type': 'array',
        'items': {'$ref': '#/$defs/TestClass'}
    }
}

# Test processing of a required integer property
props, req = {}, {}
class TestClass:
    "Test class"
    def __init__(
        self,
        x: int, # First thing
        y: list[float], # Second thing
        z: str = "default", # Third thing
    ): store_attr()

d = docments(TestClass, full=True)
_process_property('x', d.x, props, req, defs)
assert 'x' in props
assert props['x']['type'] == 'integer'
assert 'x' in req

# Test processing of a required list property
_process_property('y', d.y, props, req, defs)
assert 'y' in props
assert props['y']['type'] == 'array'
assert props['y']['items']['type'] == 'number'
assert 'y' in req

# Test processing of an optional string property with default
_process_property('z', d.z, props, req, defs)
assert 'z' in props
assert props['z']['type'] == 'string'
assert props['z']['default'] == "default"
assert 'z' not in req

#| exports
def get_schema(f:Union[callable,dict], pname='input_schema')->dict:
    "Generate JSON schema for a class, function, or method"
    if isinstance(f, dict): return f
    schema = _get_nested_schema(f)
    desc = f.__doc__
    assert desc, "Docstring missing!"
    d = docments(f, full=True)
    ret = d.pop('return')
    if (ret.anno is not empty) and (ret.anno is not None): desc += f'\n\nReturns:\n- type: {_types(ret.anno)[0]}'
    return {"name": f.__name__, "description": desc, pname: schema}

s = get_schema(silly_sum)
desc = s.pop('description')
print(desc)
s

def silly_test(
    a: 'int',  # quoted type hint
):
    "Mandatory docstring"
    return a

get_schema(silly_test)

class Dummy:
    def sums(
        self,
        a:int,  # First thing to sum
        b:int=1 # Second thing to sum
    ) -> int: # The sum of the inputs
        "Adds a + b."
        print(f"Finding the sum of {a} and {b}")
        return a + b

get_schema(Dummy.sums)

class Turn:
    "Turn between two speakers"
    def __init__(
        self,
        speaker_a:str, # First speaker's message
        speaker_b:str,  # Second speaker's message
    ): store_attr()

class Conversation:
    "A conversation between two speakers"
    def __init__(
        self,
        turns:list[Turn], # Turns of the conversation
    ): store_attr()

get_schema(Conversation)

class DictConversation:
    "A conversation between two speakers"
    def __init__(
        self,
        turns:dict[str,list[Turn]], # dictionary of topics and the Turns of the conversation
    ): store_attr()

get_schema(DictConversation)

class SetConversation:
    "A conversation between two speakers"
    def __init__(
        self,
        turns:set[Turn], # the unique Turns of the conversation
    ): store_attr()

get_schema(SetConversation)

#| exports
def PathArg(
    path: str  # A filesystem path
): return Path(path)

def path_test(
    a: PathArg,  # a type hint
    b: PathArg   # b type hint
):
    "Mandatory docstring"
    return a/b

get_schema(path_test)

def path_test2(
    a: Path,  # a type hint
    b: Path   # b type hint
):
    "Mandatory docstring"
    return a/b

get_schema(path_test2)

#| exports
def dateArg(
    datestr: str  # YYYY-MM-DD
): 
    parts = datestr.split('-')
    return date(*map(int,parts))

dateArg('2025-12-02')

def date_test(
    a: dateArg, # a type hint
    b: dateArg  # b type hint
):
    "Mandatory docstring"
    return b-a

get_schema(date_test)

def date_test2(
    a: date, # a type hint
    b: date  # b type hint
):
    "Mandatory docstring"
    return b-a

get_schema(date_test)