"""
PyTorch/XLA Integration Module

Specialized utilities for PyTorch/XLA integration, including
device management, distributed training, and XLA-specific optimizations.
"""

import logging
import os
import time
import warnings
from typing import Any

import torch
import torch.nn as nn

from torchbridge.core.config import (
    TPUConfig,
    TPUVersion,
)

from . import xla_compat

logger = logging.getLogger(__name__)


class XLADeviceManager:
    """
    XLA device management for TPU operations.

    Provides device coordination, memory management, and distributed setup
    for PyTorch/XLA environments.
    """

    def __init__(self, config: TPUConfig):
        """
        Initialize XLA device manager.

        Args:
            config: TPU configuration
        """
        self.config = config
        self._devices: list[torch.device] = []
        self._current_device: torch.device | None = None
        self._world_size = 1
        self._rank = 0

        self._setup_xla_devices()

    def _setup_xla_devices(self) -> None:
        """Set up XLA devices for TPU."""
        try:
            import torch_xla  # noqa: F401

            # Get current device using compatibility layer
            self._current_device = xla_compat.get_xla_device()
            self._devices = [self._current_device]  # Simplified device list

            # Set up distributed environment using compatibility layer
            world_size = xla_compat.get_world_size()
            if world_size > 1:
                self._world_size = world_size
                self._rank = xla_compat.get_ordinal()

            logger.info(
                "XLA Device Manager initialized: devices=%d, current_device=%s, world_size=%d, rank=%d",
                len(self._devices),
                self._current_device,
                self._world_size,
                self._rank
            )

        except ImportError:
            warnings.warn("PyTorch/XLA not available. Using CPU fallback.", stacklevel=2)
            self._current_device = torch.device("cpu")

    @property
    def device(self) -> torch.device | None:
        """Get current XLA device."""
        return self._current_device

    @property
    def devices(self) -> list[torch.device]:
        """Get all available XLA devices."""
        return self._devices

    @property
    def world_size(self) -> int:
        """Get world size for distributed training."""
        return self._world_size

    @property
    def rank(self) -> int:
        """Get current process rank."""
        return self._rank

    def set_device(self, device_id: int) -> None:
        """Set current XLA device."""
        if device_id < len(self._devices):
            self._current_device = self._devices[device_id]
        else:
            warnings.warn(f"Device {device_id} not available, using default device", stacklevel=2)

    def sync_all_devices(self) -> None:
        """Synchronize all XLA devices."""
        try:
            xla_compat.sync()
            if self.world_size > 1:
                xla_compat.rendezvous('sync_all')
        except Exception:
            logger.debug("XLA device synchronization failed", exc_info=True)
            pass

    def get_device_stats(self) -> dict[str, Any]:
        """Get device statistics."""
        stats = {
            'current_device': str(self._current_device),
            'available_devices': len(self._devices),
            'world_size': self._world_size,
            'rank': self._rank
        }

        try:
            stats.update({
                'xla_device_count': xla_compat.get_device_count(),
            })
        except Exception:
            logger.debug("XLA device count query failed", exc_info=True)
            pass

        return stats


