"""
Project initialization command for TorchBridge CLI.

Scaffolds a backend-agnostic PyTorch project from templates.
"""

import argparse
import sys
from pathlib import Path

# ---------------------------------------------------------------------------
# Template content as embedded string literals
# ---------------------------------------------------------------------------

_GITIGNORE = """\
# Python
__pycache__/
*.py[cod]
*$py.class
*.so
*.egg-info/
dist/
build/
*.egg

# Virtual environments
.venv/
venv/
env/

# IDE
.vscode/
.idea/
*.swp
*.swo

# Model artifacts
*.pt
*.pth
*.onnx
*.safetensors
checkpoints/

# Logs and reports
*.log
reports/
benchmark_results/
tensorboard_runs/

# Environment
.env
.env.local
"""

_REQUIREMENTS_BASE = """\
torch>=2.0.0
torchbridge>=0.5.0
numpy>=1.21.0
"""

_REQUIREMENTS_DISTRIBUTED = """\
torch>=2.0.0
torchbridge>=0.5.0
numpy>=1.21.0
# For distributed training with NCCL backend, ensure NCCL is installed:
# pip install torch with CUDA support, NCCL is bundled.
"""

_REQUIREMENTS_SERVING = """\
torch>=2.0.0
torchbridge>=0.5.0
numpy>=1.21.0
fastapi>=0.103.0
uvicorn[standard]>=0.23.0
"""


def _config_for_training(name: str, backend: str) -> str:
    return f"""\
# {name} - Training Configuration
# Generated by tb-init

project:
  name: "{name}"
  type: training

hardware:
  backend: {backend}  # auto-detected at runtime if 'auto'

training:
  batch_size: 32
  learning_rate: 0.001
  epochs: 10
  optimizer: adamw
  scheduler: cosine

model:
  hidden_size: 768
  num_layers: 6

export:
  formats:
    - torchscript
    - safetensors
"""


def _config_for_inference(name: str, backend: str) -> str:
    return f"""\
# {name} - Inference Configuration
# Generated by tb-init

project:
  name: "{name}"
  type: inference

hardware:
  backend: {backend}  # auto-detected at runtime if 'auto'

inference:
  batch_size: 1
  precision: fp32
  optimization_level: production

export:
  formats:
    - torchscript
    - onnx
"""


def _config_for_distributed(name: str, backend: str) -> str:
    return f"""\
# {name} - Distributed Training Configuration
# Generated by tb-init

project:
  name: "{name}"
  type: distributed

hardware:
  backend: {backend}  # auto-detected at runtime if 'auto'

training:
  batch_size: 32
  learning_rate: 0.001
  epochs: 10
  optimizer: adamw
  scheduler: cosine

distributed:
  strategy: ddp
  num_gpus: auto
  backend: nccl

model:
  hidden_size: 768
  num_layers: 6
"""


def _train_py(name: str) -> str:
    return f'''\
"""
{name} — Training Script
Generated by tb-init. Uses TorchBridge for backend-agnostic training.
"""

import torch
import torch.nn as nn
from torch.utils.data import DataLoader, TensorDataset

from torchbridge import TorchBridgeConfig, get_config
from torchbridge.hardware.abstraction.hal_core import HardwareAbstractionLayer


def detect_best_backend():
    """Detect the best available backend."""
    hal = HardwareAbstractionLayer()
    return hal


def build_model(hidden_size: int = 768, num_layers: int = 6):
    """Build a simple model for demonstration."""
    layers = []
    for _ in range(num_layers):
        layers.extend([
            nn.Linear(hidden_size, hidden_size),
            nn.GELU(),
        ])
    layers.append(nn.Linear(hidden_size, 10))
    return nn.Sequential(*layers)


def main():
    # Configuration
    config = TorchBridgeConfig.for_training()

    # Hardware detection
    hal = detect_best_backend()
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"Using device: {{device}}")

    # Model
    model = build_model().to(device)
    print(f"Model parameters: {{sum(p.numel() for p in model.parameters()):,}}")

    # Dummy dataset for demonstration
    X = torch.randn(1000, 768)
    y = torch.randint(0, 10, (1000,))
    dataset = TensorDataset(X, y)
    loader = DataLoader(dataset, batch_size=32, shuffle=True)

    # Training
    optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)
    criterion = nn.CrossEntropyLoss()

    for epoch in range(10):
        model.train()
        total_loss = 0
        for batch_X, batch_y in loader:
            batch_X, batch_y = batch_X.to(device), batch_y.to(device)
            optimizer.zero_grad()
            output = model(batch_X)
            loss = criterion(output, batch_y)
            loss.backward()
            optimizer.step()
            total_loss += loss.item()
        print(f"Epoch {{epoch+1}}/10  Loss: {{total_loss / len(loader):.4f}}")

    # Export
    model.eval()
    sample_input = torch.randn(1, 768, device=device)
    traced = torch.jit.trace(model, sample_input)
    traced.save(f"{name}_model.pt")
    print(f"Model saved to {name}_model.pt")


if __name__ == "__main__":
    main()
'''


