<?php

namespace Laravel\Ai\Schema;

use Illuminate\JsonSchema\Types\AnyOfType;

class SchemaNormalizer
{
    /**
     * Type-specific constraint keywords the deserializer rejects on a multi-type union.
     */
    private const TYPE_KEYWORDS = [
        'minLength', 'maxLength', 'pattern', 'format',
        'minimum', 'maximum', 'multipleOf',
        'items', 'minItems', 'maxItems', 'uniqueItems',
        'properties', 'required', 'additionalProperties',
    ];

    /**
     * Keywords the deserializer cannot represent and that are dropped.
     */
    private const UNSUPPORTED = [
        '$schema', '$id', '$anchor', '$comment', 'not', 'if', 'then', 'else',
        'patternProperties', 'dependentSchemas', 'dependentRequired', 'unevaluatedProperties',
        'contains', 'minContains', 'maxContains', 'prefixItems', 'examples', 'deprecated',
        'readOnly', 'writeOnly', 'minProperties', 'maxProperties', 'exclusiveMinimum', 'exclusiveMaximum',
    ];

    /**
     * Scalar types that can be combined into a multi-type union.
     */
    private const SCALAR_TYPES = ['string', 'integer', 'number', 'boolean'];

    /**
     * The JSON Schema type strings the deserializer understands.
     */
    private const TYPES = ['string', 'integer', 'number', 'boolean', 'object', 'array', 'null'];

    /**
     * Rewrite a raw JSON Schema into the subset Illuminate\JsonSchema can deserialize.
     *
     * @param  array<string, mixed>  $schema
     * @return array<string, mixed>
     */
    public static function normalize(array $schema): array
    {
        return (new self)->node($schema, $schema, []);
    }

    /**
     * Normalize a single schema node into the deserializable subset.
     *
     * @param  array<string, mixed>  $schema
     * @param  array<string, mixed>  $root
     * @param  array<string, true>  $seen
     * @return array<string, mixed>
     */
    private function node(array $schema, array $root, array $seen): array
    {
        [$schema, $seen] = $this->inlineRefs($schema, $root, $seen);

        $schema = $this->mergeAllOf($schema, $root, $seen);
        $schema = $this->preserveAnyOf($schema, $root, $seen);
        $schema = $this->collapseUnions($schema, $root, $seen);
        $schema = $this->collapseMultiType($schema);
        $schema = $this->normalizeKeywords($schema);
        $schema = $this->normalizeChildren($schema, $root, $seen);

        return $this->ensureType($schema);
    }

    /**
     * Inline local "$ref" pointers (cycle-safe); drop remote or unresolvable ones.
     *
     * @param  array<string, mixed>  $schema
     * @param  array<string, mixed>  $root
     * @param  array<string, true>  $seen
     * @return array{0: array<string, mixed>, 1: array<string, true>}
     */
    private function inlineRefs(array $schema, array $root, array $seen): array
    {
        while (isset($schema['$ref']) && is_string($schema['$ref'])) {
            $ref = $schema['$ref'];

            unset($schema['$ref']);

            if (isset($seen[$ref]) || ($resolved = $this->lookupRef($ref, $root)) === null) {
                break;
            }

            $seen[$ref] = true;
            $schema = array_merge($resolved, $schema);
        }

        return [$schema, $seen];
    }

    /**
     * Resolve a local JSON pointer against the root schema.
     *
     * @param  array<string, mixed>  $root
     * @return array<string, mixed>|null
     */
    private function lookupRef(string $ref, array $root): ?array
    {
        if (! str_starts_with($ref, '#/')) {
            return null;
        }

        $target = $root;

        foreach (explode('/', substr($ref, 2)) as $segment) {
            $segment = str_replace(['~1', '~0'], ['/', '~'], rawurldecode($segment));

            if (! is_array($target) || ! array_key_exists($segment, $target)) {
                return null;
            }

            $target = $target[$segment];
        }

        return is_array($target) ? $target : null;
    }

