# 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

"""Picoschema parser and JSON Schema compiler.

Picoschema is a compact, YAML-optimized schema definition format designed for
describing structured data in generative AI prompts. It compiles to JSON Schema.

Key Features:
    - Basic types: string, number, integer, boolean, null, any
    - Optional fields with `?` suffix (e.g., `name?: string`)
    - Inline descriptions with comma separator (e.g., `age: integer, User's age`)
    - Arrays: `tags(array): string`
    - Nested objects: `address(object): ...`
    - Enums: `status(enum): [PENDING, APPROVED, REJECTED]`
    - Wildcards: `(*): any` for additional properties

Example:
    ```yaml
    product:
      id: string, Unique identifier
      price: number, Current price
      tags(array): string
      category(enum): [ELECTRONICS, CLOTHING]
    ```

See Also:
    Full Picoschema reference: https://google.github.io/dotprompt/extending/picoschema/
"""

import enum
import logging
import re
from typing import Any, Union, cast

from pydantic import BaseModel, ConfigDict, Field, create_model

from dotpromptz.resolvers import resolve_json_schema
from dotpromptz.typing import JsonSchema, SchemaResolver

JSON_SCHEMA_SCALAR_TYPES = [
    'any',
    'boolean',
    'integer',
    'null',
    'number',
    'string',
]

WILDCARD_PROPERTY_NAME = '(*)'


def _is_json_schema(schema: dict[str, Any]) -> bool:
    """Checks if a schema is already in JSON Schema format.

    Args:
        schema: The schema to check.

    Returns:
        True if the schema is already in JSON Schema format, False otherwise.
    """
    types = JSON_SCHEMA_SCALAR_TYPES + ['object', 'array']
    return (
        isinstance(schema, dict)  # force format
        and isinstance(schema.get('type'), str)  # force format
        and schema['type'] in types  # force format
    )


def picoschema_to_json_schema(schema: Any, schema_resolver: SchemaResolver | None = None) -> JsonSchema | None:
    """Parses a Picoschema definition into a JSON Schema.

    Args:
        schema: The Picoschema definition (can be a dict or string).
        schema_resolver: Optional callable to resolve named schema references.

    Returns:
        The equivalent JSON Schema, or None if the input schema is None.
    """
    return PicoschemaParser(schema_resolver).parse(schema)