def _serve_py(name: str) -> str:
    return f'''\
"""
{name} — Inference / Serving Script
Generated by tb-init. Uses TorchBridge for backend-agnostic inference.
"""

import torch
import torch.nn as nn

from torchbridge import TorchBridgeConfig
from torchbridge.hardware.abstraction.hal_core import HardwareAbstractionLayer


def detect_best_backend():
    """Detect the best available backend."""
    hal = HardwareAbstractionLayer()
    return hal


def load_model(model_path: str, device: torch.device) -> nn.Module:
    """Load a TorchScript model."""
    model = torch.jit.load(model_path, map_location=device)
    model.eval()
    return model


def predict(model: nn.Module, input_tensor: torch.Tensor) -> torch.Tensor:
    """Run inference on input."""
    with torch.no_grad():
        return model(input_tensor)


def main():
    config = TorchBridgeConfig.for_inference()

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"Using device: {{device}}")

    # For demonstration, create a simple model
    model = nn.Sequential(
        nn.Linear(768, 768),
        nn.GELU(),
        nn.Linear(768, 10),
    ).to(device).eval()

    # Run inference
    sample = torch.randn(1, 768, device=device)
    output = predict(model, sample)
    print(f"Output shape: {{output.shape}}")
    print(f"Predictions: {{output.argmax(dim=-1).tolist()}}")


if __name__ == "__main__":
    main()
'''


def _dockerfile_training(name: str) -> str:
    return f"""\
# {name} — Training Dockerfile
# Generated by tb-init

FROM python:3.12-slim

WORKDIR /app

COPY requirements.txt .
RUN pip install --no-cache-dir -r requirements.txt

COPY . .

CMD ["python", "train.py"]
"""


def _dockerfile_inference(name: str) -> str:
    return f"""\
# {name} — Inference Dockerfile
# Generated by tb-init

FROM python:3.12-slim

WORKDIR /app

COPY requirements.txt .
RUN pip install --no-cache-dir -r requirements.txt

COPY . .

CMD ["python", "serve.py"]
"""


def _dockerfile_distributed(name: str) -> str:
    return f"""\
# {name} — Distributed Training Dockerfile
# Generated by tb-init
# NOTE: For multi-GPU training, use nvidia/cuda base image and install NCCL.

FROM python:3.12-slim

WORKDIR /app

COPY requirements.txt .
RUN pip install --no-cache-dir -r requirements.txt

COPY . .

CMD ["python", "train.py"]
"""


def _dockerfile_serving(name: str) -> str:
    return f"""\
# {name} — Serving Dockerfile
# Generated by tb-init

FROM python:3.12-slim

WORKDIR /app

COPY requirements.txt .
RUN pip install --no-cache-dir -r requirements.txt

COPY . .

EXPOSE 8000

CMD ["uvicorn", "serve:app", "--host", "0.0.0.0", "--port", "8000"]
"""


def _readme(name: str, template: str) -> str:
    return f"""\
# {name}

A backend-agnostic PyTorch project scaffolded with [TorchBridge](https://github.com/CloudlyIO/torchbridge).

## Quickstart

```bash
# Install dependencies
pip install -r requirements.txt

# Run
{"python train.py" if template in ("training", "distributed") else "python serve.py"}
```

## Docker

```bash
docker build -t {name} .
docker run {name}
```

## TorchBridge

This project uses TorchBridge for hardware-agnostic PyTorch development.
TorchBridge automatically detects your hardware (NVIDIA, AMD, Trainium, TPU)
and applies the best optimizations.

- Run `tb-doctor` to check your system compatibility
- Run `tb-validate --level quick` to validate your setup
- Run `tb-benchmark --predefined optimization --quick` to benchmark
"""


# ---------------------------------------------------------------------------
# InitCommand
# ---------------------------------------------------------------------------

