#!/usr/bin/env python3
# Copyright 2025 Google LLC
# SPDX-License-Identifier: Apache-2.0

"""JSON Schema -> Pydantic v2 generator for Genkit."""

from __future__ import annotations

import json
import re
import sys
from datetime import datetime
from pathlib import Path

# Do NOT add EvalFnResponse or EvalResponse - they are autogenerated and required by the evaluator API.
EXCLUDED = frozenset({
    'ModelRequest',
    'GenerateRequest',
    'OutputModel',
    'GenerateResponse',
    'Request',
    'ModelResponse',
    'RankedDocumentData',
    'RankedDocumentMetadata',
    'CommonRerankerOptions',
    'RerankerRequest',
    'RerankerResponse',
    'CommonRetrieverOptions',
    'RetrieverRequest',
    'RetrieverResponse',
    # SpanMetadata is hand-written in genkit._core._tracing so we can add Python-only
    # input fields (type, subtype, telemetry_labels) that drive OTel attributes.
    'SpanMetadata',
    # Do NOT add EvalFnResponse or EvalResponse - autogenerated and required by evaluator API
})
PRIM = {'string': 'str', 'number': 'float', 'integer': 'int', 'boolean': 'bool'}
# Schema type transformations: rename and/or omit fields before emission.
# Keys: schema type name. Values: {'output_name': str} and/or {'suffix': str, 'omit': [str]}.
# - output_name: emit and reference as this name (e.g. Message -> MessageData)
# - suffix: emit as {name}{suffix}, omit listed fields (hand-written subclass adds them back)
TRANSFORMATIONS = {
    'Message': {'output_name': 'MessageData'},
    'GenerateActionOptions': {'suffix': 'Data', 'omit': ['messages']},
}


def _output_name(name: str) -> str:
    """Resolve schema type name to output type name for refs and emission."""
    if name not in TRANSFORMATIONS:
        return name
    cfg = TRANSFORMATIONS[name]
    out = cfg.get('output_name')
    if isinstance(out, str):
        return out
    suf = cfg.get('suffix', '')
    return name + (suf if isinstance(suf, str) else '')


# Emit early to avoid Pydantic forward-ref issues (Schema/ConfigSchema for OutputConfig; Metadata for MessageData etc.)
PREFERRED_FIRST = ('Schema', 'ConfigSchema', 'Metadata', 'Custom')
# anyOf/oneOf defs emitted as RootModel (have .root) so Part(root=TextPart(...)) works
ROOT_MODEL_UNIONS = frozenset({'Part', 'DocumentPart', 'TraceEvent'})
HEADER = '''# Copyright {year} 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
#
# DO NOT EDIT: Generated by `generate_schema_typing` from `{schema_name}`.

"""Schema types module defining the core data models for Genkit."""

from __future__ import annotations

import warnings
from typing import Any, ClassVar, Literal

from pydantic import ConfigDict, Field, RootModel
from pydantic.alias_generators import to_camel

from genkit._core._base import GenkitModel
from genkit._core._compat import StrEnum

warnings.filterwarnings('ignore', message='Field name "schema" in "OutputConfig" shadows an attribute in parent', category=UserWarning)

'''


def _camel_to_snake(s: str) -> str:
    return re.sub(r'(?<!^)(?=[A-Z])', '_', s).lower()


def _pascal(s: str) -> str:
    return s[0].upper() + s[1:] if s else 'X'


def _resolve_ref(schema: dict, ref: str) -> dict:
    """Resolve $ref (e.g. #/$defs/X or #/$defs/X/properties/Y) to schema fragment."""
    if not ref.startswith('#/$defs/'):
        return {}
    obj = schema.get('$defs', {})
    for part in ref[8:].split('/'):
        obj = obj.get(part, {}) if isinstance(obj, dict) else {}
    return obj if isinstance(obj, dict) else {}


def _models_allowing_extra(schema: dict) -> set[str]:
    """Names of models with additionalProperties: true (extra='allow')."""
    result = set()
    defs = schema.get('$defs') or {}
    for name, defn in defs.items():
        if isinstance(defn, dict) and defn.get('additionalProperties') is True:
            result.add(name)
        if not isinstance(defn, dict):
            continue
        for prop_name, prop_def in (defn.get('properties') or {}).items():
            if (
                isinstance(prop_def, dict)
                and prop_def.get('type') == 'object'
                and prop_def.get('additionalProperties') is True
            ):
                result.add(_pascal(prop_name))
    return result