class PicoschemaParser:
    """Parses Picoschema definitions into JSON Schema.

    Handles basic types, optional fields, descriptions, arrays, objects,
    enums, wildcards, and named schema resolution.
    """

    def __init__(self, schema_resolver: SchemaResolver | None = None):
        """Initializes the PicoschemaParser.

        Args:
            schema_resolver: Optional callable to resolve named schema references.
        """
        self._schema_resolver = schema_resolver

    def must_resolve_schema(self, schema_name: str) -> JsonSchema:
        """Resolves a named schema using the configured resolver.

        Args:
            schema_name: The name of the schema to resolve.

        Returns:
            The resolved JSON Schema.

        Raises:
            ValueError: If no schema resolver is configured or the schema
                        name is not found.
        """
        if not self._schema_resolver:
            raise ValueError(f"Picoschema: unsupported scalar type '{schema_name}'.")

        val = resolve_json_schema(schema_name, self._schema_resolver)
        if not val:
            raise ValueError(f"Picoschema: could not find schema with name '{schema_name}'")
        return val

    def parse(self, schema: Any) -> JsonSchema | None:
        """Parses a schema, detecting if it's Picoschema or JSON Schema.

        If the input looks like standard JSON Schema (contains top-level 'type'
        or 'properties'), it's returned directly. Otherwise, it's parsed as
        Picoschema.

        Args:
            schema: The schema definition to parse.

        Returns:
            The resulting JSON Schema, or None if the input is None.
        """
        if not schema:
            return None

        if isinstance(schema, str):
            type_name, description = extract_description(schema)
            if type_name in JSON_SCHEMA_SCALAR_TYPES:
                out: JsonSchema = {'type': type_name}
                if description:
                    out['description'] = description
                return out
            resolved_schema = self.must_resolve_schema(type_name)
            return {**resolved_schema, 'description': description} if description else resolved_schema

        if isinstance(schema, dict) and _is_json_schema(schema):
            return cast(JsonSchema, schema)

        if isinstance(schema, dict) and isinstance(schema.get('properties'), dict):
            return {**cast(JsonSchema, schema), 'type': 'object'}

        # If the schema is not a JSON Schema, parse it as Picoschema.
        return self.parse_pico(schema)

    def parse_pico(self, obj: Any, path: list[str] | None = None) -> JsonSchema:
        """Recursively parses a Picoschema object or string fragment.

        Args:
            obj: The Picoschema fragment (dict or string).
            path: The current path within the schema structure (for error reporting).

        Returns:
            The JSON Schema representation of the fragment.

        Raises:
            ValueError: If the schema structure is invalid.
        """
        if path is None:
            path = []

        if isinstance(obj, str):
            type_name, description = extract_description(obj)
            if type_name not in JSON_SCHEMA_SCALAR_TYPES:
                resolved_schema = self.must_resolve_schema(type_name)
                return {**resolved_schema, 'description': description} if description else resolved_schema

            if type_name == 'any':
                return {'description': description} if description else {}

            return {'type': type_name, 'description': description} if description else {'type': type_name}
        elif not isinstance(obj, dict):
            raise ValueError(f'Picoschema: only consists of objects and strings. Got: {obj}')

        schema: dict[str, Any] = {
            'type': 'object',
            'properties': {},
            'required': [],
            'additionalProperties': False,
        }

        for key, value in obj.items():
            if key == WILDCARD_PROPERTY_NAME:
                schema['additionalProperties'] = self.parse_pico(value, [*path, key])
                continue

            parts = key.split('(')
            name = parts[0]
            type_info = parts[1][:-1] if len(parts) > 1 else None
            is_optional = name.endswith('?')
            property_name = name[:-1] if is_optional else name

            if not is_optional:
                schema['required'].append(property_name)

            if not type_info:
                prop = self.parse_pico(value, [*path, key])
                if is_optional and isinstance(prop.get('type'), str):
                    prop['type'] = [prop['type'], 'null']
                schema['properties'][property_name] = prop
                continue

            type_name, description = extract_description(type_info)
            if type_name == 'array':
                prop = self.parse_pico(value, [*path, key])
                schema['properties'][property_name] = {
                    'type': ['array', 'null'] if is_optional else 'array',
                    'items': prop,
                }
            elif type_name == 'object':
                prop = self.parse_pico(value, [*path, key])
                if is_optional:
                    # prop['type'] may already be a list (e.g. from a nested
                    # optional object), so avoid producing [[...], 'null'].
                    existing_type = prop.get('type', 'object')
                    if isinstance(existing_type, list):
                        if 'null' not in existing_type:
                            existing_type.append('null')
                        prop['type'] = existing_type
                    else:
                        prop['type'] = [existing_type, 'null']
                schema['properties'][property_name] = prop
            elif type_name == 'enum':
                prop = {'enum': value}
                if is_optional and None not in prop['enum']:
                    prop['enum'].append(None)
                schema['properties'][property_name] = prop
            else:
                raise ValueError(f"Picoschema: parenthetical types must be 'object' or 'array', got: {type_name}")

            if description:
                schema['properties'][property_name]['description'] = description

        if not schema['required']:
            del schema['required']
        return schema


def extract_description(input_str: str) -> tuple[str, str | None]:
    """Extracts the type/name and optional description from a Picoschema string.

    Splits a string like "type, description" into ("type", "description").

    Args:
        input_str: The Picoschema string definition.

    Returns:
        A tuple containing the type/name and the description (or None).
    """
    if ',' not in input_str:
        return input_str, None

    match = re.match(r'(.*?), *(.*)$', input_str)
    if match:
        return match.group(1), match.group(2)
    else:
        return input_str, None