class InitCommand:
    """Project initialization command implementation."""

    @staticmethod
    def register(subparsers) -> None:
        """Register the init command with argument parser."""
        parser = subparsers.add_parser(
            'init',
            help='Scaffold a new TorchBridge project',
            description='Generate a backend-agnostic PyTorch project from templates',
            formatter_class=argparse.RawDescriptionHelpFormatter,
            epilog="""
Templates:
  training     - Training pipeline with model, data loader, optimizer
  inference    - Inference/serving pipeline
  distributed  - Distributed training (DDP/FSDP)
  serving      - REST API serving with FastAPI

Examples:
  tb-init --name my_project --template training
  tb-init --name api_server --template serving
  tb-init --name big_model --template distributed --backend nvidia
            """
        )

        parser.add_argument(
            '--name',
            type=str,
            required=True,
            help='Project name (required)'
        )

        parser.add_argument(
            '--template',
            choices=['training', 'inference', 'distributed', 'serving'],
            default='training',
            help='Project template (default: training)'
        )

        parser.add_argument(
            '--backend',
            choices=['auto', 'nvidia', 'amd', 'trainium', 'tpu', 'cpu'],
            default='auto',
            help='Backend hint for config (default: auto)'
        )

        parser.add_argument(
            '--output-dir',
            type=str,
            default='.',
            help='Parent directory for generated project (default: .)'
        )

        parser.add_argument(
            '--force',
            action='store_true',
            help='Overwrite existing directory'
        )

        parser.add_argument(
            '--verbose', '-v',
            action='store_true',
            help='Enable verbose output'
        )

    @staticmethod
    def execute(args) -> int:
        """Execute the init command."""
        name = args.name
        template = getattr(args, 'template', 'training')
        backend = getattr(args, 'backend', 'auto')
        output_dir = getattr(args, 'output_dir', '.')
        force = getattr(args, 'force', False)
        verbose = getattr(args, 'verbose', False)

        project_dir = Path(output_dir) / name

        print(f" Initializing TorchBridge project: {name}")
        print(f"   Template: {template}")
        print(f"   Backend:  {backend}")
        print(f"   Location: {project_dir}")

        try:
            # Create project directory
            if project_dir.exists():
                if not force:
                    print(f" Directory already exists: {project_dir}")
                    print("   Use --force to overwrite")
                    return 1
                if verbose:
                    print("   Overwriting existing directory")

            project_dir.mkdir(parents=True, exist_ok=True)

            # Generate files based on template
            files = InitCommand._generate_files(name, template, backend)

            for filename, content in files.items():
                filepath = project_dir / filename
                filepath.parent.mkdir(parents=True, exist_ok=True)
                filepath.write_text(content)
                if verbose:
                    print(f"   Created: {filepath}")

            print(f"\n Project created: {project_dir}")
            print(f"   {len(files)} files generated")
            print(f"\n   cd {project_dir}")
            print("   pip install -r requirements.txt")
            if template in ('training', 'distributed'):
                print("   python train.py")
            else:
                print("   python serve.py")

            return 0

        except Exception as e:
            print(f" Project initialization failed: {e}")
            if verbose:
                import traceback
                traceback.print_exc()
            return 1

    @staticmethod
    def _generate_files(name: str, template: str, backend: str) -> dict[str, str]:
        """Generate project files based on template."""
        files: dict[str, str] = {}

        # Common files
        files['.gitignore'] = _GITIGNORE
        files['README.md'] = _readme(name, template)

        if template == 'training':
            files['train.py'] = _train_py(name)
            files['config.yaml'] = _config_for_training(name, backend)
            files['requirements.txt'] = _REQUIREMENTS_BASE
            files['Dockerfile'] = _dockerfile_training(name)

        elif template == 'inference':
            files['serve.py'] = _serve_py(name)
            files['config.yaml'] = _config_for_inference(name, backend)
            files['requirements.txt'] = _REQUIREMENTS_BASE
            files['Dockerfile'] = _dockerfile_inference(name)

        elif template == 'distributed':
            files['train.py'] = _train_py(name)
            files['config.yaml'] = _config_for_distributed(name, backend)
            files['requirements.txt'] = _REQUIREMENTS_DISTRIBUTED
            files['Dockerfile'] = _dockerfile_distributed(name)

        elif template == 'serving':
            files['serve.py'] = _serve_py(name)
            files['config.yaml'] = _config_for_inference(name, backend)
            files['requirements.txt'] = _REQUIREMENTS_SERVING
            files['Dockerfile'] = _dockerfile_serving(name)

        return files


def main():
    """Standalone entry point for tb-init."""
    parser = argparse.ArgumentParser(
        prog='tb-init',
        description='Scaffold a new TorchBridge project',
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )

    parser.add_argument(
        '--name',
        type=str,
        required=True,
        help='Project name (required)'
    )
    parser.add_argument(
        '--template',
        choices=['training', 'inference', 'distributed', 'serving'],
        default='training',
        help='Project template (default: training)'
    )
    parser.add_argument(
        '--backend',
        choices=['auto', 'nvidia', 'amd', 'trainium', 'tpu', 'cpu'],
        default='auto',
        help='Backend hint for config (default: auto)'
    )
    parser.add_argument(
        '--output-dir',
        type=str,
        default='.',
        help='Parent directory for generated project (default: .)'
    )
    parser.add_argument(
        '--force',
        action='store_true',
        help='Overwrite existing directory'
    )
    parser.add_argument(
        '--verbose', '-v',
        action='store_true',
        help='Enable verbose output'
    )

    args = parser.parse_args()
    try:
        return InitCommand.execute(args)
    except KeyboardInterrupt:
        print("\nOperation cancelled by user")
        return 130
    except Exception as e:
        from torchbridge.cli import _print_error
        return _print_error(e, verbose=getattr(args, 'verbose', False))


if __name__ == '__main__':
    sys.exit(main())
