from pathlib import Path

import torch

from .archs import build_mlp
from .templates import (
    render_config,
    render_data,
    render_init,
    render_model,
    render_main,
    render_predict,
    render_requirements,
    render_train,
)


def detect_device():
    if torch.backends.mps.is_available():
        return "mps"
    if torch.cuda.is_available():
        return "cuda"
    return "cpu"


def generate(
    name,
    arch="mlp",
    input_dim=784,
    output_dim=10,
    hidden=None,
    activation="relu",
    metrics=None,
    prod=False,
    lr=0.001,
    epochs=10,
    batch_size=32,
):
    if hidden is None:
        hidden = [256, 128]

    device = detect_device()
    has_amp = device == "cuda"

    if arch == "mlp":
        layers_code = build_mlp(input_dim, hidden, output_dim, activation)
    else:
        raise ValueError(f"Unsupported arch: {arch}. Supported: 'mlp'")

    src_dir = Path("src")
    src_dir.mkdir(exist_ok=True)

    files = {
        src_dir / "__init__.py": render_init(name),
        src_dir / "config.py": render_config(
            name, input_dim, hidden, output_dim, lr, epochs, batch_size, device
        ),
        src_dir / "models.py": render_model(name, layers_code),
        src_dir / "data.py": render_data(name),
        src_dir / "train.py": render_train(name, device, prod, metrics, has_amp),
        src_dir / "predict.py": render_predict(name),
        Path("requirements.txt"): render_requirements(prod, metrics),
        Path("main.py"): render_main(),
    }

    for path, content in files.items():
        path.write_text(content)
        print(f"  {path}")

    print(f"\nGenerated {name} ({arch}) on {device}")
    if prod:
        print("  View metrics: mlflow ui")
    else:
        print("  View metrics: tensorboard --logdir=runs")