# ---------------------------------------------------------------------------
# JSON Schema -> Python type mapping for Pydantic model generation
# ---------------------------------------------------------------------------

_JSON_TYPE_MAP: dict[str, type] = {
    'string': str,
    'number': float,
    'integer': int,
    'boolean': bool,
    'null': type(None),
}


def _resolve_json_schema_type(
    schema: dict[str, Any],
    name: str,
    parent_name: str,
) -> type:
    """Resolve a JSON Schema node to a Python type for Pydantic fields.

    Handles scalars, enums, arrays, nested objects, nullable types,
    and ``any`` (no type constraint).

    Args:
        schema: A JSON Schema node.
        name: The field/property name (used for nested model naming).
        parent_name: The parent model name (used for nested model naming).

    Returns:
        A Python type suitable for ``pydantic.create_model``.
    """
    # Enum
    if 'enum' in schema:
        values = [v for v in schema['enum'] if v is not None]
        has_none = None in schema['enum']
        if values:
            members: dict[str, Any] = {}
            for v in values:
                member_name = str(v).upper().replace(' ', '_').replace('-', '_')
                base_member = member_name
                idx = 1
                while member_name in members:
                    member_name = f'{base_member}_{idx}'
                    idx += 1
                members[member_name] = v
            enum_cls = enum.Enum(f'{parent_name}_{name}_enum', members)  # type: ignore[misc]
            if has_none:
                return enum_cls | None  # type: ignore[return-value]
            return enum_cls  # type: ignore[return-value]
        return type(None)

    json_type = schema.get('type')

    # Nullable / union types: {"type": ["string", "null"]}
    if isinstance(json_type, list):
        py_types: list[type] = []
        for t in json_type:
            if t == 'null':
                py_types.append(type(None))
            elif t == 'object':
                if 'properties' in schema:
                    py_types.append(_build_pydantic_model(schema, f'{parent_name}_{name}'))
                else:
                    py_types.append(dict)
            elif t == 'array':
                items = schema.get('items', {})
                if items:
                    item_type = _resolve_json_schema_type(items, f'{name}_item', parent_name)
                    py_types.append(list[item_type])  # type: ignore[valid-type]
                else:
                    py_types.append(list)
            else:
                mapped = _JSON_TYPE_MAP.get(t)
                if mapped is not None:
                    py_types.append(mapped)
                else:
                    py_types.append(Any)  # type: ignore[arg-type]
        if len(py_types) == 1:
            return py_types[0]
        return Union[tuple(py_types)]  # type: ignore[return-value]

    # No type specified — "any"
    if json_type is None:
        if 'properties' in schema:
            return _build_pydantic_model(schema, f'{parent_name}_{name}')
        return Any  # type: ignore[return-value]

    # Object
    if json_type == 'object':
        if 'properties' in schema:
            return _build_pydantic_model(schema, f'{parent_name}_{name}')
        additional = schema.get('additionalProperties')
        if isinstance(additional, dict):
            val_type = _resolve_json_schema_type(additional, name, parent_name)
            return dict[str, val_type]  # type: ignore[valid-type]
        return dict  # type: ignore[return-value]

    # Array
    if json_type == 'array':
        items = schema.get('items', {})
        if items:
            item_type = _resolve_json_schema_type(items, f'{name}_item', parent_name)
            return list[item_type]  # type: ignore[valid-type]
        return list  # type: ignore[return-value]

    # Scalar
    mapped = _JSON_TYPE_MAP.get(json_type)
    if mapped is not None:
        return mapped

    return Any  # type: ignore[return-value]


