"""Stage env- and trainer-Space directories from the repo root.

Each Space needs a *single* directory containing the full repo plus the
right Dockerfile + README front-matter at its root. This script copies
the repo into a staging directory, drops in the Space-specific
``Dockerfile`` / ``README.md``, and prints the staging path.
"""

from __future__ import annotations

import argparse
import shutil
import sys
from pathlib import Path


REPO_ROOT = Path(__file__).resolve().parent.parent

EXCLUDES = {
    ".venv",
    "__pycache__",
    ".git",
    ".cursor",
    ".DS_Store",
    "runs",
    "wandb",
    "node_modules",
    ".pytest_cache",
    ".mypy_cache",
}


def _ignore(_dir: str, names):
    return [n for n in names if n in EXCLUDES or n.endswith((".pyc", ".log"))]


def _stage(stage_dir: Path) -> None:
    if stage_dir.exists():
        shutil.rmtree(stage_dir)
    shutil.copytree(REPO_ROOT, stage_dir, ignore=_ignore, symlinks=False)


def build_env_space(stage_dir: Path) -> None:
    _stage(stage_dir)

    dockerfile = """\
# CERNenv environment Space (Docker, CPU)
FROM python:3.11-slim

ENV PYTHONUNBUFFERED=1 \\
    PIP_NO_CACHE_DIR=1 \\
    PYTHONPATH=/home/user/app

RUN apt-get update && apt-get install -y --no-install-recommends \\
        git curl ca-certificates build-essential \\
    && rm -rf /var/lib/apt/lists/*

RUN useradd -ms /bin/bash user
USER user
WORKDIR /home/user/app

COPY --chown=user:user space/env/requirements.txt /tmp/requirements.txt
RUN python -m pip install --upgrade pip && \\
    python -m pip install --user -r /tmp/requirements.txt

COPY --chown=user:user . /home/user/app

EXPOSE 7860

CMD [\"python\", \"-m\", \"uvicorn\", \"server.app:app\", \"--host\", \"0.0.0.0\", \"--port\", \"7860\"]
"""
    (stage_dir / "Dockerfile").write_text(dockerfile)

    readme = (stage_dir / "space" / "env" / "README.md").read_text()
    (stage_dir / "README.md").write_text(readme)


def build_trainer_space(stage_dir: Path) -> None:
    _stage(stage_dir)

    dockerfile = """\
# CERNenv trainer Space (Docker, A100)
FROM nvidia/cuda:12.1.1-cudnn8-devel-ubuntu22.04

ENV DEBIAN_FRONTEND=noninteractive \\
    PYTHONUNBUFFERED=1 \\
    PIP_NO_CACHE_DIR=1 \\
    HF_HOME=/home/user/.cache/huggingface \\
    TRANSFORMERS_CACHE=/home/user/.cache/huggingface/transformers \\
    PYTHONPATH=/home/user/app

RUN apt-get update && apt-get install -y --no-install-recommends \\
        python3.11 python3.11-venv python3.11-dev python3-pip \\
        git curl ca-certificates build-essential \\
    && rm -rf /var/lib/apt/lists/* \\
    && ln -sf /usr/bin/python3.11 /usr/local/bin/python \\
    && ln -sf /usr/bin/python3.11 /usr/local/bin/python3

RUN useradd -ms /bin/bash user
USER user
ENV PATH=\"/home/user/.local/bin:${PATH}\"
WORKDIR /home/user/app

COPY --chown=user:user space/training/requirements.txt /tmp/requirements.txt
RUN python -m pip install --upgrade pip && \\
    python -m pip install --user -r /tmp/requirements.txt

COPY --chown=user:user . /home/user/app

EXPOSE 7860

CMD [\"python\", \"-m\", \"uvicorn\", \"space.training.app:app\", \"--host\", \"0.0.0.0\", \"--port\", \"7860\"]
"""
    (stage_dir / "Dockerfile").write_text(dockerfile)

    readme = (stage_dir / "space" / "training" / "README.md").read_text()
    (stage_dir / "README.md").write_text(readme)


def main() -> None:  # pragma: no cover
    parser = argparse.ArgumentParser()
    parser.add_argument("kind", choices=["env", "trainer"])
    parser.add_argument("--stage_dir", required=True)
    args = parser.parse_args()

    stage_dir = Path(args.stage_dir).resolve()
    if args.kind == "env":
        build_env_space(stage_dir)
    else:
        build_trainer_space(stage_dir)
    print(stage_dir)


if __name__ == "__main__":  # pragma: no cover
    main()