"""
System diagnostics commands for TorchBridge CLI.

Provides comprehensive system compatibility checking and backend recommendations.
"""

import argparse
import platform
import subprocess
import sys
from dataclasses import dataclass

import torch

import torchbridge


@dataclass
class DiagnosticResult:
    """Result from a single diagnostic check."""
    name: str
    status: str  # "pass", "warning", "fail"
    message: str
    details: str | None = None
    recommendation: str | None = None


class DoctorCommand:
    """System diagnostics command implementation."""

    @staticmethod
    def register(subparsers) -> None:
        """Register the doctor command with argument parser."""
        parser = subparsers.add_parser(
            'doctor',
            help='Diagnose system compatibility and backend readiness',
            description='Check system configuration for optimal TorchBridge performance',
            formatter_class=argparse.RawDescriptionHelpFormatter,
            epilog="""
Check Categories:
  basic      - Python, PyTorch, and basic dependencies
  hardware   - GPU detection and capabilities
  optimization - Backend framework availability
  advanced   - Advanced features (Triton, CUDA kernels)

Examples:
  tb-doctor                    # Quick system check
  tb-doctor --full-report      # Comprehensive diagnostics
  tb-doctor --category hardware # Hardware-specific checks
  tb-doctor --fix             # Attempt to fix detected issues
            """
        )

        parser.add_argument(
            '--category',
            choices=['basic', 'hardware', 'optimization', 'advanced'],
            help='Focus on specific diagnostic category'
        )

        parser.add_argument(
            '--full-report',
            action='store_true',
            help='Run comprehensive diagnostics (all categories)'
        )

        parser.add_argument(
            '--fix',
            action='store_true',
            help='Print remediation steps for detected issues (does not modify your system)'
        )

        parser.add_argument(
            '--output', '-o',
            type=str,
            help='Save diagnostic report to file'
        )

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

        parser.add_argument(
            '--ci',
            action='store_true',
            help='CI mode: JSON to stdout, no color, structured exit codes'
        )

    @staticmethod
    def execute(args) -> int:
        """Execute the doctor command."""
        ci_mode = getattr(args, 'ci', False)

        if not ci_mode:
            print(" TorchBridge System Diagnostics")
            print("=" * 50)

        try:
            results = []

            # Determine which checks to run
            if args.category:
                if args.category == 'basic':
                    results.extend(DoctorCommand._check_basic_requirements(args.verbose))
                elif args.category == 'hardware':
                    results.extend(DoctorCommand._check_hardware(args.verbose))
                elif args.category == 'optimization':
                    results.extend(DoctorCommand._check_optimization_frameworks(args.verbose))
                elif args.category == 'advanced':
                    results.extend(DoctorCommand._check_advanced_features(args.verbose))
                else:
                    raise ValueError(f"Invalid category: {args.category}")
            elif args.full_report:
                results.extend(DoctorCommand._check_basic_requirements(args.verbose))
                results.extend(DoctorCommand._check_hardware(args.verbose))
                results.extend(DoctorCommand._check_optimization_frameworks(args.verbose))
                results.extend(DoctorCommand._check_advanced_features(args.verbose))
            else:
                # Quick check - basic + hardware
                results.extend(DoctorCommand._check_basic_requirements(args.verbose))
                results.extend(DoctorCommand._check_hardware(args.verbose))

            # CI mode: JSON output and structured exit codes
            if ci_mode:
                return DoctorCommand._output_ci_json(results)

            # Display results
            DoctorCommand._display_results(results, args.verbose)

            # Generate summary
            summary = DoctorCommand._generate_summary(results)
            print(summary)

            # Fix issues if requested
            if args.fix:
                DoctorCommand._attempt_fixes(results, args.verbose)

            # Save report if requested
            if args.output:
                DoctorCommand._save_report(results, args.output, args.verbose)

            # Return appropriate exit code
            has_failures = any(r.status == 'fail' for r in results)
            return 1 if has_failures else 0

        except Exception as e:
            if ci_mode:
                import json
                print(json.dumps({"error": str(e)}))
                return 1
            print(f" Diagnostics failed: {e}")
            if args.verbose:
                import traceback
                traceback.print_exc()
            return 1

    @staticmethod
    def _check_basic_requirements(verbose: bool) -> list[DiagnosticResult]:
        """Check basic Python and PyTorch requirements."""
        if verbose:
            print(" Checking basic requirements...")

        results = []

        # Python version
        python_version = platform.python_version()
        python_major, python_minor = map(int, python_version.split('.')[:2])

        if python_major >= 3 and python_minor >= 8:
            results.append(DiagnosticResult(
                "Python Version",
                "pass",
                f"Python {python_version} ( Compatible)"
            ))
        else:
            results.append(DiagnosticResult(
                "Python Version",
                "fail",
                f"Python {python_version} (Requires Python 3.10+)",
                recommendation="Upgrade to Python 3.10 or later"
            ))

        # PyTorch version
        try:
            torch_version = torch.__version__
            torch_major, torch_minor = map(int, torch_version.split('.')[:2])

            if torch_major >= 2:
                results.append(DiagnosticResult(
                    "PyTorch Version",
                    "pass",
                    f"PyTorch {torch_version} ( Compatible)"
                ))
            elif torch_major == 1 and torch_minor >= 12:
                results.append(DiagnosticResult(
                    "PyTorch Version",
                    "warning",
                    f"PyTorch {torch_version} ( Recommend 2.0+)",
                    recommendation="Upgrade to PyTorch 2.0+ for best performance"
                ))
            else:
                results.append(DiagnosticResult(
                    "PyTorch Version",
                    "fail",
                    f"PyTorch {torch_version} ( Requires 1.12+)",
                    recommendation="Upgrade to PyTorch 2.0+ for full compatibility"
                ))
        except Exception as e:
            results.append(DiagnosticResult(
                "PyTorch Installation",
                "fail",
                f"PyTorch not found ({e})",
                recommendation="Install PyTorch: pip install torch"
            ))

        # NumPy version
        try:
            import numpy as np
            numpy_version = np.__version__
            results.append(DiagnosticResult(
                "NumPy Version",
                "pass",
                f"NumPy {numpy_version} ( Available)"
            ))
        except ImportError:
            results.append(DiagnosticResult(
                "NumPy Installation",
                "fail",
                "NumPy not found",
                recommendation="Install NumPy: pip install numpy"
            ))

        # TorchBridge installation
        try:
            kpt_version = torchbridge.__version__
            results.append(DiagnosticResult(
                "TorchBridge Version",
                "pass",
                f"TorchBridge {kpt_version} ( Available)"
            ))
        except Exception as e:
            results.append(DiagnosticResult(
                "TorchBridge Installation",
                "fail",
                f"TorchBridge not properly installed ({e})",
                recommendation="Reinstall: pip install -e ."
            ))

        return results

    @staticmethod
    def _check_hardware(verbose: bool) -> list[DiagnosticResult]:
        """Check hardware capabilities and GPU availability."""
        if verbose:
            print(" Checking hardware capabilities...")

        results = []

        # CUDA availability
        if torch.cuda.is_available():
            gpu_name = torch.cuda.get_device_name(0)
            gpu_memory = torch.cuda.get_device_properties(0).total_memory / 1e9
            cuda_version = torch.version.cuda

            results.append(DiagnosticResult(
                "CUDA GPU",
                "pass",
                f"{gpu_name} with {gpu_memory:.1f} GB ( Available)",
                details=f"CUDA {cuda_version}"
            ))

            # Check GPU compute capability
            major, minor = torch.cuda.get_device_properties(0).major, torch.cuda.get_device_properties(0).minor
            compute_capability = f"{major}.{minor}"

            if major >= 7:  # Volta or newer
                results.append(DiagnosticResult(
                    "GPU Compute Capability",
                    "pass",
                    f"Compute {compute_capability} ( Excellent for HAL features)",
                    details="Supports Tensor Cores and advanced backend features"
                ))
            elif major >= 6:  # Pascal
                results.append(DiagnosticResult(
                    "GPU Compute Capability",
                    "warning",
                    f"Compute {compute_capability} ( Good, but older)",
                    recommendation="Consider upgrading for best performance"
                ))
            else:
                results.append(DiagnosticResult(
                    "GPU Compute Capability",
                    "warning",
                    f"Compute {compute_capability} ( Limited backend feature support)",
                    recommendation="GPU may not support all backend features"
                ))

            # Check memory
            if gpu_memory >= 8.0:
                results.append(DiagnosticResult(
                    "GPU Memory",
                    "pass",
                    f"{gpu_memory:.1f} GB ( Sufficient for most workloads)"
                ))
            elif gpu_memory >= 4.0:
                results.append(DiagnosticResult(
                    "GPU Memory",
                    "warning",
                    f"{gpu_memory:.1f} GB ( May limit large models)",
                    recommendation="Consider larger GPU for production workloads"
                ))
            else:
                results.append(DiagnosticResult(
                    "GPU Memory",
                    "warning",
                    f"{gpu_memory:.1f} GB ( Limited memory for advanced features)",
                    recommendation="Upgrade GPU or use CPU-only backend"
                ))
        else:
            results.append(DiagnosticResult(
                "CUDA GPU",
                "warning",
                "No CUDA GPU detected ( CPU-only mode)",
                recommendation="Install CUDA-compatible PyTorch for GPU acceleration"
            ))

        # Apple Silicon MPS
        if hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
            results.append(DiagnosticResult(
                "Apple Silicon GPU",
                "pass",
                "Apple Silicon GPU ( Available)",
                details="MPS backend available for acceleration"
            ))

        # CPU info
        import multiprocessing
        cpu_count = multiprocessing.cpu_count()
        results.append(DiagnosticResult(
            "CPU Cores",
            "pass",
            f"{cpu_count} cores ( Available)",
            details=f"Platform: {platform.machine()}"
        ))

        return results

    @staticmethod
    def _check_optimization_frameworks(verbose: bool) -> list[DiagnosticResult]:
        """Check availability of backend frameworks and compilation tools."""
        if verbose:
            print(" Checking backend frameworks...")

        results = []

        # torch.compile (PyTorch 2.0+)
        try:
            if hasattr(torch, 'compile'):
                # Test torch.compile
                test_model = torch.nn.Linear(10, 1)
                torch.compile(test_model, mode='default')
                results.append(DiagnosticResult(
                    "torch.compile",
                    "pass",
                    "torch.compile ( Available and working)",
                    details="PyTorch 2.0+ compilation framework"
                ))
            else:
                results.append(DiagnosticResult(
                    "torch.compile",
                    "warning",
                    "torch.compile not available ( Requires PyTorch 2.0+)",
                    recommendation="Upgrade to PyTorch 2.0+ for compilation features"
                ))
        except Exception as e:
            results.append(DiagnosticResult(
                "torch.compile",
                "fail",
                f"torch.compile not working ({e})",
                recommendation="Check PyTorch installation"
            ))

        # TorchScript
        try:
            test_model = torch.nn.Linear(10, 1)
            test_input = torch.randn(1, 10)
            torch.jit.trace(test_model, test_input)
            results.append(DiagnosticResult(
                "TorchScript",
                "pass",
                "TorchScript JIT ( Available and working)"
            ))
        except Exception as e:
            results.append(DiagnosticResult(
                "TorchScript",
                "warning",
                f"TorchScript issues detected ({e})",
                recommendation="Check for model compatibility issues"
            ))

        return results

    @staticmethod
    def _check_advanced_features(verbose: bool) -> list[DiagnosticResult]:
        """Check availability of advanced backend features."""
        if verbose:
            print(" Checking advanced features...")

        results = []

        # Triton
        try:
            import triton
            triton_version = triton.__version__
            results.append(DiagnosticResult(
                "Triton Kernels",
                "pass",
                f"Triton {triton_version} ( Available)",
                details="Python-based GPU kernel development"
            ))
        except ImportError:
            if torch.cuda.is_available():
                results.append(DiagnosticResult(
                    "Triton Kernels",
                    "warning",
                    "Triton not available ( Optional for advanced kernels)",
                    recommendation="Install Triton: pip install triton"
                ))
            else:
                results.append(DiagnosticResult(
                    "Triton Kernels",
                    "warning",
                    "Triton not available ( Requires CUDA GPU)",
                    details="Triton requires CUDA-compatible hardware"
                ))

        # Flash Attention
        try:
            import flash_attn  # noqa: F401
            results.append(DiagnosticResult(
                "Flash Attention",
                "pass",
                "Flash Attention ( Available)",
                details="Optimized attention implementation"
            ))
        except ImportError:
            if torch.cuda.is_available():
                results.append(DiagnosticResult(
                    "Flash Attention",
                    "warning",
                    "Flash Attention not available ( Optional optimization)",
                    recommendation="Install: pip install flash-attn"
                ))

        # CUDA toolkit (for custom kernels)
        try:
            result = subprocess.run(['nvcc', '--version'], capture_output=True, text=True)
            if result.returncode == 0:
                results.append(DiagnosticResult(
                    "CUDA Toolkit",
                    "pass",
                    "NVCC compiler ( Available)",
                    details="Can compile custom CUDA kernels"
                ))
            else:
                results.append(DiagnosticResult(
                    "CUDA Toolkit",
                    "warning",
                    "CUDA Toolkit not found ( Optional for custom kernels)",
                    recommendation="Install CUDA Toolkit for kernel development"
                ))
        except FileNotFoundError:
            results.append(DiagnosticResult(
                "CUDA Toolkit",
                "warning",
                "NVCC not in PATH ( Optional for custom kernels)",
                recommendation="Add CUDA to PATH or install CUDA Toolkit"
            ))

        return results

    @staticmethod
    def _output_ci_json(results: list[DiagnosticResult]) -> int:
        """Output results as JSON for CI mode with structured exit codes.

        Exit codes: 0=all pass, 1=failures, 2=warnings only.
        """
        import json
        import time

        report = {
            'timestamp': time.time(),
            'system_info': {
                'platform': platform.platform(),
                'python_version': platform.python_version(),
                'pytorch_version': torch.__version__,
                'cuda_available': torch.cuda.is_available(),
                'gpu_name': torch.cuda.get_device_name(0) if torch.cuda.is_available() else None,
            },
            'diagnostics': [
                {
                    'name': r.name,
                    'status': r.status,
                    'message': r.message,
                    'details': r.details,
                    'recommendation': r.recommendation,
                }
                for r in results
            ],
            'summary': {
                'total': len(results),
                'passed': sum(1 for r in results if r.status == 'pass'),
                'warnings': sum(1 for r in results if r.status == 'warning'),
                'failures': sum(1 for r in results if r.status == 'fail'),
            },
        }

        print(json.dumps(report, indent=2))

        has_failures = any(r.status == 'fail' for r in results)
        has_warnings = any(r.status == 'warning' for r in results)

        if has_failures:
            return 1
        if has_warnings:
            return 2
        return 0

    @staticmethod
    def _display_results(results: list[DiagnosticResult], verbose: bool) -> None:
        """Display diagnostic results in a formatted way."""
        print("\n Diagnostic Results:")
        print("-" * 60)

        for result in results:
            status_icon = {
                'pass': '',
                'warning': '',
                'fail': ''
            }.get(result.status, '')

            print(f"{status_icon} {result.name}: {result.message}")

            if verbose and result.details:
                print(f"   Details: {result.details}")

            if result.recommendation:
                print(f"    Recommendation: {result.recommendation}")

            if verbose:
                print()

    @staticmethod
    def _generate_summary(results: list[DiagnosticResult]) -> str:
        """Generate a summary of diagnostic results."""
        total = len(results)
        passed = sum(1 for r in results if r.status == 'pass')
        warnings = sum(1 for r in results if r.status == 'warning')
        failed = sum(1 for r in results if r.status == 'fail')

        summary = f"\n Summary: {passed}/{total} checks passed"
        if warnings > 0:
            summary += f", {warnings} warnings"
        if failed > 0:
            summary += f", {failed} failures"

        if failed > 0:
            summary += "\n Critical issues detected - system may not work optimally"
        elif warnings > 0:
            summary += "\n  Some backend features may not be available"
        else:
            summary += "\n System is ready for optimal TorchBridge performance!"

        return summary

    @staticmethod
    def _attempt_fixes(results: list[DiagnosticResult], verbose: bool) -> None:
        """Attempt to fix detected issues where possible."""
        print("\n Attempting to fix detected issues...")

        fixable_issues = [r for r in results if r.status in ['fail', 'warning'] and r.recommendation]

        if not fixable_issues:
            print("   No fixable issues detected.")
            return

        for issue in fixable_issues:
            if 'pip install' in issue.recommendation:
                print(f"   Attempting to install missing package for {issue.name}...")
                # Note: In a real implementation, you might want to be more careful about automatic installation
                print(f"   ℹ  Manual action required: {issue.recommendation}")

        print("    Some fixes require manual intervention. See recommendations above.")

    @staticmethod
    def _save_report(results: list[DiagnosticResult], output_path: str, verbose: bool) -> None:
        """Save diagnostic report to file."""
        if verbose:
            print(f" Saving diagnostic report to: {output_path}")

        import json
        import time

        report = {
            'timestamp': time.time(),
            'system_info': {
                'platform': platform.platform(),
                'python_version': platform.python_version(),
                'pytorch_version': torch.__version__,
                'cuda_available': torch.cuda.is_available(),
                'gpu_name': torch.cuda.get_device_name(0) if torch.cuda.is_available() else None
            },
            'diagnostics': [
                {
                    'name': r.name,
                    'status': r.status,
                    'message': r.message,
                    'details': r.details,
                    'recommendation': r.recommendation
                }
                for r in results
            ]
        }

        with open(output_path, 'w') as f:
            json.dump(report, f, indent=2)

        if verbose:
            print(f"   Report saved ({len(results)} diagnostics)")


def main():
    """Standalone entry point for tb-doctor."""
    parser = argparse.ArgumentParser(
        prog='tb-doctor',
        description='Diagnose system compatibility and backend readiness',
        formatter_class=argparse.RawDescriptionHelpFormatter
    )

    # Add the doctor command arguments directly to the parser
    parser.add_argument(
        '--category',
        choices=['basic', 'hardware', 'optimization', 'advanced'],
        help='Focus on specific category'
    )
    parser.add_argument(
        '--full-report',
        action='store_true',
        help='Run comprehensive diagnostics (all categories)'
    )
    parser.add_argument(
        '--fix',
        action='store_true',
        help='Attempt to fix detected issues (where possible)'
    )
    parser.add_argument(
        '--output', '-o',
        type=str,
        help='Save diagnostic report to file (JSON format)'
    )
    parser.add_argument(
        '--verbose', '-v',
        action='store_true',
        help='Enable verbose output'
    )
    parser.add_argument(
        '--ci',
        action='store_true',
        help='CI mode: JSON to stdout, no color, structured exit codes'
    )

    args = parser.parse_args()
    try:
        return DoctorCommand.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())