def _typed_map_aliases(defs: dict) -> dict[str, str]:
    """Inline object schemas with typed scalar ``additionalProperties`` -> Python dict alias.

    e.g. ``ReflectionRunActionParams.telemetryLabels``:
      ``{type: object, additionalProperties: {type: string}}`` -> ``dict[str, str]``.

    Emitting these as type aliases (mirroring ``Metadata`` / ``Custom``) keeps the
    symbol exported and importable while letting callers pass plain Python dicts —
    a class with no fields and ``extra='forbid'`` would reject every key on the
    Dev UI's ``{'genkitx:ignore-trace': 'true'}`` payload.
    """

    result: dict[str, str] = {}
    for name, defn in defs.items():
        if not isinstance(defn, dict) or defn.get('type') != 'object':
            continue
        ap = defn.get('additionalProperties')
        if not isinstance(ap, dict):
            continue
        ap_type = ap.get('type')
        if isinstance(ap_type, str) and ap_type in PRIM:
            result[name] = f'dict[str, {PRIM[ap_type]}]'
    return result


def _extract_inline_classes(schema: dict) -> dict[str, dict]:
    """Extract inline object schemas to named classes (e.g. Score.details -> Details)."""
    result = {}
    defs = schema.get('$defs') or {}

    def walk(props: dict) -> None:
        for prop_name, prop_schema in (props or {}).items():
            if isinstance(prop_schema, dict) and prop_schema.get('type') == 'object' and '$ref' not in prop_schema:
                class_name = _pascal(prop_name)
                if class_name not in defs and class_name not in result:
                    result[class_name] = prop_schema
                walk(prop_schema.get('properties', {}))

    for defn in defs.values():
        if isinstance(defn, dict):
            walk(defn.get('properties', {}))
    return result


def _py_type(prop: dict, schema: dict, defs: dict, class_name: str, field_name: str) -> str:
    """Resolve JSON schema property to Python type string (e.g. list[DocumentData], Role | str)."""
    if '$ref' in prop:
        ref = prop['$ref']
        path = ref[8:].split('/') if ref.startswith('#/$defs/') else []
        target = _resolve_ref(schema, ref)
        ref_name = path[-1] if path else ''
        # Top-level def (#/$defs/X) -> use class name
        if len(path) == 1:
            return _output_name(ref_name)
        # Nested ref (#/$defs/X/properties/Y) -> resolve target; empty schema -> Any
        if not target or (not target.get('type') and not target.get('properties') and 'enum' not in target):
            return 'Any'
        if 'enum' in target:
            vals = target.get('enum', [])
            if field_name in ('tool_choice', 'toolChoice') and set(vals) == {'auto', 'required', 'none'}:
                return 'ToolChoice'
            if field_name in ('constrained',) and set(vals) == {'none', 'all', 'no-tools'}:
                return 'Constrained'
            if field_name in ('stage',) and set(vals) == {'featured', 'stable', 'unstable', 'legacy', 'deprecated'}:
                return 'Stage'
            return _output_name(ref_name)
        if target.get('type') == 'array':
            inner = _py_type(target.get('items', {}), schema, defs, class_name, field_name) or 'Any'
            return f'list[{inner}]'
        if target.get('type') == 'object':
            # Flexible dict-like fields: Metadata and Custom (SDK uses .get(), [], etc.)
            if ref_name == 'metadata' and 'additionalProperties' in target:
                return 'Metadata'
            if ref_name == 'custom' and 'additionalProperties' in target:
                return 'Custom'
            return 'dict[str, Any]'
        if target.get('type'):
            t = target['type']
            return PRIM.get(t, 'Any') if isinstance(t, str) else 'Any'
        return _output_name(ref_name)
    # anyOf / oneOf -> Union of refs or resolved types
    for key in ('anyOf', 'oneOf'):
        if key in prop:
            opts = prop[key]
            refs = [o.get('$ref', '').split('/')[-1] for o in opts if o.get('$ref')]
            if refs:
                return ' | '.join(_output_name(r) for r in refs)
            types = sorted({_py_type(o, schema, defs, class_name, field_name) for o in opts} - {''})
            return ' | '.join(types) if types else 'Any'
    if prop.get('type') == 'array':
        return f'list[{_py_type(prop.get("items", {}), schema, defs, class_name, field_name) or "Any"}]'
    if prop.get('type') == 'object':
        # Flexible custom field (additionalProperties) -> use Custom (type alias for dict)
        if field_name in ('custom',) and prop.get('additionalProperties') is not None:
            return 'Custom'
        if _pascal(field_name) in defs:
            return _pascal(field_name)
        return 'dict[str, Any]'
    if 'enum' in prop:
        vals = prop['enum']
        if field_name in ('tool_choice', 'toolChoice') and set(vals) == {'auto', 'required', 'none'}:
            return 'ToolChoice'
        if field_name in ('constrained',) and set(vals) == {'none', 'all', 'no-tools'}:
            return 'Constrained'
        if field_name in ('stage',) and set(vals) == {'featured', 'stable', 'unstable', 'legacy', 'deprecated'}:
            return 'Stage'
        return 'Literal[' + ', '.join(repr(v) for v in vals) + ']'
    t = prop.get('type')
    if isinstance(t, list):
        return ' | '.join(sorted(set(PRIM.get(x, 'Any') for x in t))) or 'Any'
    return PRIM.get(t, 'Any') if t else 'Any'


