# Copyright 2025 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# SPDX-License-Identifier: Apache-2.0

"""Tests for picoschema functionality."""

import enum
import unittest
from typing import Any

from pydantic import BaseModel

from dotpromptz import picoschema
from dotpromptz.typing import JsonSchema


class TestPicoschemaParser(unittest.TestCase):
    """Picoshema parser functionality tests."""

    def setUp(self) -> None:
        """Set up the test case."""
        self.parser = picoschema.PicoschemaParser()

    def test_must_resolve_schema_success(self) -> None:
        """Test resolving a schema successfully."""

        def mock_resolver(name: str) -> JsonSchema | None:
            if name == 'MySchema':
                return {'type': 'string', 'description': 'Resolved schema'}
            return None

        parser_with_resolver = picoschema.PicoschemaParser(schema_resolver=mock_resolver)
        resolved = parser_with_resolver.must_resolve_schema('MySchema')
        self.assertEqual(resolved, {'type': 'string', 'description': 'Resolved schema'})

    def test_must_resolve_schema_not_found(self) -> None:
        """Test resolving a non-existent schema."""

        def mock_resolver(name: str) -> JsonSchema | None:
            return None

        parser = picoschema.PicoschemaParser(schema_resolver=mock_resolver)
        with self.assertRaises(LookupError) as context:
            parser.must_resolve_schema('NonExistentSchema')
        self.assertEqual(str(context.exception), "schema resolver for 'NonExistentSchema' returned None")

    def test_must_resolve_schema_no_resolver(self) -> None:
        """Test error when resolving without a resolver."""
        with self.assertRaises(ValueError) as context:
            self.parser.must_resolve_schema('AnySchema')
        self.assertEqual(
            str(context.exception),
            "Picoschema: unsupported scalar type 'AnySchema'.",
        )

    def test_parse_no_schema(self) -> None:
        """Test parsing None returns None."""
        self.assertIsNone(self.parser.parse(None))

    def test_parse_scalar_type_schema(self) -> None:
        """Test parsing a scalar type string."""
        result = self.parser.parse('string')
        self.assertEqual(result, {'type': 'string'})

    def test_parse_object_schema(self) -> None:
        """Test parsing a standard JSON object schema."""
        schema = {'type': 'object', 'properties': {'name': {'type': 'string'}}}
        expected_schema = {
            'type': 'object',
            'properties': {'name': {'type': 'string'}},
        }
        result = self.parser.parse(schema)
        self.assertEqual(result, expected_schema)

    def test_parse_invalid_schema_type(self) -> None:
        """Test error on invalid schema type."""
        with self.assertRaises(ValueError):
            self.parser.parse(123)

    def test_parse_named_schema(self) -> None:
        """Test parsing a named schema reference."""

        def mock_resolver(name: str) -> JsonSchema | None:
            if name == 'CustomType':
                return {'type': 'integer'}
            return None

        parser_with_resolver = picoschema.PicoschemaParser(schema_resolver=mock_resolver)
        result = parser_with_resolver.parse('CustomType')
        self.assertEqual(result, {'type': 'integer'})

    def test_parse_named_schema_with_description(self) -> None:
        """Test parsing a named schema reference with a description."""

        def mock_resolver(name: str) -> JsonSchema | None:
            if name == 'DescribedType':
                return {'type': 'boolean'}
            return None

        parser_with_resolver = picoschema.PicoschemaParser(schema_resolver=mock_resolver)
        result = parser_with_resolver.parse('DescribedType, this is a description')
        self.assertEqual(
            result,
            {'type': 'boolean', 'description': 'this is a description'},
        )

    def test_parse_scalar_type_schema_with_description(self) -> None:
        """Test parsing a scalar type string with description."""
        result = self.parser.parse('string, a string')
        self.assertEqual(result, {'type': 'string', 'description': 'a string'})

    def test_parse_properties_object_shorthand(self) -> None:
        """Test parsing an object schema using the properties shorthand."""
        schema = {'name': 'string'}
        expected = {
            'type': 'object',
            'properties': {'name': {'type': 'string'}},
            'required': ['name'],
            'additionalProperties': False,
        }
        result = self.parser.parse(schema)
        self.assertEqual(result, expected)

    def test_parse_pico_scalar_type(self) -> None:
        """Test parsing a Picoschema object with scalar types."""
        result = self.parser.parse_pico('string')
        self.assertEqual(result, {'type': 'string'})

    def test_parse_pico_object_type(self) -> None:
        """Test parsing a Picoschema object type."""
        schema = {'name': 'string'}
        expected = {
            'type': 'object',
            'properties': {'name': {'type': 'string'}},
            'required': ['name'],
            'additionalProperties': False,
        }
        result = self.parser.parse_pico(schema)
        self.assertEqual(result, expected)

    def test_parse_pico_array_type(self) -> None:
        """Test parsing a Picoschema array type."""
        schema = {'names(array)': 'string'}
        expected = {
            'type': 'object',
            'properties': {'names': {'type': 'array', 'items': {'type': 'string'}}},
            'required': ['names'],
            'additionalProperties': False,
        }
        result = self.parser.parse_pico(schema)
        self.assertEqual(result, expected)

    def test_parse_pico_enum_type(self) -> None:
        """Test parsing a Picoschema enum type."""
        schema = {'status(enum)': ['active', 'inactive']}
        expected = {
            'type': 'object',
            'properties': {'status': {'enum': ['active', 'inactive']}},
            'required': ['status'],
            'additionalProperties': False,
        }
        result = self.parser.parse_pico(schema)
        self.assertEqual(result, expected)

    def test_parse_pico_optional_property(self) -> None:
        """Test parsing Picoschema with optional properties."""
        schema = {'name?': 'string'}
        expected = {
            'type': 'object',
            'properties': {'name': {'type': ['string', 'null']}},
            'additionalProperties': False,
        }
        result = self.parser.parse_pico(schema)
        self.assertEqual(result, expected)

    def test_parse_pico_wildcard_property(self) -> None:
        """Test parsing Picoschema with wildcard properties."""
        schema = {'(*)': 'string'}
        expected = {
            'type': 'object',
            'properties': {},
            'additionalProperties': {'type': 'string'},
        }
        result = self.parser.parse_pico(schema)
        self.assertEqual(result, expected)

    def test_parse_pico_nested_object(self) -> None:
        """Test parsing a Picoschema nested object."""
        schema = {'address(object)': {'street': 'string'}}
        expected = {
            'type': 'object',
            'properties': {
                'address': {
                    'type': 'object',
                    'properties': {'street': {'type': 'string'}},
                    'required': ['street'],
                    'additionalProperties': False,
                }
            },
            'required': ['address'],
            'additionalProperties': False,
        }
        result = self.parser.parse_pico(schema)
        self.assertEqual(result, expected)

    def test_parse_pico_nested_array(self) -> None:
        """Test parsing a Picoschema nested array."""
        schema = {'items(array)': {'props(array)': 'string'}}
        expected = {
            'type': 'object',
            'properties': {
                'items': {
                    'type': 'array',
                    'items': {
                        'type': 'object',
                        'properties': {
                            'props': {
                                'type': 'array',
                                'items': {
                                    'type': 'string',
                                },
                            }
                        },
                        'required': ['props'],
                        'additionalProperties': False,
                    },
                }
            },
            'required': ['items'],
            'additionalProperties': False,
        }
        result = self.parser.parse_pico(schema)
        self.assertEqual(result, expected)

    def test_parse_pico_enum_with_optional_and_null(self) -> None:
        """Test parsing a Picoschema enum with optional and null values."""
        schema = {'status?(enum)': ['active', 'inactive']}
        expected = {
            'type': 'object',
            'properties': {'status': {'enum': ['active', 'inactive', None]}},
            'additionalProperties': False,
        }
        result = self.parser.parse_pico(schema)
        self.assertEqual(result, expected)

    def test_parse_pico_description_on_type(self) -> None:
        """Test parsing Picoschema with descriptions on types."""
        schema = {'name': 'string, a name'}
        expected = {
            'type': 'object',
            'properties': {'name': {'type': 'string', 'description': 'a name'}},
            'required': ['name'],
            'additionalProperties': False,
        }
        result = self.parser.parse_pico(schema)
        self.assertEqual(result, expected)

    def test_parse_pico_description_on_optional_array(self) -> None:
        """Test parsing Picoschema description on optional array."""
        schema = {'items?(array, list of items)': 'string'}
        expected = {
            'type': 'object',
            'properties': {
                'items': {
                    'type': ['array', 'null'],
                    'items': {'type': 'string'},
                    'description': 'list of items',
                }
            },
            'additionalProperties': False,
        }
        result = self.parser.parse_pico(schema)
        self.assertEqual(result, expected)

    def test_parse_pico_description_on_enum(self) -> None:
        """Test parsing Picoschema description on an enum."""
        schema = {'status(enum, the status)': ['active', 'inactive']}
        expected = {
            'type': 'object',
            'properties': {
                'status': {
                    'enum': ['active', 'inactive'],
                    'description': 'the status',
                }
            },
            'required': ['status'],
            'additionalProperties': False,
        }
        result = self.parser.parse_pico(schema)
        self.assertEqual(result, expected)

    def test_parse_pico_description_on_custom_schema(self) -> None:
        """Test parsing Picoschema description on a custom schema type."""

        def mock_resolver(name: str) -> JsonSchema | None:
            if name == 'CustomSchema':
                return {'type': 'string'}
            return None

        parser_with_resolver = picoschema.PicoschemaParser(schema_resolver=mock_resolver)
        schema = {'field1': 'CustomSchema, a custom field'}
        expected: JsonSchema = {
            'type': 'object',
            'properties': {'field1': {'type': 'string', 'description': 'a custom field'}},
            'required': ['field1'],
            'additionalProperties': False,
        }
        self.assertEqual(parser_with_resolver.parse_pico(schema), expected)

    def test_invalid_input_type(self) -> None:
        """Test error on invalid input type to parse."""
        with self.assertRaises(ValueError):
            self.parser.parse_pico(123)