    /**
     * Flatten "allOf" branches into the node; the deserializer has no intersection type.
     *
     * @param  array<string, mixed>  $schema
     * @param  array<string, mixed>  $root
     * @param  array<string, true>  $seen
     * @return array<string, mixed>
     */
    private function mergeAllOf(array $schema, array $root, array $seen): array
    {
        while (is_array($schema['allOf'] ?? null)) {
            $branches = $schema['allOf'];

            unset($schema['allOf']);

            $merged = [];

            foreach ($branches as $branch) {
                if (is_array($branch)) {
                    [$branch, $seen] = $this->inlineRefs($branch, $root, $seen);
                    $merged = $this->mergeSchema($merged, $branch);
                }
            }

            $schema = $this->mergeSchema($merged, $schema);
        }

        return $schema;
    }

    /**
     * Merge two schema fragments, unioning "required" and recursively merging "properties".
     *
     * @param  array<string, mixed>  $base
     * @param  array<string, mixed>  $overlay
     * @return array<string, mixed>
     */
    private function mergeSchema(array $base, array $overlay): array
    {
        $required = array_values(array_unique(array_merge(
            is_array($base['required'] ?? null) ? $base['required'] : [],
            is_array($overlay['required'] ?? null) ? $overlay['required'] : [],
        )));

        $properties = is_array($base['properties'] ?? null) ? $base['properties'] : [];

        foreach (is_array($overlay['properties'] ?? null) ? $overlay['properties'] : [] as $key => $value) {
            $properties[$key] = isset($properties[$key]) && is_array($properties[$key]) && is_array($value)
                ? $this->mergeSchema($properties[$key], $value)
                : $value;
        }

        $merged = array_merge($base, $overlay);

        if ($required !== []) {
            $merged['required'] = $required;
        }

        if ($properties !== []) {
            $merged['properties'] = $properties;
        }

        return $merged;
    }

    /**
     * Collapse "anyOf"/"oneOf" into a deserializable form.
     *
     * @param  array<string, mixed>  $schema
     * @param  array<string, mixed>  $root
     * @param  array<string, true>  $seen
     * @return array<string, mixed>
     */
    private function collapseUnions(array $schema, array $root, array $seen): array
    {
        foreach (['anyOf', 'oneOf'] as $key) {
            if (! is_array($schema[$key] ?? null)) {
                continue;
            }

            if ($key === 'anyOf' && $this->supportsAnyOf()) {
                continue;
            }

            $branches = $schema[$key];

            unset($schema[$key]);

            $schema = $this->mergeUnion($schema, $branches, $root, $seen);
        }

        return $schema;
    }

    /**
     * Keep anyOf compositions intact when the installed deserializer supports them.
     *
     * @param  array<string, mixed>  $schema
     * @param  array<string, mixed>  $root
     * @param  array<string, true>  $seen
     * @return array<string, mixed>
     */
    private function preserveAnyOf(array $schema, array $root, array $seen): array
    {
        if (! $this->supportsAnyOf() || ! is_array($schema['anyOf'] ?? null)) {
            return $schema;
        }

        $carry = ['anyOf' => true, 'title' => true, 'description' => true, 'enum' => true, 'default' => true];
        $base = array_diff_key($schema, $carry);

        $branches = [];
        $resolved = false;

        foreach ($schema['anyOf'] as $branch) {
            if (! is_array($branch)) {
                continue;
            }
            if ($branch === []) {
                continue;
            }
            [$branch, $branchSeen] = $this->inlineRefs($branch, $root, $seen);

            if ($branch === []) {
                continue;
            }

            if ($this->isNullBranch($branch)) {
                $branches[] = ['type' => 'null'];

                continue;
            }

            $branches[] = $this->node($base === [] ? $branch : $this->mergeSchema($base, $branch), $root, $branchSeen);
            $resolved = true;
        }

        if (! $resolved) {
            unset($schema['anyOf']);

            return $schema;
        }

        $schema = array_intersect_key($schema, $carry);
        $schema['anyOf'] = $branches;

        return $schema;
    }

