# Copyright The Lightning AI team.
#
# 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.

"""Generate the list of supported public models for type hints.

This script reads the authoritative `PUBLIC_MODELS` mapping from
`lightning_sdk.llm.public_assistants` and generates the file:

    src/litai/utils/supported_public_models.py

That file contains:
    - `MODELS`: a tuple of all supported public model names.
    - `ModelLiteral`: a `typing.Literal` of the same names for IDE autocompletion
      and static type checking.

The generated file includes licensing information, a clear auto-generated warning,
and is formatted to comply with typical linting line-length limits.

Usage:
    python src/litai/utils/gen_supported_models.py

This script is intended for internal development and release processes
and is usually run via a pre-commit hook or during the build step before
publishing to PyPI.
"""

import re
from pathlib import Path

from lightning_sdk.llm.public_assistants import PUBLIC_MODELS


def main() -> int:
    """Generate the list of supported public models for type hints."""
    path = Path("src/litai/utils/supported_public_models.py")

    new_models = sorted(PUBLIC_MODELS.keys())

    existing_models = []
    if path.exists():
        content = path.read_text(encoding="utf-8")
        match = re.search(r"MODELS: Final = \(([\s\S]*?)\)", content)
        if match:
            tuple_content = match.group(1)
            # Use eval to safely parse the tuple content
            try:
                # Add brackets to make it a list for a consistent eval
                existing_models = eval(f"[{tuple_content.strip()}]")
            except (SyntaxError, NameError):
                print("Warning: Could not parse models from existing file.")

    if new_models == existing_models:
        print("supported_public_models.py is up to date.")
        return 0

    model_literals = ",\n    ".join(repr(m) for m in new_models)

    code = f"""# AUTO-GENERATED. DO NOT EDIT.
# Generated by litai/utils/gen_supported_models.py
# Copyright The Lightning AI team.
#
# 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.

\"\"\"Supported public model names and their type literal.

This file is auto-generated by `litai/utils/gen_supported_models.py`
from the authoritative `PUBLIC_MODELS` in `lightning_sdk.llm.public_assistants`.

It contains:
- `MODELS`: a tuple of all supported public model names.
- `ModelLiteral`: a `typing.Literal` of the same names, for autocompletion and type safety.

Do not edit this file by hand. Regenerate it instead.
\"\"\"

from typing import Final, Literal

MODELS: Final = (
    {model_literals},
)
ModelLiteral = Literal[
    {model_literals}
]
"""
    with open(path, "w", encoding="utf-8", newline="\n") as f:
        f.write(code + "\n")
    print("Wrote src/litai/utils/supported_public_models.py")
    return 1  # Return a non-zero exit code to signal a change was made.


if __name__ == "__main__":
    raise SystemExit(main())