def _build_pydantic_model(
    schema: dict[str, Any],
    model_name: str,
) -> type[BaseModel]:
    """Build a Pydantic model from a JSON Schema object node.

    Args:
        schema: A JSON Schema dict with ``type: 'object'`` and ``properties``.
        model_name: The name for the generated Pydantic model class.

    Returns:
        A dynamically created ``BaseModel`` subclass.
    """
    properties = schema.get('properties', {})
    required_fields = set(schema.get('required', []))
    additional_props = schema.get('additionalProperties', False)

    fields: dict[str, Any] = {}

    for prop_name, prop_schema in properties.items():
        py_type = _resolve_json_schema_type(prop_schema, prop_name, model_name)
        description = prop_schema.get('description')

        field_kwargs: dict[str, Any] = {}
        if description:
            field_kwargs['description'] = description

        if prop_name in required_fields:
            fields[prop_name] = (py_type, Field(..., **field_kwargs))
        else:
            # Optional field — wrap in Optional if not already nullable
            if not _type_allows_none(py_type):
                py_type = py_type | None  # type: ignore[assignment]
            fields[prop_name] = (py_type, Field(default=None, **field_kwargs))

    extra = 'allow' if additional_props is not False else 'forbid'

    config = ConfigDict(extra=extra)  # type: ignore[typeddict-item]

    return create_model(
        model_name,
        __config__=config,
        **fields,
    )


def _type_allows_none(py_type: Any) -> bool:
    """Check if a type annotation already allows None."""
    origin = getattr(py_type, '__origin__', None)
    if origin is Union:
        return type(None) in py_type.__args__
    return py_type is type(None)


def picoschema_to_pydantic_model(
    schema: Any,
    schema_resolver: SchemaResolver | None = None,
    *,
    model_name: str = 'PicoschemaModel',
) -> type[BaseModel] | None:
    """Convert a Picoschema definition to a dynamic Pydantic ``BaseModel`` class.

    Internally converts the picoschema to JSON Schema first, then builds a
    Pydantic model from the resulting schema using ``pydantic.create_model()``.

    Supports all picoschema features: scalar types, objects, arrays, enums,
    optional fields (nullable), nested objects, wildcards, and descriptions.

    Args:
        schema: The Picoschema definition (can be a dict or string).
        schema_resolver: Optional callable to resolve named schema references.
        model_name: The name for the top-level generated model class.

    Returns:
        A dynamically created ``BaseModel`` subclass, or ``None`` if the
        input schema is ``None``.

    Example::

        schema = {
            'name': 'string, User name',
            'age?': 'integer, Age in years',
            'tags(array)': 'string',
            'role(enum)': ['ADMIN', 'USER'],
        }
        Model = picoschema_to_pydantic_model(schema)
        instance = Model(name='Alice', tags=['dev'], role='ADMIN')
    """
    json_schema = picoschema_to_json_schema(schema, schema_resolver)
    if json_schema is None:
        return None

    # Scalar schemas (e.g. picoschema='string') don't produce object types.
    # Wrap them in a single-field model for usability.
    if json_schema.get('type') != 'object' and 'properties' not in json_schema:
        scalar_type = _resolve_json_schema_type(json_schema, 'value', model_name)
        description = json_schema.get('description')
        field_kwargs: dict[str, Any] = {}
        if description:
            field_kwargs['description'] = description
        return create_model(
            model_name,
            value=(scalar_type, Field(..., **field_kwargs)),
        )

    return _build_pydantic_model(json_schema, model_name)


# ---------------------------------------------------------------------------
# JSON Schema -> Picoschema reverse conversion
# ---------------------------------------------------------------------------

_reverse_logger = logging.getLogger(__name__)

# JSON Schema scalar types that map directly to Picoschema type strings.
_SCALAR_TYPES = frozenset({'string', 'number', 'integer', 'boolean', 'null'})