    /**
     * Reduce union branches into the node: scalar multi-type union, merged object variants,
     * or the first branch (lossy fallback) since the deserializer only accepts a single schema plus null.
     *
     * @param  array<string, mixed>  $schema
     * @param  array<int, mixed>  $branches
     * @param  array<string, mixed>  $root
     * @param  array<string, true>  $seen
     * @return array<string, mixed>
     */
    private function mergeUnion(array $schema, array $branches, array $root, array $seen): array
    {
        $nullable = false;
        $resolved = [];

        foreach ($branches as $branch) {
            if (! is_array($branch)) {
                continue;
            }

            [$branch, $branchSeen] = $this->inlineRefs($branch, $root, $seen);

            if ($branch === []) {
                continue;
            }

            if ($this->isNullBranch($branch)) {
                $nullable = true;
            } else {
                $resolved[] = $this->node($branch, $root, $branchSeen);
            }
        }

        if (($scalarTypes = $this->scalarUnionTypes($resolved)) !== null) {
            $schema['type'] = array_values(array_unique($nullable ? [...$scalarTypes, 'null'] : $scalarTypes));

            return $schema;
        }

        if ($resolved !== []) {
            $schema = $this->mergeObjectVariants($schema, $resolved)
                ?? $this->mergeSchema($schema, $this->firstObjectBranch($resolved) ?? $resolved[0]);
        }

        if (! $nullable) {
            return $schema;
        }

        $type = $schema['type'] ?? $this->baseType($schema);
        $schema['type'] = array_values(array_unique([...(is_array($type) ? $type : [$type]), 'null']));

        return $schema;
    }

    /**
     * Union every plain-object variant's properties, ignoring scalar branches; null when fewer than two objects remain.
     *
     * @param  array<string, mixed>  $schema
     * @param  array<int, array<string, mixed>>  $branches
     * @return array<string, mixed>|null
     */
    private function mergeObjectVariants(array $schema, array $branches): ?array
    {
        $objects = array_values(array_filter($branches, $this->isObjectBranch(...)));

        if (count($objects) < 2) {
            return null;
        }

        return $this->reduceObjectGroup($objects, $schema);
    }

    /**
     * Merge a group of object variants into one object, layering an optional conjunctive outer schema on top.
     *
     * @param  array<int, array<string, mixed>>  $variants
     * @param  array<string, mixed>  $outer
     * @return array<string, mixed>
     */
    private function reduceObjectGroup(array $variants, array $outer = []): array
    {
        $propertyLists = [];
        $requiredSets = [];
        $nullable = false;
        $description = null;
        $closed = true;

        foreach ($variants as $variant) {
            $nullable = $nullable || is_array($variant['type'] ?? null);

            foreach (is_array($variant['properties'] ?? null) ? $variant['properties'] : [] as $key => $prop) {
                $propertyLists[$key][] = $prop;
            }

            if (is_array($variant['required'] ?? null)) {
                $requiredSets[] = $variant['required'];
            }

            if ($description === null && is_string($variant['description'] ?? null)) {
                $description = $variant['description'];
            }

            $closed = $closed && (($variant['additionalProperties'] ?? null) === false);
        }

        $properties = [];

        foreach ($propertyLists as $key => $list) {
            $properties[$key] = $this->mergeVariantValues($list);
        }

        foreach (is_array($outer['properties'] ?? null) ? $outer['properties'] : [] as $key => $prop) {
            $properties[$key] = isset($properties[$key]) && is_array($properties[$key]) && is_array($prop)
                ? $this->mergeVariantValues([$properties[$key], $prop])
                : $prop;
        }

        $result = $outer;
        $result['type'] = $nullable ? ['object', 'null'] : 'object';

        if ($properties !== []) {
            $result['properties'] = $properties;
        } else {
            unset($result['properties']);
        }

        $outerRequired = is_array($outer['required'] ?? null) ? $outer['required'] : [];
        $branchRequired = count($requiredSets) >= 2 ? array_values(array_intersect(...$requiredSets)) : [];
        $required = array_values(array_unique(array_merge($outerRequired, $branchRequired)));

        if ($required !== []) {
            $result['required'] = $required;
        } else {
            unset($result['required']);
        }

        if ($description !== null && ! isset($result['description'])) {
            $result['description'] = $description;
        }

        if ($closed && ! isset($result['additionalProperties'])) {
            $result['additionalProperties'] = false;
        }

        return $result;
    }

    /**
     * Reduce the schemas a single property takes across variants into one schema, recursing on object variants.
     *
     * @param  array<int, mixed>  $list
     * @return array<string, mixed>
     */
    private function mergeVariantValues(array $list): array
    {
        $list = array_values(array_filter($list, is_array(...)));

        if (count($list) <= 1) {
            return $list[0] ?? [];
        }

        $objects = array_values(array_filter($list, $this->isObjectBranch(...)));

        if (count($objects) >= 2) {
            return $this->reduceObjectGroup($objects);
        }

        if (count($objects) === 1) {
            return $objects[0];
        }

        return $this->unionScalarValues($list);
    }

