import json

from stream import Stream
from tool import ToolSpec, ToolCall
from inference import Message

# TODO: Move to module 'format'
def encode_value(value):
    if isinstance(value, ToolSpec):
        return value.to_openai_dict()
    elif isinstance(value, Message):
        return {'role': value.role, 'content': encode_value(value.content)}
    elif isinstance(value, ToolCall):
        json = {
            'id': value.id,
            'tool_name': value.tool_name,
            'arguments': encode_value(value.arguments),
        }
        if value.result != None:
            json['result'] = encode_value(value.result),
        return json
    elif isinstance(value, Stream):
        if value.open:
            return {'streamId': value.id}
        else:
            return value.content
    else:
        return value

def encode_json(value):
    def encode_unknown(value):
        # All complex values are transformed to a dict
        value = encode_value(value)
        # otherwise simply show the repr
        if isinstance(value, dict):
            return value
        else:
            return repr(value)
    return json.dumps(value, default=encode_unknown)