from contextlib import redirect_stdout
from io import StringIO

import pytest
from fastapi import FastAPI
from fastapi.testclient import TestClient

from swarms.structs.agent import Agent
from swarms.utils.class_args_wrapper import print_class_parameters

app = FastAPI()


def test_print_class_parameters_agent():
    f = StringIO()
    with redirect_stdout(f):
        print_class_parameters(Agent)
    output = f.getvalue().strip()
    # Replace with the expected output for Agent class
    expected_output = (
        "Parameter: name, Type: <class 'str'>\nParameter: age, Type:"
        " <class 'int'>"
    )
    assert output == expected_output


def test_print_class_parameters_error():
    with pytest.raises(TypeError):
        print_class_parameters("Not a class")


@app.get("/parameters/{class_name}")
def get_parameters(class_name: str):
    classes = {"Agent": Agent}
    if class_name in classes:
        return print_class_parameters(classes[class_name], api_format=True)
    else:
        return {"error": "Class not found"}


client = TestClient(app)


def test_get_parameters_agent():
    response = client.get("/parameters/Agent")
    assert response.status_code == 200
    # Replace with the expected output for Agent class
    expected_output = {"x": "<class 'int'>", "y": "<class 'int'>"}
    assert response.json() == expected_output


def test_get_parameters_not_found():
    response = client.get("/parameters/NonexistentClass")
    assert response.status_code == 200
    assert response.json() == {"error": "Class not found"}