    /**
     * Union a property's scalar variants, combining their type and enum constraints.
     *
     * @param  array<int, array<string, mixed>>  $list
     * @return array<string, mixed>
     */
    private function unionScalarValues(array $list): array
    {
        $merged = [];
        $types = [];
        $enums = [];

        foreach ($list as $value) {
            $merged = array_merge($merged, $value);
            $types = [...$types, ...$this->typeList($value['type'] ?? null)];

            if (is_array($value['enum'] ?? null)) {
                $enums = [...$enums, ...$value['enum']];
            }
        }

        $types = array_values(array_unique(array_filter($types, fn ($type): bool => $type !== null)));

        if (! isset($merged['items']) && $types !== []) {
            $merged['type'] = count($types) === 1 ? $types[0] : $types;
        }

        if ($enums !== []) {
            $merged['enum'] = array_values(array_unique($enums));
        }

        return $merged;
    }

    /**
     * Normalize a "type" declaration into a list of type strings.
     *
     * @return array<int, mixed>
     */
    private function typeList(mixed $type): array
    {
        return is_array($type) ? $type : [$type];
    }

    /**
     * Find the first branch whose type resolves to a plain object, used as a structure-preserving fallback.
     *
     * @param  array<int, array<string, mixed>>  $branches
     * @return array<string, mixed>|null
     */
    private function firstObjectBranch(array $branches): ?array
    {
        foreach ($branches as $branch) {
            if ($this->isObjectBranch($branch)) {
                return $branch;
            }
        }

        return null;
    }

    /**
     * Determine whether a normalized branch's type is a plain object, optionally nullable.
     *
     * @param  array<string, mixed>  $branch
     */
    private function isObjectBranch(array $branch): bool
    {
        $type = $branch['type'] ?? null;
        $nonNull = array_values(array_filter(is_array($type) ? $type : [$type], fn ($value): bool => $value !== 'null'));

        return $nonNull === ['object'];
    }

    /**
     * Determine if the installed Illuminate JSON schema package supports anyOf.
     */
    private function supportsAnyOf(): bool
    {
        return class_exists(AnyOfType::class);
    }

    /**
     * Determine whether the branch only represents null.
     *
     * @param  array<string, mixed>  $schema
     */
    private function isNullBranch(array $schema): bool
    {
        return in_array($schema['type'] ?? null, ['null', ['null']], true);
    }

    /**
     * Get the unioned scalar types when every branch is a plain scalar, else null.
     *
     * @param  array<int, array<string, mixed>>  $branches
     * @return array<int, string>|null
     */
    private function scalarUnionTypes(array $branches): ?array
    {
        if (count($branches) < 2) {
            return null;
        }

        $types = [];

        foreach ($branches as $branch) {
            $branchTypes = is_array($branch['type'] ?? null) ? $branch['type'] : [$branch['type'] ?? null];

            foreach ($branchTypes as $type) {
                if (! in_array($type, self::SCALAR_TYPES, true)) {
                    return null;
                }

                $types[] = $type;
            }
        }

        return $types;
    }

    /**
     * Keep multi-type unions deserializable and drop unrepresentable "null"-only types.
     *
     * @param  array<string, mixed>  $schema
     * @return array<string, mixed>
     */
    private function collapseMultiType(array $schema): array
    {
        $type = $schema['type'] ?? null;

        if (! is_array($type)) {
            if ($type === 'null' || ($type !== null && ! in_array($type, self::TYPES, true))) {
                unset($schema['type']);
            }

            return $schema;
        }

        $valid = array_values(array_filter($type, fn ($value): bool => in_array($value, self::TYPES, true)));
        $nonNull = array_values(array_filter($valid, fn ($value): bool => $value !== 'null'));

        if ($nonNull === []) {
            unset($schema['type']);

            return $schema;
        }

        $schema['type'] = count($valid) === 1 ? $valid[0] : $valid;

        if (count($nonNull) > 1) {
            foreach (self::TYPE_KEYWORDS as $keyword) {
                unset($schema[$keyword]);
            }
        }

        return $schema;
    }

