from typing import Mapping

from jschon import JSON, JSONPointer, JSONSchema, URI, create_catalog
from jschon.json import JSONCompatible
from jschon.jsonschema import Result
from jschon.vocabulary import Keyword
from jschon.vocabulary.applicator import AnyOfKeyword, OneOfKeyword


class DiscriminatorKeyword(Keyword):
    key = 'discriminator'
    depends_on = 'anyOf', 'oneOf',
    instance_types = 'object',

    def __init__(self, parentschema: JSONSchema, value: Mapping[str, JSONCompatible]):
        super().__init__(parentschema, value)
        self.any_of_targets = {
            JSONPointer.parse_uri_fragment(URI(subschema['$ref'].value).fragment)[-1]:
                (index, subschema)
            for index, subschema in enumerate(parentschema['anyOf'])
        }

    def evaluate(self, instance: JSON, result: Result) -> None:
        discriminator_property = self.json['propertyName'].value
        discriminator = instance[discriminator_property].value
        target_index, target_schema = self.any_of_targets[discriminator]
        with result(instance, str(target_index)) as subresult:
            if not target_schema.evaluate(instance, subresult).passed:
                result.fail(f'The instance is invalid against the {discriminator} subschema')


class DiscriminatedAnyOfKeyword(AnyOfKeyword):
    key = 'anyOf'

    def evaluate(self, instance: JSON, result: Result) -> None:
        if not self.parentschema.get('discriminator'):
            super().evaluate(instance, result)


class DiscriminatedOneOfKeyword(OneOfKeyword):
    key = 'oneOf'

    def evaluate(self, instance: JSON, result: Result) -> None:
        if not self.parentschema.get('discriminator'):
            super().evaluate(instance, result)


catalog = create_catalog('2020-12')

metaschema = catalog.get_schema(URI('https://json-schema.org/draft/2020-12/schema'))
metaschema.kwclasses['anyOf'] = DiscriminatedAnyOfKeyword
metaschema.kwclasses['oneOf'] = DiscriminatedOneOfKeyword
metaschema.kwclasses['discriminator'] = DiscriminatorKeyword

schema = JSONSchema({
    "$schema": "https://json-schema.org/draft/2020-12/schema",
    "$id": "https://example.com/schema",
    "anyOf": [
        {"$ref": "#/$defs/simpleObject"},
        {"$ref": "#/$defs/complexObject"}
    ],
    "discriminator": {
        "propertyName": "objectType"
    },
    "$defs": {
        "simpleObject": {
            "properties": {
                "objectType": {"const": "simpleObject"},
                "value": {"type": ["number", "string"]}
            }
        },
        "complexObject": {
            "properties": {
                "objectType": {"const": "complexObject"},
                "value": {"type": ["array", "object"]}
            }
        }
    }
})

valid_simple_object = JSON({
    "objectType": "simpleObject",
    "value": 10
})
invalid_simple_object = JSON({
    "objectType": "simpleObject",
    "value": ["nope"]
})
valid_complex_object = JSON({
    "objectType": "complexObject",
    "value": ["super", "complex"]
})
invalid_complex_object = JSON({
    "objectType": "complexObject",
    "value": None
})

print(schema.evaluate(valid_simple_object).output('basic'))
print(schema.evaluate(invalid_simple_object).output('basic'))
print(schema.evaluate(valid_complex_object).output('basic'))
print(schema.evaluate(invalid_complex_object).output('basic'))