def json_schema_to_picoschema(schema: JsonSchema) -> Any:
    """Convert a JSON Schema to Picoschema notation.

    Args:
        schema: A JSON Schema dict.

    Returns:
        A Picoschema representation (str, dict, or list depending on the
        schema structure). Returns ``None`` for empty or unsupported schemas.
    """
    if not schema or not isinstance(schema, dict):
        return None

    return _convert_node(schema, required=True)


def _convert_node(node: dict[str, Any], required: bool = True) -> Any:
    """Recursively convert a JSON Schema node to Picoschema.

    Args:
        node: A JSON Schema node.
        required: Whether this node is required by its parent.

    Returns:
        Picoschema representation of the node.
    """
    schema_type = node.get('type')
    description = node.get('description')

    # Handle nullable types: {"type": ["string", "null"]} -> optional string
    if isinstance(schema_type, list):
        non_null = [t for t in schema_type if t != 'null']
        _ = 'null' in schema_type  # nullable detection (unused for now)
        schema_type = non_null[0] if len(non_null) == 1 else None

    # Enum
    if 'enum' in node:
        enum_values = [v for v in node['enum'] if v is not None]
        return enum_values

    # Scalar types
    if isinstance(schema_type, str) and schema_type in _SCALAR_TYPES:
        type_str = schema_type
        if description:
            type_str = f'{schema_type}, {description}'
        return type_str

    # "any" type (no type specified, no properties)
    if schema_type is None and 'properties' not in node and 'items' not in node:
        if description:
            return f'any, {description}'
        return 'any'

    # Array
    if schema_type == 'array':
        items = node.get('items', {})
        if items:
            return _convert_node(items)
        return 'any'

    # Object
    if schema_type == 'object' or 'properties' in node:
        return _convert_object(node)

    # Fallback: return the type as a string if known
    if isinstance(schema_type, str):
        if description:
            return f'{schema_type}, {description}'
        return schema_type

    _reverse_logger.warning('json_schema_to_picoschema: unsupported schema node, node=%s', node)
    return None


def _convert_object(node: dict[str, Any]) -> dict[str, Any]:
    """Convert a JSON Schema object node to a Picoschema dict.

    Args:
        node: A JSON Schema object node.

    Returns:
        A Picoschema dict.
    """
    properties = node.get('properties', {})
    required_fields = set(node.get('required', []))
    additional = node.get('additionalProperties')
    result: dict[str, Any] = {}

    for prop_name, prop_schema in properties.items():
        is_required = prop_name in required_fields
        prop_type = prop_schema.get('type')

        # Detect nullable from type list
        if isinstance(prop_type, list):
            non_null = [t for t in prop_type if t != 'null']
            _ = 'null' in prop_type  # nullable detection (unused for now)
            prop_type = non_null[0] if len(non_null) == 1 else None

        # Build the key: add ? suffix for optional fields
        key = prop_name if is_required else f'{prop_name}?'

        # Enum property
        if 'enum' in prop_schema:
            enum_values = [v for v in prop_schema['enum'] if v is not None]
            key_with_type = f'{key}(enum)'
            result[key_with_type] = enum_values
            # Note: Picoschema doesn't support per-enum descriptions in the key
            # notation, so any description on the enum schema is intentionally dropped.
            continue

        # Array property
        if prop_type == 'array':
            items = prop_schema.get('items', {})
            key_with_type = f'{key}(array)'
            result[key_with_type] = _convert_node(items) if items else 'any'
            continue

        # Nested object property
        if prop_type == 'object' or 'properties' in prop_schema:
            key_with_type = f'{key}(object)'
            result[key_with_type] = _convert_object(prop_schema)
            continue

        # Scalar or reference
        value = _convert_node(prop_schema, required=is_required)
        result[key] = value

    # Wildcard for additionalProperties
    if additional and isinstance(additional, dict):
        result['(*)'] = _convert_node(additional)
    elif additional is True:
        result['(*)'] = 'any'

    return result
