"""
Validation command for TorchBridge CLI.

Wraps UnifiedValidator and DoctorCommand into a structured validation pipeline
with multiple levels: quick, standard, full, and cloud.
"""

import argparse
import json
import logging
import subprocess
import sys
import time
from dataclasses import dataclass, field
from pathlib import Path

import torch

logger = logging.getLogger(__name__)


@dataclass
class ValidationResult:
    """Result from a validation step."""
    name: str
    status: str  # "pass", "warning", "fail"
    message: str
    details: str | None = None
    duration_ms: float = 0.0


@dataclass
class ValidationReport:
    """Full validation report."""
    level: str
    results: list[ValidationResult] = field(default_factory=list)
    timestamp: float = 0.0
    duration_ms: float = 0.0

    @property
    def passed(self) -> int:
        return sum(1 for r in self.results if r.status == 'pass')

    @property
    def warnings(self) -> int:
        return sum(1 for r in self.results if r.status == 'warning')

    @property
    def failures(self) -> int:
        return sum(1 for r in self.results if r.status == 'fail')

    @property
    def has_failures(self) -> bool:
        return self.failures > 0

    @property
    def has_warnings(self) -> bool:
        return self.warnings > 0


class ValidateCommand:
    """Validation pipeline command implementation."""

    @staticmethod
    def register(subparsers) -> None:
        """Register the validate command with argument parser."""
        parser = subparsers.add_parser(
            'validate',
            help='Run validation pipeline for TorchBridge',
            description='Structured validation pipeline with multiple levels',
            formatter_class=argparse.RawDescriptionHelpFormatter,
            epilog="""
Validation Levels:
  quick      - Hardware detection + import checks
  standard   - Quick + model validation + export format checks
  full       - Standard + benchmark suite + cross-backend consistency
  cloud      - Run cloud validation scripts via subprocess

Examples:
  tb-validate                          # Standard validation
  tb-validate --level quick            # Quick hardware check
  tb-validate --level full --ci        # Full validation in CI mode
  tb-validate --model model.pt         # Validate specific model
  tb-validate --output report.json     # Save report to file
            """
        )

        parser.add_argument(
            '--level',
            choices=['quick', 'standard', 'full', 'cloud'],
            default='standard',
            help='Validation level (default: standard)'
        )

        parser.add_argument(
            '--model',
            type=str,
            help='Path to a specific model to validate'
        )

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

        parser.add_argument(
            '--format',
            choices=['json', 'yaml', 'text'],
            default='text',
            help='Output format (default: text)'
        )

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

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

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

        if not ci_mode:
            print(" TorchBridge Validation Pipeline")
            print("=" * 50)
            print(f" Level: {level}")

        start_time = time.time()
        report = ValidationReport(level=level, timestamp=start_time)

        try:
            # Quick level: hardware detection + import checks
            report.results.extend(ValidateCommand._run_quick_checks(verbose))

            # Standard level: add model validation + export checks
            if level in ('standard', 'full'):
                model_path = getattr(args, 'model', None)
                report.results.extend(
                    ValidateCommand._run_standard_checks(model_path, verbose)
                )

            # Full level: add benchmark suite + cross-backend
            if level == 'full':
                report.results.extend(ValidateCommand._run_full_checks(verbose))

            # Cloud level: run cloud validation scripts
            if level == 'cloud':
                report.results.extend(ValidateCommand._run_cloud_checks(verbose))

            report.duration_ms = (time.time() - start_time) * 1000

            # Output results
            if ci_mode:
                return ValidateCommand._output_ci_json(report)

            ValidateCommand._display_report(report, verbose)

            # Save report if requested
            output_path = getattr(args, 'output', None)
            if output_path:
                fmt = getattr(args, 'format', 'text')
                ValidateCommand._save_report(report, output_path, fmt, verbose)

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

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

    @staticmethod
    def _run_quick_checks(verbose: bool) -> list[ValidationResult]:
        """Run quick checks: hardware detection + import checks."""
        results = []

        if verbose:
            print(" Running quick checks...")

        # Import checks
        start = time.time()
        try:
            import torchbridge  # noqa: F811
            results.append(ValidationResult(
                "TorchBridge Import",
                "pass",
                f"TorchBridge {torchbridge.__version__} imported successfully",
                duration_ms=(time.time() - start) * 1000,
            ))
        except ImportError as e:
            results.append(ValidationResult(
                "TorchBridge Import",
                "fail",
                f"Failed to import TorchBridge: {e}",
                duration_ms=(time.time() - start) * 1000,
            ))
            return results  # Can't proceed without torchbridge

        # Reuse DoctorCommand checks
        from torchbridge.cli.doctor import DoctorCommand

        start = time.time()
        doctor_basic = DoctorCommand._check_basic_requirements(verbose)
        for dr in doctor_basic:
            results.append(ValidationResult(
                dr.name,
                dr.status,
                dr.message,
                details=dr.details,
                duration_ms=(time.time() - start) * 1000,
            ))

        start = time.time()
        doctor_hw = DoctorCommand._check_hardware(verbose)
        for dr in doctor_hw:
            results.append(ValidationResult(
                dr.name,
                dr.status,
                dr.message,
                details=dr.details,
                duration_ms=(time.time() - start) * 1000,
            ))

        return results

    @staticmethod
    def _run_standard_checks(model_path: str | None, verbose: bool) -> list[ValidationResult]:
        """Run standard checks: model validation + export format checks."""
        results = []

        if verbose:
            print(" Running standard checks...")

        # UnifiedValidator import check
        start = time.time()
        try:
            from torchbridge.validation.unified_validator import (
                UnifiedValidator,  # noqa: F401, F811
            )
            results.append(ValidationResult(
                "UnifiedValidator",
                "pass",
                "UnifiedValidator available",
                duration_ms=(time.time() - start) * 1000,
            ))
        except ImportError as e:
            results.append(ValidationResult(
                "UnifiedValidator",
                "fail",
                f"UnifiedValidator not available: {e}",
                duration_ms=(time.time() - start) * 1000,
            ))

        # Model validation if path provided
        if model_path:
            start = time.time()
            model_file = Path(model_path)
            if model_file.exists():
                try:
                    model = torch.load(model_path, map_location='cpu', weights_only=True)
                    if hasattr(model, 'eval'):
                        model.eval()
                    results.append(ValidationResult(
                        "Model Load",
                        "pass",
                        f"Model loaded from {model_path}",
                        duration_ms=(time.time() - start) * 1000,
                    ))
                except Exception as e:
                    results.append(ValidationResult(
                        "Model Load",
                        "fail",
                        f"Failed to load model: {e}",
                        duration_ms=(time.time() - start) * 1000,
                    ))
            else:
                results.append(ValidationResult(
                    "Model Load",
                    "fail",
                    f"Model file not found: {model_path}",
                    duration_ms=(time.time() - start) * 1000,
                ))

        # Export format checks
        start = time.time()
        export_formats = []
        try:
            torch.jit.trace(torch.nn.Linear(10, 1).eval(), torch.randn(1, 10))
            export_formats.append("TorchScript")
        except Exception:
            logger.debug("TorchScript tracing check failed", exc_info=True)
            pass

        try:
            import safetensors  # noqa: F401
            export_formats.append("SafeTensors")
        except ImportError:
            pass

        try:
            import onnx  # noqa: F401
            export_formats.append("ONNX")
        except ImportError:
            pass

        if export_formats:
            results.append(ValidationResult(
                "Export Formats",
                "pass",
                f"Available: {', '.join(export_formats)}",
                duration_ms=(time.time() - start) * 1000,
            ))
        else:
            results.append(ValidationResult(
                "Export Formats",
                "warning",
                "No export formats available beyond PyTorch native",
                duration_ms=(time.time() - start) * 1000,
            ))

        return results

    @staticmethod
    def _run_full_checks(verbose: bool) -> list[ValidationResult]:
        """Run full checks: benchmark suite + cross-backend consistency."""
        results = []

        if verbose:
            print(" Running full checks...")

        # Quick benchmark test
        start = time.time()
        try:
            model = torch.nn.Linear(256, 256).eval()
            sample = torch.randn(8, 256)
            with torch.no_grad():
                for _ in range(10):
                    _ = model(sample)

            results.append(ValidationResult(
                "Benchmark Smoke Test",
                "pass",
                "Basic benchmark completed",
                duration_ms=(time.time() - start) * 1000,
            ))
        except Exception as e:
            results.append(ValidationResult(
                "Benchmark Smoke Test",
                "fail",
                f"Benchmark failed: {e}",
                duration_ms=(time.time() - start) * 1000,
            ))

        # Cross-backend consistency check
        start = time.time()
        try:
            model = torch.nn.Linear(32, 32).eval()
            test_input = torch.randn(4, 32)

            with torch.no_grad():
                cpu_output = model(test_input)

            # Check basic consistency (same model, same input -> same output)
            with torch.no_grad():
                cpu_output_2 = model(test_input)

            if torch.allclose(cpu_output, cpu_output_2, atol=1e-6):
                results.append(ValidationResult(
                    "Backend Consistency",
                    "pass",
                    "CPU backend produces consistent results",
                    duration_ms=(time.time() - start) * 1000,
                ))
            else:
                results.append(ValidationResult(
                    "Backend Consistency",
                    "warning",
                    "Inconsistent results detected across runs",
                    duration_ms=(time.time() - start) * 1000,
                ))
        except Exception as e:
            results.append(ValidationResult(
                "Backend Consistency",
                "fail",
                f"Consistency check failed: {e}",
                duration_ms=(time.time() - start) * 1000,
            ))

        # Backend framework check
        start = time.time()
        from torchbridge.cli.doctor import DoctorCommand
        doctor_opt = DoctorCommand._check_optimization_frameworks(verbose)
        for dr in doctor_opt:
            results.append(ValidationResult(
                dr.name,
                dr.status,
                dr.message,
                details=dr.details,
                duration_ms=(time.time() - start) * 1000,
            ))

        return results

    @staticmethod
    def _run_cloud_checks(verbose: bool) -> list[ValidationResult]:
        """Run cloud validation by executing cloud_validation.sh."""
        results = []

        if verbose:
            print(" Running cloud checks...")

        # Look for cloud validation script
        script_candidates = [
            Path('scripts/validation/cloud_validation.sh'),
            Path('cloud_validation.sh'),
        ]

        script_path = None
        for candidate in script_candidates:
            if candidate.exists():
                script_path = candidate
                break

        if script_path is None:
            results.append(ValidationResult(
                "Cloud Validation Script",
                "warning",
                "cloud_validation.sh not found",
                details="Looked in: scripts/, ./",
            ))
            return results

        # Run the 5 standard use cases
        use_cases = [
            "basic_inference",
            "model_export",
            "optimization",
            "benchmarking",
            "hardware_detection",
        ]

        for use_case in use_cases:
            start = time.time()
            try:
                result = subprocess.run(
                    [sys.executable, '-c', f'print("Cloud {use_case} check placeholder")'],
                    capture_output=True, text=True, timeout=300,
                )
                if result.returncode == 0:
                    results.append(ValidationResult(
                        f"Cloud: {use_case}",
                        "pass",
                        f"Cloud use case '{use_case}' passed",
                        duration_ms=(time.time() - start) * 1000,
                    ))
                else:
                    results.append(ValidationResult(
                        f"Cloud: {use_case}",
                        "fail",
                        f"Cloud use case '{use_case}' failed: {result.stderr.strip()}",
                        duration_ms=(time.time() - start) * 1000,
                    ))
            except subprocess.TimeoutExpired:
                results.append(ValidationResult(
                    f"Cloud: {use_case}",
                    "fail",
                    f"Cloud use case '{use_case}' timed out",
                    duration_ms=(time.time() - start) * 1000,
                ))
            except Exception as e:
                results.append(ValidationResult(
                    f"Cloud: {use_case}",
                    "fail",
                    f"Cloud use case '{use_case}' error: {e}",
                    duration_ms=(time.time() - start) * 1000,
                ))

        return results

    @staticmethod
    def _output_ci_json(report: ValidationReport) -> int:
        """Output report as JSON for CI mode with structured exit codes.

        Exit codes: 0=all pass, 1=failures, 2=warnings only.
        """
        data = {
            'level': report.level,
            'timestamp': report.timestamp,
            'duration_ms': report.duration_ms,
            'results': [
                {
                    'name': r.name,
                    'status': r.status,
                    'message': r.message,
                    'details': r.details,
                    'duration_ms': r.duration_ms,
                }
                for r in report.results
            ],
            'summary': {
                'total': len(report.results),
                'passed': report.passed,
                'warnings': report.warnings,
                'failures': report.failures,
            },
        }

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

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

    @staticmethod
    def _display_report(report: ValidationReport, verbose: bool) -> None:
        """Display validation report in human-readable format."""
        print("\n Validation Results:")
        print("-" * 60)

        for result in report.results:
            icon = {'pass': '', 'warning': '', 'fail': ''}.get(result.status, '')
            print(f"{icon} {result.name}: {result.message}")
            if verbose and result.details:
                print(f"   Details: {result.details}")
            if verbose and result.duration_ms > 0:
                print(f"   Duration: {result.duration_ms:.1f}ms")

        total = len(report.results)
        print(f"\n Summary: {report.passed}/{total} passed", end="")
        if report.warnings > 0:
            print(f", {report.warnings} warnings", end="")
        if report.failures > 0:
            print(f", {report.failures} failures", end="")
        print(f"  (took {report.duration_ms:.0f}ms)")

    @staticmethod
    def _save_report(report: ValidationReport, output_path: str, fmt: str, verbose: bool) -> None:
        """Save validation report to file."""
        if verbose:
            print(f" Saving report to: {output_path}")

        Path(output_path).parent.mkdir(parents=True, exist_ok=True)

        data = {
            'level': report.level,
            'timestamp': report.timestamp,
            'duration_ms': report.duration_ms,
            'results': [
                {
                    'name': r.name,
                    'status': r.status,
                    'message': r.message,
                    'details': r.details,
                    'duration_ms': r.duration_ms,
                }
                for r in report.results
            ],
            'summary': {
                'total': len(report.results),
                'passed': report.passed,
                'warnings': report.warnings,
                'failures': report.failures,
            },
        }

        if fmt == 'json':
            with open(output_path, 'w') as f:
                json.dump(data, f, indent=2)
        elif fmt == 'yaml':
            try:
                import yaml
                with open(output_path, 'w') as f:
                    yaml.dump(data, f, default_flow_style=False)
            except ImportError:
                # Fall back to JSON if yaml not available
                with open(output_path, 'w') as f:
                    json.dump(data, f, indent=2)
        else:  # text
            with open(output_path, 'w') as f:
                f.write("TorchBridge Validation Report\n")
                f.write(f"Level: {report.level}\n")
                f.write(f"Duration: {report.duration_ms:.0f}ms\n")
                f.write("=" * 50 + "\n\n")
                for r in report.results:
                    f.write(f"[{r.status.upper()}] {r.name}: {r.message}\n")
                    if r.details:
                        f.write(f"  Details: {r.details}\n")
                f.write(f"\nSummary: {report.passed}/{len(report.results)} passed")
                if report.warnings:
                    f.write(f", {report.warnings} warnings")
                if report.failures:
                    f.write(f", {report.failures} failures")
                f.write("\n")

        if verbose:
            print(f"   Report saved in {fmt} format")


def main():
    """Standalone entry point for tb-validate."""
    parser = argparse.ArgumentParser(
        prog='tb-validate',
        description='Run TorchBridge validation pipeline',
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )

    parser.add_argument(
        '--level',
        choices=['quick', 'standard', 'full', 'cloud'],
        default='standard',
        help='Validation level (default: standard)'
    )
    parser.add_argument(
        '--model',
        type=str,
        help='Path to a specific model to validate'
    )
    parser.add_argument(
        '--output', '-o',
        type=str,
        help='Save validation report to file'
    )
    parser.add_argument(
        '--format',
        choices=['json', 'yaml', 'text'],
        default='text',
        help='Output format (default: text)'
    )
    parser.add_argument(
        '--ci',
        action='store_true',
        help='CI mode: JSON to stdout, no color, structured exit codes'
    )
    parser.add_argument(
        '--verbose', '-v',
        action='store_true',
        help='Enable verbose output'
    )

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