def _emit_enum(name: str, d: dict) -> list[str]:
    lines = [f'class {name}(StrEnum):', f'    """{name} data type class."""', '']
    for v in d.get('enum', []):
        m = str(v).upper().replace('-', '_')
        if m and m[0].isdigit():
            m = '_' + m
        lines.append(f'    {m} = {repr(v)}')
    return lines + ['']


def _emit_model(
    name: str, d: dict, schema: dict, defs: dict, allow: set[str], omit: set[str] | None = None
) -> list[str]:
    props, req = d.get('properties', {}), set(d.get('required', []))
    if omit:
        props = {k: v for k, v in props.items() if k not in omit and _camel_to_snake(k) not in omit}
        req = req - omit - {_camel_to_snake(k) for k in omit}
    ext = ', protected_namespaces=()' if any(_camel_to_snake(k) in ('schema', 'schema_') for k in props) else ''
    frz = ', frozen=True' if name == 'PathMetadata' else ''
    cfg = f"ConfigDict(alias_generator=to_camel, extra='{'allow' if name in allow else 'forbid'}', populate_by_name=True{ext}{frz})"
    lines = [
        f'class {name}(GenkitModel):',
        f'    """Model for {name.lower().replace("_", " ")} data."""',
        f'    model_config: ClassVar[ConfigDict] = {cfg}',
    ]
    for k, v in props.items():
        # Use schema_ for OutputConfig.schema to avoid shadowing GenkitModel.schema
        snake = _camel_to_snake(k)
        if name == 'OutputConfig' and snake == 'schema':
            field_name = 'schema_'
            alias_extra = ", alias='schema'"
        elif snake in ('schema_', 'schema'):
            field_name = 'schema' if name != 'OutputConfig' else 'schema_'
            alias_extra = ", alias='schema'" if name == 'OutputConfig' else ''
        else:
            field_name = snake
            alias_extra = ''
        py_type_str = _py_type(v, schema, defs, name, k)
        # OutputConfig.schema is free-form JSON schema object; use dict for direct use
        if name == 'OutputConfig' and snake == 'schema':
            py_type_str = 'dict[str, Any]'
        if name == 'MessageData' and k == 'role':
            py_type_str = 'Role | str'
        # OutputConfig.schema is free-form JSON schema dict (not Schema model)
        if name == 'OutputConfig' and snake == 'schema':
            py_type_str = 'dict[str, Any]'
        desc = v.get('description')
        desc_extra = f', description={repr(desc)}' if desc else ''
        if k in req:
            lines.append(f'    {field_name}: {py_type_str} = Field(...{desc_extra}{alias_extra})')
        else:
            default_val = (
                f'Field(default=None{desc_extra}{alias_extra})'
                if '|' in py_type_str or py_type_str == 'Any'
                else 'None'
            )
            lines.append(f'    {field_name}: {py_type_str} | None = {default_val}')
    if name == 'GenerateActionOutputConfig':
        lines.extend([
            '    # Store Pydantic type for runtime validation (excluded from JSON)',
            '    schema_type: Any = Field(default=None, exclude=True)',
        ])
    return lines + ['']