class XLADistributedTraining:
    """
    PyTorch/XLA distributed training utilities.

    Provides distributed training setup, gradient synchronization,
    and collective operations for TPU pods.
    """

    def __init__(self, device_manager: XLADeviceManager):
        """
        Initialize distributed training.

        Args:
            device_manager: XLA device manager
        """
        self.device_manager = device_manager
        self._process_group = None
        self._is_initialized = False

        if self.device_manager.world_size > 1:
            self._setup_distributed()

    def _setup_distributed(self) -> None:
        """Set up distributed training environment."""
        try:
            import torch_xla.distributed.xla_backend as xla_backend  # noqa: F401

            if not torch.distributed.is_initialized():
                # Initialize process group
                torch.distributed.init_process_group(
                    backend='xla',
                    rank=self.device_manager.rank,
                    world_size=self.device_manager.world_size
                )
                self._is_initialized = True

                logger.info(
                    "Distributed training initialized: backend=xla, world_size=%d, rank=%d",
                    self.device_manager.world_size,
                    self.device_manager.rank
                )

        except ImportError:
            warnings.warn("XLA distributed backend not available", stacklevel=2)

    @property
    def is_distributed(self) -> bool:
        """Check if distributed training is enabled."""
        return self._is_initialized and self.device_manager.world_size > 1

    def wrap_model(self, model: nn.Module) -> nn.Module:
        """
        Wrap model for distributed training.

        Args:
            model: Model to wrap

        Returns:
            Wrapped model for distributed training
        """
        if not self.is_distributed:
            return model

        try:
            # Use XLA's parallel wrapper if available
            import torch_xla.distributed.parallel_loader as pl  # noqa: F401

            # For now, just move model to device and return
            # Full DDP implementation would require more XLA-specific code
            model = model.to(self.device_manager.device)
            return model

        except ImportError:
            warnings.warn("XLA parallel loader not available", stacklevel=2)
            return model

    def all_reduce(self, tensor: torch.Tensor, op: str = "sum") -> torch.Tensor:
        """
        Perform all-reduce operation across TPU cores.

        Args:
            tensor: Tensor to reduce
            op: Reduction operation ('sum', 'mean', 'max', 'min')

        Returns:
            Reduced tensor
        """
        if not self.is_distributed:
            return tensor

        try:
            import torch_xla.core.xla_model as xm

            if op == "sum":
                return xm.all_reduce('sum', tensor)
            elif op == "mean":
                result = xm.all_reduce('sum', tensor)
                return result / self.device_manager.world_size
            elif op == "max":
                return xm.all_reduce('max', tensor)
            elif op == "min":
                return xm.all_reduce('min', tensor)
            else:
                raise ValueError(f"Unsupported reduction operation: {op}")

        except ImportError:
            return tensor

    def all_gather(self, tensor: torch.Tensor) -> list[torch.Tensor]:
        """
        Gather tensors from all processes.

        Args:
            tensor: Tensor to gather

        Returns:
            List of tensors from all processes
        """
        if not self.is_distributed:
            return [tensor]

        try:
            import torch_xla.core.xla_model as xm
            return xm.all_gather(tensor)
        except ImportError:
            return [tensor]

    def barrier(self) -> None:
        """Synchronize all processes."""
        if self.is_distributed:
            try:
                import torch_xla.core.xla_model as xm
                xm.rendezvous('barrier')
            except ImportError:
                pass


class XLAOptimizations:
    """
    XLA-specific optimization utilities.

    Provides model transformations, compilation hints,
    and performance optimizations for XLA/TPU.
    """

    def __init__(self, config: TPUConfig):
        """
        Initialize XLA optimizations.

        Args:
            config: TPU configuration
        """
        self.config = config

    def optimize_model_for_xla(self, model: nn.Module) -> nn.Module:
        """
        Apply XLA-specific model optimizations.

        Args:
            model: Model to optimize

        Returns:
            Optimized model
        """
        # Apply layer-specific optimizations
        model = self._optimize_linear_layers(model)
        model = self._optimize_attention_layers(model)
        model = self._optimize_activation_functions(model)

        return model

    def _optimize_linear_layers(self, model: nn.Module) -> nn.Module:
        """Optimize linear layers for XLA."""
        for name, module in model.named_modules():
            if isinstance(module, nn.Linear):
                # Ensure weight dimensions are optimal for TPU matrix units
                in_features, out_features = module.in_features, module.out_features

                if in_features % 8 != 0 or out_features % 8 != 0:
                    warnings.warn(
                        f"Linear layer {name} dimensions ({in_features}x{out_features}) "
                        "not optimal for TPU. Consider padding to multiples of 8.",
                    stacklevel=2,
                    )

        return model

    def _optimize_attention_layers(self, model: nn.Module) -> nn.Module:
        """Optimize attention mechanisms for XLA."""
        for name, module in model.named_modules():
            # Look for attention patterns
            if hasattr(module, 'num_attention_heads'):
                head_dim = getattr(module, 'attention_head_size', None)

                if head_dim and head_dim % 8 != 0:
                    warnings.warn(
                        f"Attention layer {name} head dimension {head_dim} "
                        "not optimal for TPU. Consider using dimensions divisible by 8.",
                    stacklevel=2,
                    )

        return model

    def _optimize_activation_functions(self, model: nn.Module) -> nn.Module:
        """Optimize activation functions for XLA."""
        # XLA has optimized implementations for certain activations
        for _name, module in model.named_modules():
            if isinstance(module, nn.ReLU):
                # ReLU is well-optimized in XLA
                continue
            elif isinstance(module, nn.GELU):
                # GELU has good XLA support
                continue
            elif isinstance(module, nn.SiLU):
                # SiLU (Swish) is also well-supported
                continue

        return model

    def add_compilation_hints(self, model: nn.Module) -> nn.Module:
        """Add compilation hints for better XLA optimization."""

        # Add static shape hints where possible
        for module in model.modules():
            # Mark modules that have static shapes
            if hasattr(module, 'forward'):
                # Add metadata for XLA compiler
                module._xla_static_shapes = True  # type: ignore[assignment]

        return model

    def enable_dynamic_shapes(self, model: nn.Module) -> nn.Module:
        """Enable dynamic shape support for flexible inputs."""

        if self.config.enable_xla_dynamic_shapes:
            # Configure model for dynamic shapes
            for module in model.modules():
                if hasattr(module, '_xla_static_shapes'):
                    delattr(module, '_xla_static_shapes')

        return model


