"""PyTorch Profiler environment configuration.

Sets up environment variables and provides helpers for using
torch.profiler.profile() to collect distributed training traces.

Reference:
- PyTorch Profiler: https://pytorch.org/tutorials/recipes/recipes/profiler_recipe.html
- TensorBoard: https://pytorch.org/tutorials/intermediate/tensorboard_profiler_tutorial.html
"""

import os
from dataclasses import dataclass
from pathlib import Path
from typing import Optional


@dataclass
class PyTorchProfilerConfig:
    """Configuration for PyTorch Profiler.

    Attributes:
        output_dir: Directory for trace output
        record_shapes: Record tensor shapes
        profile_memory: Profile memory usage
        with_stack: Record Python stack traces
        with_flops: Estimate FLOPs
        with_modules: Record module hierarchy
        export_chrome_trace: Export Chrome trace format
        export_stacks: Export stack traces
    """

    output_dir: str = "./pytorch_traces"
    record_shapes: bool = True
    profile_memory: bool = True
    with_stack: bool = True
    with_flops: bool = False
    with_modules: bool = True
    export_chrome_trace: bool = True
    export_stacks: bool = False


class PyTorchCollector:
    """Helper for collecting traces with PyTorch Profiler."""

    def __init__(self, config: Optional[PyTorchProfilerConfig] = None) -> None:
        """Initialize the collector.

        Args:
            config: PyTorch Profiler configuration
        """
        self.config = config or PyTorchProfilerConfig()

    def get_env_vars(self) -> dict[str, str]:
        """Get environment variables for PyTorch Profiler.

        Returns:
            Dictionary of environment variables
        """
        env = {}

        # Trace output directory
        output_dir = Path(self.config.output_dir)
        output_dir.mkdir(parents=True, exist_ok=True)

        env["TORCH_TRACE_OUTPUT_DIR"] = str(output_dir.absolute())

        # Enable profiler-friendly settings
        env["TORCH_SHOW_CPP_STACKTRACES"] = "1"

        # CUDA synchronization for accurate timing
        env["CUDA_LAUNCH_BLOCKING"] = "0"  # Keep async for normal operation

        return env

    def get_profiler_code(
        self,
        wait_steps: int = 1,
        warmup_steps: int = 1,
        active_steps: int = 3,
        repeat: int = 1,
    ) -> str:
        """Generate Python code for using torch.profiler.

        This generates code that can be inserted into training scripts
        to enable profiling.

        Args:
            wait_steps: Steps to wait before profiling
            warmup_steps: Warmup steps
            active_steps: Steps to actively profile
            repeat: Number of cycles to repeat

        Returns:
            Python code string
        """
        output_dir = Path(self.config.output_dir).absolute()

        code = f'''
import torch
from torch.profiler import profile, record_function, ProfilerActivity, schedule, tensorboard_trace_handler

# Profiler configuration
profiler_schedule = schedule(
    wait={wait_steps},
    warmup={warmup_steps},
    active={active_steps},
    repeat={repeat}
)

# Activities to profile
activities = [
    ProfilerActivity.CPU,
    ProfilerActivity.CUDA,
]

# Profiler context
with profile(
    activities=activities,
    schedule=profiler_schedule,
    on_trace_ready=lambda p: p.export_chrome_trace(
        "{output_dir}/trace_rank{{rank}}_step{{step}}.json".format(
            rank=torch.distributed.get_rank() if torch.distributed.is_initialized() else 0,
            step=p.step_num
        )
    ),
    record_shapes={self.config.record_shapes},
    profile_memory={self.config.profile_memory},
    with_stack={self.config.with_stack},
    with_flops={self.config.with_flops},
    with_modules={self.config.with_modules},
) as prof:
    # Training loop
    for step, batch in enumerate(dataloader):
        with record_function("forward"):
            output = model(batch)
        
        with record_function("backward"):
            loss.backward()
        
        with record_function("optimizer"):
            optimizer.step()
            optimizer.zero_grad()
        
        # Step the profiler
        prof.step()
'''
        return code.strip()

    def get_minimal_profiler_code(self) -> str:
        """Generate minimal profiling code for quick integration.

        Returns:
            Python code string for minimal profiling
        """
        output_dir = Path(self.config.output_dir).absolute()

        code = f'''
import torch
from torch.profiler import profile, ProfilerActivity

def profile_training_step(model, batch, optimizer, loss_fn):
    """Profile a single training step."""
    with profile(
        activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
        record_shapes=True,
        with_stack=True,
    ) as prof:
        output = model(batch)
        loss = loss_fn(output, batch.targets)
        loss.backward()
        optimizer.step()
        optimizer.zero_grad()
    
    # Export trace
    rank = torch.distributed.get_rank() if torch.distributed.is_initialized() else 0
    prof.export_chrome_trace("{output_dir}/trace_rank{{rank}}.json".format(rank=rank))
    
    return prof
'''
        return code.strip()

    def setup_output_dir(self) -> Path:
        """Create and return the output directory.

        Returns:
            Path to output directory
        """
        output_dir = Path(self.config.output_dir)
        output_dir.mkdir(parents=True, exist_ok=True)
        return output_dir

    def find_trace_files(self) -> list[Path]:
        """Find trace files in the output directory.

        Returns:
            List of trace file paths
        """
        output_dir = Path(self.config.output_dir)
        if not output_dir.exists():
            return []

        files = []
        for pattern in ["*.json", "*.json.gz", "*.pt.trace.json"]:
            files.extend(output_dir.glob(pattern))

        return sorted(files)


def get_pytorch_collection_env(
    output_dir: str = "./traces",
) -> dict[str, str]:
    """Convenience function to get PyTorch collection environment variables.

    Args:
        output_dir: Output directory for traces

    Returns:
        Dictionary of environment variables
    """
    config = PyTorchProfilerConfig(output_dir=output_dir)
    collector = PyTorchCollector(config)
    return collector.get_env_vars()
