"""Picoschema parser and JSON Schema compiler.

Picoschema is a compact, YAML-optimized schema format that compiles to
JSON Schema. Supports basic types, optional fields, arrays, nested objects,
enums, inline descriptions, and wildcard properties.
"""

import enum
import functools
import operator
import re
import types
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 py_types[0] | py_types[1] if len(py_types) == 2 else functools.reduce(operator.or_, 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 or isinstance(py_type, types.UnionType):
        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)