class XLAUtilities:
    """
    General XLA utilities and helpers.

    Provides debugging, profiling, and utility functions
    for XLA/TPU development.
    """

    @staticmethod
    def get_xla_env_info() -> dict[str, Any]:
        """Get XLA environment information."""
        env_info: dict[str, Any] = {
            'XLA_FLAGS': os.environ.get('XLA_FLAGS', ''),
            'XLA_PYTHON_CLIENT_MEM_FRACTION': os.environ.get('XLA_PYTHON_CLIENT_MEM_FRACTION', ''),
            'TPU_TYPE': os.environ.get('TPU_TYPE', ''),
            'TPU_NAME': os.environ.get('TPU_NAME', ''),
        }

        try:
            import torch_xla
            env_info.update({
                'torch_xla_version': torch_xla.__version__,
                'xla_available': True
            })

            env_info.update({
                'xla_device_count': xla_compat.get_device_count(),
                'xrt_world_size': xla_compat.get_world_size()
            })
        except Exception:
            logger.debug("XLA environment info detection failed", exc_info=True)
            env_info['xla_available'] = False

        return env_info

    @staticmethod
    def profile_xla_compilation(model: nn.Module,
                              sample_input: torch.Tensor) -> dict[str, Any]:
        """Profile XLA compilation performance."""

        try:
            # Use compatibility layer for device access
            device = xla_compat.get_xla_device()
            model = model.to(device)
            sample_input = sample_input.to(device)

            # Profile compilation
            start_time = time.time()

            with torch.no_grad():
                output = model(sample_input)
                xla_compat.sync()  # Force compilation

            compilation_time = time.time() - start_time

            # Profile execution (post-compilation)
            start_time = time.time()
            for _ in range(10):
                with torch.no_grad():
                    output = model(sample_input)
                    xla_compat.sync()

            execution_time = (time.time() - start_time) / 10

            return {
                'compilation_time': compilation_time,
                'avg_execution_time': execution_time,
                'output_shape': list(output.shape),
                'device': str(device)
            }

        except Exception:
            logger.debug("XLA compilation profiling failed", exc_info=True)
            return {'error': 'PyTorch/XLA not available'}

    @staticmethod
    def debug_xla_graph(model: nn.Module,
                       sample_input: torch.Tensor) -> str:
        """Get XLA computation graph for debugging."""

        try:
            import torch_xla.debug.metrics as met

            device = xla_compat.get_xla_device()
            model = model.to(device)
            sample_input = sample_input.to(device)

            # Clear existing metrics
            met.clear_all()

            # Run model to generate graph
            with torch.no_grad():
                model(sample_input)
                xla_compat.sync()

            # Get metrics report
            metrics_report = met.metrics_report()
            return metrics_report

        except Exception:
            logger.debug("XLA graph debugging failed", exc_info=True)
            return "PyTorch/XLA debug utilities not available"

    @staticmethod
    def optimize_xla_flags(tpu_version: TPUVersion) -> dict[str, str]:
        """Get optimized XLA flags for specific TPU versions."""

        base_flags = {
            'XLA_FLAGS': '--xla_optimization_level=2 --xla_force_host_platform_device_count=1'
        }

        # Version-specific optimizations
        if tpu_version in [TPUVersion.V5P, TPUVersion.V6E, TPUVersion.V7]:
            # High-performance TPUs
            base_flags['XLA_FLAGS'] += ' --xla_enable_async_collectives=true'
            base_flags['XLA_FLAGS'] += ' --xla_tpu_enable_async_collective_fusion=true'

        elif tpu_version == TPUVersion.V5E:
            # Cost-optimized TPUs
            base_flags['XLA_FLAGS'] += ' --xla_optimization_level=1'

        elif tpu_version == TPUVersion.V4:
            # Legacy TPUs
            base_flags['XLA_FLAGS'] = '--xla_optimization_level=1 --xla_force_host_platform_device_count=1'

        return base_flags


# Integration factory function
def create_xla_integration(config: TPUConfig) -> tuple[XLADeviceManager, XLADistributedTraining, XLAOptimizations]:
    """
    Create complete XLA integration setup.

    Args:
        config: TPU configuration

    Returns:
        Tuple of (device_manager, distributed_training, optimizations)
    """
    device_manager = XLADeviceManager(config)
    distributed_training = XLADistributedTraining(device_manager)
    optimizations = XLAOptimizations(config)

    return device_manager, distributed_training, optimizations
