"""Command line utilities for managing OpenReward environments."""
from __future__ import annotations

import argparse
import re
from dataclasses import dataclass, field
from importlib import resources
from pathlib import Path
from string import Template
from typing import Callable, Mapping


@dataclass(frozen=True)
class TemplateFiles:
    """Collection of scaffold files for an environment."""

    dockerfile: str
    server_py: str
    extra_files: Mapping[str, str] = field(default_factory=dict)


def _pascal_case(name: str) -> str:
    parts = re.split(r"[^0-9a-zA-Z]+", name)
    cleaned = "".join(part.capitalize() for part in parts if part)
    if not cleaned:
        cleaned = "Environment"
    if cleaned[0].isdigit():
        cleaned = f"Env{cleaned}"
    return cleaned


def _load_template_file(template_name: str, filename: str) -> str:
    package = f"openreward.templates.{template_name}"
    file = resources.files(package).joinpath(filename)
    return file.read_text(encoding="utf-8")


_README_TEMPLATE = """\
# {env_name}

[![OpenReward Environment](https://img.shields.io/badge/%E2%AD%90%20OpenReward-Environment-f7e6cc)](https://openreward.ai/YourOrg/{env_name})

## Description

<!-- 2-4 sentences: what does this environment evaluate? What domain does it cover?
     Mention the verification method and the source dataset. -->

## Capabilities

<!-- Bulleted list of skills or capabilities the environment tests. -->

- Capability 1
- Capability 2
- Capability 3

## Compute Requirements

<!-- Does the environment need a sandbox? GPU? Extra memory?
     If it runs without a sandbox, say so. -->

## License

<!-- License that covers the use of the environment, as well as any underlying data or code it depends on.
     Link to the license text. If the environment and its dependencies have different licenses, list both. -->

## Tasks

<!-- How many tasks are there? What splits (train/test)?
     Briefly describe the structure of each task. -->

## Reward Structure

<!-- How are rewards computed?
     Sparse or dense? Binary or continuous?
     Is verification programmatic or does it use an LLM grader? Does it use rubrics?
     If there are multiple validation steps, list them in order. -->

## Data

<!-- Where does the task data come from? How is it stored?
     Link to the source dataset if applicable. -->

## Tools

<!-- Explain the tools available to the agent. -->

## Time Horizon

<!-- Single-turn or multi-turn?
     If multi-turn, roughly how many tool calls does a typical task require? -->

## Environment Difficulty

<!-- Solve rates, baseline model performance, or other difficulty statistics. -->

## Other Environment Requirements

<!-- Any external dependencies: API keys, secrets, third-party services. -->

## Safety

<!-- Does the agent have access to external systems (network, file system, APIs)?
     Are there dual-use risks in the domain (e.g. chemistry, cybersecurity)?
     Is there a possibility of goal misspecification?
     What mitigations are in place? -->

## Citations

<!-- BibTeX entries for the environment itself and any underlying datasets or papers. -->

```bibtex
@dataset{{YourCitation,
  author    = {{Your Name or Team}},
  title     = {{{env_name}}},
  year      = {{2026}},
  publisher = {{OpenReward}},
  url       = {{https://openreward.ai/YourOrg/{env_name}}}
}}
```
"""


def _render_readme(env_name: str) -> str:
    return _README_TEMPLATE.format(env_name=env_name)


def _basic_template(env_name: str) -> TemplateFiles:
    class_name = _pascal_case(env_name)
    dockerfile = _load_template_file("basic", "Dockerfile")
    server_template = Template(_load_template_file("basic", "server.py.tmpl"))
    server_py = server_template.substitute(CLASS_NAME=class_name)
    extra_files = {
        "requirements.txt": _load_template_file("basic", "requirements.txt"),
        "README.md": _render_readme(env_name),
    }
    return TemplateFiles(dockerfile=dockerfile, server_py=server_py, extra_files=extra_files)


def _sandbox_template(env_name: str) -> TemplateFiles:
    """Create a sandbox template with Docker-in-Docker support."""
    dockerfile = _load_template_file("sandbox", "Dockerfile")
    server_py = _load_template_file("sandbox", "server.py")
    extra_files = {
        "requirements.txt": _load_template_file("sandbox", "requirements.txt"),
        "sandbox_env.py": _load_template_file("sandbox", "sandbox_env.py"),
        "README.md": _render_readme(env_name),
    }
    return TemplateFiles(dockerfile=dockerfile, server_py=server_py, extra_files=extra_files)


TemplateFactory = Callable[[str], TemplateFiles]


TEMPLATES: Mapping[str, TemplateFactory] = {
    "basic": _basic_template,
    "sandbox": _sandbox_template,
}


def _render_template(template_name: str, env_name: str) -> TemplateFiles:
    generator = TEMPLATES.get(template_name)
    if generator is None:
        raise ValueError(f"Unknown template '{template_name}'. Available templates: {', '.join(TEMPLATES)}")
    return generator(env_name)


def _write_file(path: Path, contents: str) -> None:
    if path.exists():
        raise FileExistsError(f"{path} already exists")
    path.write_text(contents, encoding="utf-8")


def command_init(environment: str, template: str) -> None:
    target_dir = Path.cwd() / environment
    target_dir.mkdir(parents=True, exist_ok=False)

    files = _render_template(template, environment)

    _write_file(target_dir / "Dockerfile", files.dockerfile)
    _write_file(target_dir / "server.py", files.server_py)
    for relative, contents in files.extra_files.items():
        _write_file(target_dir / relative, contents)


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(prog="orwd", description="OpenReward CLI utilities")
    subparsers = parser.add_subparsers(dest="command", required=True)

    init_parser = subparsers.add_parser(
        "init",
        help="Initialize a new environment template",
    )
    init_parser.add_argument("environment", help="Name of the environment to scaffold")
    init_parser.add_argument(
        "--template",
        default="basic",
        choices=sorted(TEMPLATES.keys()),
        help="Scaffold template to use (default: %(default)s)",
    )
    init_parser.set_defaults(func=lambda args: command_init(args.environment, args.template))

    return parser


def main(argv: list[str] | None = None) -> int:
    parser = build_parser()
    args = parser.parse_args(argv)
    try:
        args.func(args)
    except FileExistsError as exc:
        parser.error(str(exc))
    return 0


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