    /**
     * Drop unsupported keywords, rewrite const to enum, and normalize additionalProperties.
     *
     * @param  array<string, mixed>  $schema
     * @return array<string, mixed>
     */
    private function normalizeKeywords(array $schema): array
    {
        foreach (self::UNSUPPORTED as $keyword) {
            unset($schema[$keyword]);
        }

        unset($schema['$defs'], $schema['definitions']);

        if (array_key_exists('const', $schema)) {
            $schema['enum'] = [$schema['const']];
            unset($schema['const']);
        }

        if (array_key_exists('default', $schema) && $schema['default'] === null) {
            unset($schema['default']);
        }

        if (! isset($schema['type']) && array_key_exists('additionalProperties', $schema)) {
            $schema['type'] = 'object';
        }

        if (($schema['additionalProperties'] ?? false) !== false) {
            unset($schema['additionalProperties']);
        }

        return $schema;
    }

    /**
     * Recurse into "properties" (pruning orphaned "required") and "items".
     *
     * @param  array<string, mixed>  $schema
     * @param  array<string, mixed>  $root
     * @param  array<string, true>  $seen
     * @return array<string, mixed>
     */
    private function normalizeChildren(array $schema, array $root, array $seen): array
    {
        if (is_array($schema['properties'] ?? null)) {
            $properties = [];

            foreach ($schema['properties'] as $key => $definition) {
                if (is_array($definition)) {
                    $properties[$key] = $this->node($definition, $root, $seen);
                }
            }

            $schema['properties'] = $properties;

            if (is_array($schema['required'] ?? null)) {
                $schema['required'] = array_values(array_filter(
                    $schema['required'],
                    fn ($name): bool => is_string($name) && array_key_exists($name, $properties),
                ));
            }
        }

        if (isset($schema['items'])) {
            $schema['items'] = is_array($schema['items']) && ! array_is_list($schema['items'])
                ? $this->node($schema['items'], $root, $seen)
                : null;

            if ($schema['items'] === null) {
                unset($schema['items']);
            }
        }

        return $schema;
    }

    /**
     * Give a node an explicit type so a typeless node never reaches the deserializer.
     *
     * @param  array<string, mixed>  $schema
     * @return array<string, mixed>
     */
    private function ensureType(array $schema): array
    {
        if (isset($schema['anyOf']) || isset($schema['oneOf']) || isset($schema['allOf'])) {
            return $schema;
        }

        $type = $schema['type'] ?? null;
        $nonNull = array_filter(is_array($type) ? $type : [$type], fn ($t): bool => $t !== 'null');

        if ($nonNull !== [] && array_filter($nonNull, fn ($t): bool => ! in_array($t, self::TYPES, true)) === []) {
            return $schema;
        }

        unset($schema['type']);
        $schema['type'] = $this->baseType($schema);

        return $schema;
    }

    /**
     * Infer a single base type from a node's shape, defaulting to "string".
     *
     * @param  array<string, mixed>  $schema
     */
    private function baseType(array $schema): string
    {
        return match (true) {
            isset($schema['properties']), isset($schema['required']), isset($schema['additionalProperties']) => 'object',
            isset($schema['items']), isset($schema['minItems']), isset($schema['maxItems']), isset($schema['uniqueItems']) => 'array',
            isset($schema['enum']) && is_array($schema['enum']) => $this->inferEnumType($schema['enum']),
            isset($schema['minimum']), isset($schema['maximum']), isset($schema['multipleOf']) => 'number',
            default => 'string',
        };
    }

    /**
     * Infer the scalar type shared by an enum, defaulting to "string".
     *
     * @param  array<int, mixed>  $enum
     */
    private function inferEnumType(array $enum): string
    {
        $resolved = null;

        foreach ($enum as $value) {
            $current = match (true) {
                is_bool($value) => 'boolean',
                is_int($value) => 'integer',
                is_float($value) => 'number',
                is_string($value) => 'string',
                default => null,
            };

            if ($current === null) {
                return 'string';
            }

            if ($resolved === null || $resolved === $current) {
                $resolved = $current;

                continue;
            }

            if (in_array($resolved, ['integer', 'number'], true) && in_array($current, ['integer', 'number'], true)) {
                $resolved = 'number';

                continue;
            }

            return 'string';
        }

        return $resolved ?? 'string';
    }
}