from r1_json import JSON, Union, JsonKey


def toJSON( val: Union[ JSON, int, str ] ) -> JSON:
    if isinstance( val, str ):
        return JsonStr( val )
    if isinstance( val, int ):
        return JsonInt( val )

    return val

class JsonArray:
    def __init__( self ) -> None:
        self.arr : list[ JSON ] = []
    def get( self, key: JsonKey ) -> JSON:
        assert isinstance( key, int )
        return self.arr[ key ]
    def set( self, key: int, val: Union[ JSON, int, str ] ) -> None:
        assert isinstance( key, int )
        self.arr[ key ] = toJSON( val )
    def append( self, val: JSON ) -> None:
        self.arr.append( val )


class JsonObject:
    def __init__( self ) -> None:
        self.assoc : dict[ str, JSON ] = {}

    def get( self, key: JsonKey ) -> JSON:
        return self.assoc[ str( key ) ]

    def set( self, key: JsonKey, val: Union[ JSON, int, str ] ) -> None:
        self.assoc[ str( key ) ] = toJSON( val )


class JsonWrapper:
    def get( self, key: Union[ str, int ] ) -> JSON:
        assert False
    def set( self, key: Union[ str, int ], val: JSON ) -> None:
        assert False

class JsonInt( int, JsonWrapper ): pass
class JsonStr( str, JsonWrapper ): pass