def generate(schema_path: Path, _out: Path) -> str:
    schema = json.loads(schema_path.read_text())
    defs = dict(schema.get('$defs', {}))
    defs.update({k: v for k, v in _extract_inline_classes(schema).items() if k not in defs})
    allow_extra = _models_allowing_extra(schema)
    typed_map_aliases = _typed_map_aliases(defs)
    out = [HEADER.format(year=datetime.now().year, schema_name=schema_path.name)]
    emitted = set()

    # Pass 1: enums
    for name, defn in defs.items():
        if name in EXCLUDED or name in emitted or not isinstance(defn, dict):
            continue
        if 'enum' in defn:
            class_name = _output_name(name)
            out.extend(_emit_enum(class_name, defn))
            emitted.add(name)

    # Pass 2: object models (must precede root models that reference them)
    # Emit Schema, ConfigSchema, Metadata early (OutputConfig uses Schema; MessageData etc. use Metadata)
    for name in (*PREFERRED_FIRST, *(k for k in defs if k not in PREFERRED_FIRST)):
        defn = defs.get(name, {})
        if name in EXCLUDED or name in emitted or not isinstance(defn, dict) or defn.get('type') != 'object':
            continue
        class_name = _output_name(name)
        # Metadata and Custom: type aliases for dict (SDK uses .get(), [], passes dict)
        if name == 'Metadata':
            out.extend([
                'Metadata = dict[str, Any]  # type alias for flexible metadata',
                '',
            ])
        elif name == 'Custom':
            out.extend([
                'Custom = dict[str, Any]  # type alias for flexible custom data',
                '',
            ])
        elif name in typed_map_aliases:
            # Typed string-keyed maps (e.g. TelemetryLabels: dict[str, str]). Emitting as a
            # type alias keeps the symbol exported and lets callers pass plain dicts.
            out.extend([
                f'{class_name} = {typed_map_aliases[name]}  # type alias for {name.lower()} (typed string map)',
                '',
            ])
        elif name in TRANSFORMATIONS and (cfg := TRANSFORMATIONS[name]).get('omit'):
            omit_set = set(cfg.get('omit', []))
            out.extend(_emit_model(class_name, defn, schema, defs, allow_extra, omit=omit_set))
            emitted.add(name)
        else:
            out.extend(_emit_model(class_name, defn, schema, defs, allow_extra))
        emitted.add(name)

    # Pass 2.5: union types (anyOf/oneOf)
    # Part and DocumentPart need RootModel so Part(root=TextPart(...)) works; others get type aliases
    ROOT_MODEL_UNIONS = frozenset({'Part', 'DocumentPart'})
    for name, defn in defs.items():
        if name in EXCLUDED or name in emitted or not isinstance(defn, dict):
            continue
        for key in ('anyOf', 'oneOf'):
            if key not in defn:
                continue
            opts = defn[key]
            refs = [(o.get('$ref') or '').split('/')[-1] for o in opts if isinstance(o, dict) and o.get('$ref')]
            if not refs:
                continue
            class_name = _output_name(name)
            union_str = ' | '.join(_output_name(r) for r in refs)
            if name in ROOT_MODEL_UNIONS:
                out.extend([
                    f'class {class_name}(RootModel[{union_str}]):',
                    f'    """Root model for {name} union (Part(root=X), DocumentPart(root=X))."""',
                    '',
                ])
            else:
                out.extend([f'{class_name} = {union_str}', ''])
            emitted.add(name)
            break

    # Pass 3: root models (array at top level, e.g. EvalResponse)
    for name, defn in defs.items():
        if name in EXCLUDED or name in emitted or not isinstance(defn, dict):
            continue
        if defn.get('type') != 'array' or 'enum' in defn:
            continue
        class_name = _output_name(name)
        items_schema = defn.get('items', {})
        ref_name = (items_schema.get('$ref') or '').split('/')[-1]
        inner_type = _output_name(ref_name) if ref_name else 'Any'
        out.extend([
            f'class {class_name}(RootModel[list[{inner_type}]]):',
            f'    """Root model for {name.lower()}."""',
            f'    root: list[{inner_type}]',
            '',
        ])
        emitted.add(name)

    # Supplemental types (Python SDK; from inline schema or SDK-specific)
    out.extend([
        '',
        'class Constrained(StrEnum):',
        '    """Constrained generation support (none, all, no-tools)."""',
        '',
        '    NONE = "none"',
        '    ALL = "all"',
        '    NO_TOOLS = "no-tools"',
        '',
        'class Stage(StrEnum):',
        '    """Model stage (featured, stable, unstable, legacy, deprecated)."""',
        '',
        '    FEATURED = "featured"',
        '    STABLE = "stable"',
        '    UNSTABLE = "unstable"',
        '    LEGACY = "legacy"',
        '    DEPRECATED = "deprecated"',
        '',
        'class ToolChoice(StrEnum):',
        '    """Tool choice for generation (auto, required, none)."""',
        '',
        '    AUTO = "auto"',
        '    REQUIRED = "required"',
        '    NONE = "none"',
        '',
        'class MediaModel(RootModel[Any]):',
        '    """Wrapper for media content (flexible structure)."""',
        '',
        'class Text(RootModel[str]):',
        '    """Plain text content."""',
        '',
        'Resource1 = Resource  # alias for Resource (resource with uri)',
        '',
    ])

    content = '\n'.join(out)
    content = re.sub(r'\bMessage\b(?!Data)', 'MessageData', content)
    return content + '\n'


def main() -> None:
    # From py/tools/schema_to_typing/schema_to_typing.py -> repo root is parent.parent.parent.parent
    top = Path(__file__).resolve().parent.parent.parent.parent
    schema = top / 'genkit-tools' / 'genkit-schema.json'
    out = top / 'py' / 'packages' / 'genkit' / 'genkit' / '_core' / '_typing.py'
    if len(sys.argv) >= 2:
        schema = Path(sys.argv[1]).resolve()
        out = Path(sys.argv[2]).resolve() if len(sys.argv) > 2 else schema.parent / '_typing.py'
    if not schema.is_file():
        sys.exit(1)
    out.write_text(generate(schema, out), encoding='utf-8')


if __name__ == '__main__':
    main()
