import asyncio
import json
import unittest

from server import JSON_SCHEMA_DRAFT, Server


class FunctionTests(unittest.TestCase):
    def setUp(self):
        self.app = Server()

    def request(self, method, path, body=b""):
        async def call():
            messages = []

            async def receive():
                return {"type": "http.request", "body": body, "more_body": False}

            async def send(message):
                messages.append(message)

            await self.app(
                {"type": "http", "method": method, "path": path}, receive, send
            )
            return messages[0]["status"], messages[1]["body"]

        return asyncio.run(call())

    def test_reflection_validates_calls_and_describes_function(self):
        source = b"def double(x: int) -> int:\n    return x * 2\n"
        self.assertEqual(self.request("PATCH", "/double", source)[0], 200)
        status, body = self.request("GET", "/double")
        self.assertEqual(status, 200)
        entry = json.loads(self.request("GET", "/")[1])["functions"][0]
        self.assertEqual(json.loads(body), entry)
        self.assertEqual(entry["name"], "double")
        self.assertEqual(entry["source"], source.decode())
        input_schema = entry["schema"]["input"]
        self.assertEqual(input_schema["$schema"], JSON_SCHEMA_DRAFT)
        self.assertNotIn("$defs", input_schema)
        self.assertEqual(input_schema["anyOf"][1], {"type": "integer"})
        self.assertEqual(input_schema["anyOf"][0]["required"], ["x"])
        self.assertEqual(
            entry["schema"]["output"],
            {"$schema": JSON_SCHEMA_DRAFT, "type": "integer"},
        )
        self.assertEqual(self.request("POST", "/double", b'{"x": 4}'), (200, b"8"))
        self.assertEqual(self.request("POST", "/double", b"4"), (200, b"8"))
        for invalid in (b'{"x": "4"}', b'{"x": 4.5}', b"{}", b"[]", b"NaN"):
            self.assertEqual(self.request("POST", "/double", invalid)[0], 400)
        self.assertEqual(self.request("GET", "/double/schema")[0], 404)
        self.assertEqual(self.request("PUT", "/double/schema", b"{}")[0], 404)

        self.assertEqual(
            self.request(
                "PATCH", "/double", b"def double(x: int) -> int:\n    return str(x)"
            )[0],
            200,
        )
        self.assertEqual(self.request("POST", "/double", b'{"x": 4}')[0], 500)

    def test_invalid_update_keeps_previous_function(self):
        original = b"def double(x: int) -> int:\n    return x * 2"
        self.request("PATCH", "/double", original)
        for invalid in (
            b"return x * 2",
            b"def other(x: int) -> int:\n    return x",
            b"def double(x):\n    return x",
            b"async def double(x: int) -> int:\n    return x",
        ):
            self.assertEqual(self.request("PATCH", "/double", invalid)[0], 400)
        status, body = self.request("GET", "/double")
        self.assertEqual(status, 200)
        self.assertEqual(json.loads(body)["source"], original.decode())
        self.assertEqual(self.request("POST", "/double", b"3"), (200, b"6"))

    def test_multiple_parameters_union_and_default(self):
        source = (
            b'def label(x: int | str, prefix: str = "item") -> str:\n'
            b'    return f"{prefix}:{x}"'
        )
        self.assertEqual(self.request("PATCH", "/label", source)[0], 200)
        self.assertEqual(
            self.request("POST", "/label", b'{"x":4,"prefix":"order"}'),
            (200, b'"order:4"'),
        )
        self.assertEqual(
            self.request("POST", "/label", b'{"x":"abc"}'),
            (200, b'"item:abc"'),
        )
        for invalid in (b'{"x":4.5}', b'{"prefix":"order"}', b'{"x":4,"extra":1}'):
            self.assertEqual(self.request("POST", "/label", invalid)[0], 400)
        schema = json.loads(self.request("GET", "/")[1])["functions"][0]["schema"]
        self.assertEqual(schema["input"]["$schema"], JSON_SCHEMA_DRAFT)
        self.assertEqual(schema["input"]["required"], ["x"])
        self.assertEqual(
            {item["type"] for item in schema["input"]["properties"]["x"]["anyOf"]},
            {"integer", "string"},
        )
        self.assertEqual(schema["input"]["properties"]["prefix"]["default"], "item")

    def test_model_input_and_output(self):
        source = b"""from pydantic import BaseModel
class Order(BaseModel):
    quantity: int
    unit_price: float
def total(x: Order) -> Order:
    return x
"""
        self.assertEqual(self.request("PATCH", "/total", source)[0], 200)
        self.assertEqual(
            json.loads(
                self.request(
                    "POST", "/total", b'{"x":{"quantity":2,"unit_price":3.5}}'
                )[1]
            ),
            {"quantity": 2, "unit_price": 3.5},
        )
        self.assertEqual(
            json.loads(
                self.request("POST", "/total", b'{"quantity":2,"unit_price":3.5}')[1]
            ),
            {"quantity": 2, "unit_price": 3.5},
        )
        self.assertEqual(
            self.request("POST", "/total", b'{"x":{"quantity":"2","unit_price":3.5}}')[
                0
            ],
            400,
        )
        entry = json.loads(self.request("GET", "/")[1])["functions"][0]
        self.assertEqual(entry["schema"]["input"]["$schema"], JSON_SCHEMA_DRAFT)
        self.assertEqual(
            entry["schema"]["input"]["$defs"]["Order"]["properties"]["quantity"][
                "type"
            ],
            "integer",
        )

    def test_single_parameter_accepts_every_json_value(self):
        source = b"from typing import Any\ndef identity(x: Any) -> Any:\n    return x"
        self.assertEqual(self.request("PATCH", "/identity", source)[0], 200)
        for value in (
            b"4",
            b'"hello"',
            b"true",
            b"false",
            b"null",
            b"[1,2]",
            b'{"item":1}',
        ):
            status, response = self.request("POST", "/identity", value)
            self.assertEqual(status, 200, value)
            self.assertEqual(json.loads(response), json.loads(value))
        self.assertEqual(self.request("POST", "/identity", b'{"x":4}'), (200, b"4"))


if __name__ == "__main__":
    unittest.main()