class TestPicoschemaToModel(unittest.TestCase):
    def test_none_schema_returns_none(self) -> None:
        result = picoschema.picoschema_to_pydantic_model(None)
        self.assertIsNone(result)

    def test_scalar_string(self) -> None:
        Model = picoschema.picoschema_to_pydantic_model('string')
        self.assertIsNotNone(Model)
        assert Model is not None
        self.assertTrue(issubclass(Model, BaseModel))
        instance = Model(value='hello')
        self.assertEqual(instance.value, 'hello')

    def test_scalar_with_description(self) -> None:
        Model = picoschema.picoschema_to_pydantic_model('integer, an age')
        assert Model is not None
        self.assertTrue(issubclass(Model, BaseModel))
        instance = Model(value=42)
        self.assertEqual(instance.value, 42)

    def test_simple_object(self) -> None:
        schema = {'name': 'string', 'age': 'integer'}
        Model = picoschema.picoschema_to_pydantic_model(schema)
        assert Model is not None
        self.assertTrue(issubclass(Model, BaseModel))
        instance = Model(name='Alice', age=30)
        self.assertEqual(instance.name, 'Alice')
        self.assertEqual(instance.age, 30)

    def test_optional_fields(self) -> None:
        schema = {'name': 'string', 'nickname?': 'string'}
        Model = picoschema.picoschema_to_pydantic_model(schema)
        assert Model is not None
        # Optional field defaults to None
        instance = Model(name='Alice')
        self.assertEqual(instance.name, 'Alice')
        self.assertIsNone(instance.nickname)
        # Optional field can be set
        instance2 = Model(name='Alice', nickname='Ali')
        self.assertEqual(instance2.nickname, 'Ali')

    def test_array_field(self) -> None:
        schema = {'tags(array)': 'string'}
        Model = picoschema.picoschema_to_pydantic_model(schema)
        assert Model is not None
        instance = Model(tags=['a', 'b'])
        self.assertEqual(instance.tags, ['a', 'b'])

    def test_enum_field(self) -> None:
        schema = {'status(enum)': ['ACTIVE', 'INACTIVE']}
        Model = picoschema.picoschema_to_pydantic_model(schema)
        assert Model is not None
        instance = Model(status='ACTIVE')
        self.assertIsNotNone(instance.status)
        self.assertEqual(instance.status.value, 'ACTIVE')

    def test_optional_enum_field(self) -> None:
        schema = {'status?(enum)': ['ACTIVE', 'INACTIVE']}
        Model = picoschema.picoschema_to_pydantic_model(schema)
        assert Model is not None
        # Omitted optional enum defaults to None
        instance = Model()
        self.assertIsNone(instance.status)
        # Can set it
        instance2 = Model(status='ACTIVE')
        self.assertEqual(instance2.status.value, 'ACTIVE')

    def test_nested_object(self) -> None:
        schema = {
            'address(object)': {
                'street': 'string',
                'city': 'string',
            }
        }
        Model = picoschema.picoschema_to_pydantic_model(schema)
        assert Model is not None
        instance = Model(address={'street': '123 Main', 'city': 'NYC'})
        self.assertEqual(instance.address.street, '123 Main')
        self.assertEqual(instance.address.city, 'NYC')

    def test_descriptions_preserved(self) -> None:
        schema = {'name': 'string, User name', 'age': 'integer, User age'}
        Model = picoschema.picoschema_to_pydantic_model(schema)
        assert Model is not None
        json_schema = Model.model_json_schema()
        name_desc = json_schema['properties']['name'].get('description')
        age_desc = json_schema['properties']['age'].get('description')
        self.assertEqual(name_desc, 'User name')
        self.assertEqual(age_desc, 'User age')

    def test_wildcard_allows_extra(self) -> None:
        schema = {'name': 'string', '(*)': 'any'}
        Model = picoschema.picoschema_to_pydantic_model(schema)
        assert Model is not None
        instance = Model(name='Alice', extra_field='whatever')
        self.assertEqual(instance.name, 'Alice')

    def test_no_extra_by_default(self) -> None:
        schema = {'name': 'string'}
        Model = picoschema.picoschema_to_pydantic_model(schema)
        assert Model is not None
        from pydantic import ValidationError

        with self.assertRaises(ValidationError):
            Model(name='Alice', unknown='bad')

    def test_custom_model_name(self) -> None:
        schema = {'name': 'string'}
        Model = picoschema.picoschema_to_pydantic_model(schema, model_name='UserModel')
        assert Model is not None
        self.assertEqual(Model.__name__, 'UserModel')

    def test_complex_nested(self) -> None:
        schema = {
            'id': 'string, Unique identifier',
            'tags(array)': 'string',
            'address(object)': {
                'street': 'string',
                'zip?': 'string',
            },
            'role(enum)': ['ADMIN', 'USER'],
        }
        Model = picoschema.picoschema_to_pydantic_model(schema)
        assert Model is not None
        instance = Model(
            id='u1',
            tags=['dev'],
            address={'street': '123 Main'},
            role='ADMIN',
        )
        self.assertEqual(instance.id, 'u1')
        self.assertEqual(instance.tags, ['dev'])
        self.assertEqual(instance.address.street, '123 Main')
        self.assertIsNone(instance.address.zip)
        self.assertEqual(instance.role.value, 'ADMIN')

    def test_any_field(self) -> None:
        schema = {'data': 'any'}
        Model = picoschema.picoschema_to_pydantic_model(schema)
        assert Model is not None
        instance = Model(data={'key': 'value'})
        self.assertEqual(instance.data, {'key': 'value'})
        instance2 = Model(data=42)
        self.assertEqual(instance2.data, 42)

    def test_optional_array(self) -> None:
        schema = {'items?(array)': 'string'}
        Model = picoschema.picoschema_to_pydantic_model(schema)
        assert Model is not None
        instance = Model()
        self.assertIsNone(instance.items)
        instance2 = Model(items=['a', 'b'])
        self.assertEqual(instance2.items, ['a', 'b'])

    def test_optional_object(self) -> None:
        schema = {'meta?(object)': {'key': 'string'}}
        Model = picoschema.picoschema_to_pydantic_model(schema)
        assert Model is not None
        instance = Model()
        self.assertIsNone(instance.meta)
        instance2 = Model(meta={'key': 'val'})
        self.assertEqual(instance2.meta.key, 'val')


class TestExtractDescription(unittest.TestCase):
    """Extract description tests."""

    def test_extract(self) -> None:
        """Test extracting a description from a string."""
        input_str = 'string, a simple string'
        expected = ('string', 'a simple string')
        result = picoschema.extract_description(input_str)
        self.assertEqual(result, expected)

    def test_extract_no_description(self) -> None:
        """Test extracting a description from a string with no description."""
        input_str = 'string'
        expected = ('string', None)
        result = picoschema.extract_description(input_str)
        self.assertEqual(result, expected)


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