"""
Invariant Validator - Foundry-based Invariant Testing
======================================================

Validates synthesized invariants by generating and running Foundry tests.

Features:
- Generate Foundry test files from synthesized invariants
- Run invariant tests using Foundry's fuzzing
- Collect pass/fail results for ML feedback
- Provide validation statistics and reports

Author: Fernando Boiero <fboiero@frvm.utn.edu.ar>
Date: January 2026
Version: 1.0.0
"""

import hashlib
import json
import logging
import re
import shutil
import tempfile
from dataclasses import dataclass, field
from datetime import datetime
from pathlib import Path
from typing import Any, Dict, List, Optional, Union

from src.adapters.invariant_synthesizer import (
    InvariantCategory,
    InvariantFormat,
    SynthesizedInvariant,
)
from src.poc.validators.foundry_runner import (
    FoundryResult,
    FoundryRunner,
    TestResult,
    TestStatus,
)

logger = logging.getLogger(__name__)


@dataclass
class InvariantTestResult:
    """Result of validating a single invariant."""

    invariant_name: str
    category: InvariantCategory
    passed: bool
    test_status: TestStatus
    gas_used: Optional[int] = None
    execution_time_ms: float = 0
    error_message: Optional[str] = None
    counterexample: Optional[str] = None
    fuzzing_runs: int = 0
    traces: Optional[str] = None

    def to_dict(self) -> Dict[str, Any]:
        """Convert to dictionary."""
        return {
            "invariant_name": self.invariant_name,
            "category": self.category.value,
            "passed": self.passed,
            "test_status": self.test_status.value,
            "gas_used": self.gas_used,
            "execution_time_ms": self.execution_time_ms,
            "error_message": self.error_message,
            "counterexample": self.counterexample,
            "fuzzing_runs": self.fuzzing_runs,
        }


@dataclass
class ValidationReport:
    """Complete validation report for a set of invariants."""

    contract_name: str
    contract_hash: str
    total_invariants: int
    passed: int
    failed: int
    skipped: int
    results: List[InvariantTestResult]
    total_execution_time_ms: float
    total_gas_used: int
    timestamp: datetime = field(default_factory=datetime.now)
    foundry_version: Optional[str] = None
    errors: List[str] = field(default_factory=list)

    @property
    def pass_rate(self) -> float:
        """Calculate pass rate."""
        if self.total_invariants == 0:
            return 0.0
        return self.passed / self.total_invariants

    def to_dict(self) -> Dict[str, Any]:
        """Convert to dictionary."""
        return {
            "contract_name": self.contract_name,
            "contract_hash": self.contract_hash,
            "total_invariants": self.total_invariants,
            "passed": self.passed,
            "failed": self.failed,
            "skipped": self.skipped,
            "pass_rate": self.pass_rate,
            "results": [r.to_dict() for r in self.results],
            "total_execution_time_ms": self.total_execution_time_ms,
            "total_gas_used": self.total_gas_used,
            "timestamp": self.timestamp.isoformat(),
            "foundry_version": self.foundry_version,
            "errors": self.errors,
        }


class InvariantTestGenerator:
    """
    Generates Foundry test files for invariants.

    Creates compilable Solidity test contracts that verify
    synthesized invariants using Foundry's fuzzing engine.
    """

    # Template for invariant test contract
    TEST_CONTRACT_TEMPLATE = '''// SPDX-License-Identifier: MIT
pragma solidity ^0.8.19;

import "forge-std/Test.sol";
import "{contract_import}";

/// @title Invariant Tests for {contract_name}
/// @notice Auto-generated by MIESC Invariant Validator
/// @dev Run with: forge test --match-contract {test_contract_name}
contract {test_contract_name} is Test {{
    {contract_name} target;

    function setUp() public {{
        {setup_code}
    }}

{invariant_functions}
}}
'''

    # Template for individual invariant function
    INVARIANT_FUNCTION_TEMPLATE = '''
    /// @notice {description}
    /// @dev Category: {category}, Importance: {importance}
    function invariant_{name}() public view {{
        {assertion_code}
    }}
'''

    # Template for fuzz test function
    FUZZ_TEST_TEMPLATE = '''
    /// @notice Fuzz test for {name}
    /// @dev {description}
    function testFuzz_{name}({fuzz_params}) public {{
        {preconditions}
        {action}
        {postconditions}
    }}
'''

    def __init__(self, output_dir: Optional[Path] = None):
        """
        Initialize test generator.

        Args:
            output_dir: Directory for generated test files
        """
        self.output_dir = output_dir or Path(tempfile.gettempdir()) / "miesc_invariant_tests"
        self.output_dir.mkdir(parents=True, exist_ok=True)
        logger.info(f"InvariantTestGenerator initialized, output: {self.output_dir}")

    def generate_test_file(
        self,
        invariants: List[SynthesizedInvariant],
        contract_name: str,
        contract_path: str,
        setup_code: str = "",
    ) -> Path:
        """
        Generate Foundry test file for invariants.

        Args:
            invariants: List of invariants to test
            contract_name: Name of the target contract
            contract_path: Import path for the contract
            setup_code: Code for setUp function

        Returns:
            Path to generated test file
        """
        # Generate invariant functions
        invariant_functions = []

        for inv in invariants:
            if not inv.foundry_test and not inv.solidity_assertion:
                logger.warning(f"Skipping {inv.name}: no Foundry test or Solidity assertion")
                continue

            # Use foundry_test if available, otherwise generate from assertion
            if inv.foundry_test:
                assertion_code = inv.foundry_test
            elif inv.solidity_assertion:
                assertion_code = self._wrap_assertion(inv.solidity_assertion)
            else:
                continue

            func = self.INVARIANT_FUNCTION_TEMPLATE.format(
                name=self._sanitize_name(inv.name),
                description=inv.description.replace("\n", " "),
                category=inv.category.value,
                importance=inv.importance,
                assertion_code=assertion_code,
            )
            invariant_functions.append(func)

        if not invariant_functions:
            raise ValueError("No valid invariants to generate tests for")

        # Default setup if not provided
        if not setup_code:
            setup_code = f"target = new {contract_name}();"

        # Generate test contract
        test_contract_name = f"{contract_name}InvariantTest"
        contract_code = self.TEST_CONTRACT_TEMPLATE.format(
            contract_import=contract_path,
            contract_name=contract_name,
            test_contract_name=test_contract_name,
            setup_code=setup_code,
            invariant_functions="\n".join(invariant_functions),
        )

        # Write to file
        test_file = self.output_dir / f"{test_contract_name}.t.sol"
        test_file.write_text(contract_code)

        logger.info(f"Generated test file: {test_file} ({len(invariant_functions)} invariants)")
        return test_file

    def generate_individual_tests(
        self,
        invariants: List[SynthesizedInvariant],
        contract_name: str,
        contract_path: str,
        setup_code: str = "",
    ) -> List[Path]:
        """
        Generate individual test files for each invariant.

        Args:
            invariants: List of invariants
            contract_name: Target contract name
            contract_path: Import path
            setup_code: Setup code

        Returns:
            List of paths to generated test files
        """
        test_files = []

        for inv in invariants:
            if not inv.foundry_test and not inv.solidity_assertion:
                continue

            try:
                test_file = self.generate_test_file(
                    invariants=[inv],
                    contract_name=contract_name,
                    contract_path=contract_path,
                    setup_code=setup_code,
                )
                # Rename to include invariant name
                new_name = f"{contract_name}_{self._sanitize_name(inv.name)}.t.sol"
                new_path = self.output_dir / new_name
                test_file.rename(new_path)
                test_files.append(new_path)
            except Exception as e:
                logger.error(f"Error generating test for {inv.name}: {e}")

        return test_files

    def _wrap_assertion(self, assertion: str) -> str:
        """Wrap a Solidity assertion in proper format."""
        assertion = assertion.strip()

        # If it's already a full statement, return as-is
        if assertion.endswith(";"):
            return assertion

        # If it's a condition, wrap in assertTrue
        if not assertion.startswith("assert") and not assertion.startswith("require"):
            return f"assertTrue({assertion}, \"Invariant violated\");"

        return assertion + ";"

    def _sanitize_name(self, name: str) -> str:
        """Sanitize name for Solidity function."""
        # Remove special characters, keep alphanumeric and underscore
        sanitized = re.sub(r"[^a-zA-Z0-9_]", "_", name)
        # Ensure it starts with a letter
        if sanitized and sanitized[0].isdigit():
            sanitized = "inv_" + sanitized
        return sanitized or "unnamed"

    def cleanup(self) -> None:
        """Clean up generated test files."""
        if self.output_dir.exists():
            shutil.rmtree(self.output_dir)
            logger.info(f"Cleaned up: {self.output_dir}")


class InvariantValidator:
    """
    Validates synthesized invariants using Foundry.

    Orchestrates test generation, execution, and result collection
    for ML feedback loop.
    """

    def __init__(
        self,
        project_dir: Optional[Union[str, Path]] = None,
        fork_url: Optional[str] = None,
        fuzzing_runs: int = 256,
        fuzzing_depth: int = 100,
        timeout: int = 600,
    ):
        """
        Initialize invariant validator.

        Args:
            project_dir: Foundry project directory
            fork_url: RPC URL for mainnet forking
            fuzzing_runs: Number of fuzzing runs per invariant
            fuzzing_depth: Maximum call depth for fuzzing
            timeout: Test execution timeout in seconds
        """
        self.project_dir = Path(project_dir) if project_dir else Path.cwd()
        self.fork_url = fork_url
        self.fuzzing_runs = fuzzing_runs
        self.fuzzing_depth = fuzzing_depth
        self.timeout = timeout

        # Initialize components
        self.test_generator = InvariantTestGenerator(
            output_dir=self.project_dir / "test" / "invariants"
        )
        self.runner = FoundryRunner(
            project_dir=self.project_dir,
            fork_url=fork_url,
            verbosity=3,
            gas_report=True,
        )

        # Results storage
        self.validation_history: List[ValidationReport] = []

        logger.info(
            f"InvariantValidator initialized "
            f"(fuzzing_runs={fuzzing_runs}, depth={fuzzing_depth})"
        )

    def validate(
        self,
        invariants: List[SynthesizedInvariant],
        contract_name: str,
        contract_path: str,
        setup_code: str = "",
        run_individual: bool = False,
    ) -> ValidationReport:
        """
        Validate a list of invariants.

        Args:
            invariants: Invariants to validate
            contract_name: Target contract name
            contract_path: Import path for contract
            setup_code: Custom setUp code
            run_individual: Run each invariant separately

        Returns:
            ValidationReport with results
        """
        # Calculate contract hash
        contract_hash = hashlib.sha256(contract_name.encode()).hexdigest()[:16]

        # Filter invariants with testable assertions
        testable = [
            inv for inv in invariants
            if inv.foundry_test or inv.solidity_assertion
        ]

        if not testable:
            logger.warning("No testable invariants found")
            return ValidationReport(
                contract_name=contract_name,
                contract_hash=contract_hash,
                total_invariants=0,
                passed=0,
                failed=0,
                skipped=len(invariants),
                results=[],
                total_execution_time_ms=0,
                total_gas_used=0,
                errors=["No testable invariants (missing foundry_test or solidity_assertion)"],
            )

        try:
            if run_individual:
                return self._validate_individual(
                    testable, contract_name, contract_path, contract_hash, setup_code
                )
            else:
                return self._validate_batch(
                    testable, contract_name, contract_path, contract_hash, setup_code
                )
        except Exception as e:
            logger.error(f"Validation failed: {e}")
            return ValidationReport(
                contract_name=contract_name,
                contract_hash=contract_hash,
                total_invariants=len(testable),
                passed=0,
                failed=0,
                skipped=len(testable),
                results=[],
                total_execution_time_ms=0,
                total_gas_used=0,
                errors=[str(e)],
            )

    def _validate_batch(
        self,
        invariants: List[SynthesizedInvariant],
        contract_name: str,
        contract_path: str,
        contract_hash: str,
        setup_code: str,
    ) -> ValidationReport:
        """Run all invariants in a single test file."""
        # Generate test file
        test_file = self.test_generator.generate_test_file(
            invariants=invariants,
            contract_name=contract_name,
            contract_path=contract_path,
            setup_code=setup_code,
        )

        # Run tests
        result = self.runner.run_test(
            test_path=test_file,
            timeout=self.timeout,
        )

        # Map results to invariants
        test_results = self._map_results_to_invariants(invariants, result)

        # Calculate statistics
        passed = sum(1 for r in test_results if r.passed)
        failed = sum(1 for r in test_results if not r.passed and r.test_status != TestStatus.SKIPPED)
        skipped = sum(1 for r in test_results if r.test_status == TestStatus.SKIPPED)

        report = ValidationReport(
            contract_name=contract_name,
            contract_hash=contract_hash,
            total_invariants=len(invariants),
            passed=passed,
            failed=failed,
            skipped=skipped,
            results=test_results,
            total_execution_time_ms=result.execution_time_ms,
            total_gas_used=result.total_gas,
            foundry_version=result.forge_version,
            errors=[result.error] if result.error else [],
        )

        self.validation_history.append(report)
        return report

    def _validate_individual(
        self,
        invariants: List[SynthesizedInvariant],
        contract_name: str,
        contract_path: str,
        contract_hash: str,
        setup_code: str,
    ) -> ValidationReport:
        """Run each invariant in a separate test file."""
        test_results = []
        total_time = 0.0
        total_gas = 0

        for inv in invariants:
            try:
                # Generate individual test
                test_files = self.test_generator.generate_individual_tests(
                    invariants=[inv],
                    contract_name=contract_name,
                    contract_path=contract_path,
                    setup_code=setup_code,
                )

                if not test_files:
                    test_results.append(InvariantTestResult(
                        invariant_name=inv.name,
                        category=inv.category,
                        passed=False,
                        test_status=TestStatus.SKIPPED,
                        error_message="Could not generate test",
                    ))
                    continue

                # Run test
                result = self.runner.run_test(
                    test_path=test_files[0],
                    timeout=self.timeout // len(invariants),
                )

                total_time += result.execution_time_ms
                total_gas += result.total_gas

                # Parse result
                if result.tests:
                    test = result.tests[0]
                    test_results.append(InvariantTestResult(
                        invariant_name=inv.name,
                        category=inv.category,
                        passed=test.status == TestStatus.PASSED,
                        test_status=test.status,
                        gas_used=test.gas_used,
                        execution_time_ms=test.duration_ms,
                        error_message=test.error_message,
                        fuzzing_runs=self.fuzzing_runs,
                    ))
                else:
                    test_results.append(InvariantTestResult(
                        invariant_name=inv.name,
                        category=inv.category,
                        passed=result.success,
                        test_status=TestStatus.PASSED if result.success else TestStatus.ERROR,
                        error_message=result.error,
                    ))

            except Exception as e:
                logger.error(f"Error validating {inv.name}: {e}")
                test_results.append(InvariantTestResult(
                    invariant_name=inv.name,
                    category=inv.category,
                    passed=False,
                    test_status=TestStatus.ERROR,
                    error_message=str(e),
                ))

        # Calculate statistics
        passed = sum(1 for r in test_results if r.passed)
        failed = sum(1 for r in test_results if not r.passed and r.test_status != TestStatus.SKIPPED)
        skipped = sum(1 for r in test_results if r.test_status == TestStatus.SKIPPED)

        report = ValidationReport(
            contract_name=contract_name,
            contract_hash=contract_hash,
            total_invariants=len(invariants),
            passed=passed,
            failed=failed,
            skipped=skipped,
            results=test_results,
            total_execution_time_ms=total_time,
            total_gas_used=total_gas,
        )

        self.validation_history.append(report)
        return report

    def _map_results_to_invariants(
        self,
        invariants: List[SynthesizedInvariant],
        foundry_result: FoundryResult,
    ) -> List[InvariantTestResult]:
        """Map Foundry test results to invariants."""
        results = []

        # Create lookup by test name
        test_lookup: Dict[str, TestResult] = {}
        for test in foundry_result.tests:
            # Extract function name from full test name
            name = test.name.split("::")[-1] if "::" in test.name else test.name
            test_lookup[name.lower()] = test

        for inv in invariants:
            sanitized = self.test_generator._sanitize_name(inv.name)
            func_name = f"invariant_{sanitized}".lower()

            if func_name in test_lookup:
                test = test_lookup[func_name]
                results.append(InvariantTestResult(
                    invariant_name=inv.name,
                    category=inv.category,
                    passed=test.status == TestStatus.PASSED,
                    test_status=test.status,
                    gas_used=test.gas_used,
                    execution_time_ms=test.duration_ms,
                    error_message=test.error_message,
                    traces=test.traces,
                    fuzzing_runs=self.fuzzing_runs,
                ))
            else:
                # Test not found - might have failed to compile
                results.append(InvariantTestResult(
                    invariant_name=inv.name,
                    category=inv.category,
                    passed=False,
                    test_status=TestStatus.SKIPPED,
                    error_message="Test not found in results",
                ))

        return results

    def get_feedback_for_ml(
        self,
        report: ValidationReport,
    ) -> Dict[str, Any]:
        """
        Generate feedback data for ML model training.

        Args:
            report: Validation report

        Returns:
            Dictionary with feedback data for ML
        """
        feedback = {
            "contract_hash": report.contract_hash,
            "timestamp": report.timestamp.isoformat(),
            "pass_rate": report.pass_rate,
            "invariants": [],
        }

        for result in report.results:
            feedback["invariants"].append({
                "name": result.invariant_name,
                "category": result.category.value,
                "passed": result.passed,
                "gas_used": result.gas_used,
                "has_counterexample": result.counterexample is not None,
            })

        return feedback

    def get_validation_stats(self) -> Dict[str, Any]:
        """Get statistics from all validations."""
        if not self.validation_history:
            return {"total_validations": 0}

        total_invariants = sum(r.total_invariants for r in self.validation_history)
        total_passed = sum(r.passed for r in self.validation_history)
        total_failed = sum(r.failed for r in self.validation_history)

        # Category breakdown
        category_stats: Dict[str, Dict[str, int]] = {}
        for report in self.validation_history:
            for result in report.results:
                cat = result.category.value
                if cat not in category_stats:
                    category_stats[cat] = {"passed": 0, "failed": 0}
                if result.passed:
                    category_stats[cat]["passed"] += 1
                else:
                    category_stats[cat]["failed"] += 1

        return {
            "total_validations": len(self.validation_history),
            "total_invariants": total_invariants,
            "total_passed": total_passed,
            "total_failed": total_failed,
            "overall_pass_rate": total_passed / total_invariants if total_invariants > 0 else 0,
            "category_breakdown": category_stats,
            "avg_execution_time_ms": sum(r.total_execution_time_ms for r in self.validation_history)
            / len(self.validation_history),
        }

    def save_report(
        self,
        report: ValidationReport,
        output_path: Union[str, Path],
    ) -> None:
        """Save validation report to JSON file."""
        output_path = Path(output_path)
        output_path.parent.mkdir(parents=True, exist_ok=True)

        with open(output_path, "w") as f:
            json.dump(report.to_dict(), f, indent=2, default=str)

        logger.info(f"Saved report to: {output_path}")

    def cleanup(self) -> None:
        """Clean up generated test files."""
        self.test_generator.cleanup()


# ============================================================================
# Convenience Functions
# ============================================================================


def validate_invariants(
    invariants: List[SynthesizedInvariant],
    contract_name: str,
    contract_path: str,
    project_dir: Optional[str] = None,
    setup_code: str = "",
) -> ValidationReport:
    """
    Validate invariants using Foundry.

    Args:
        invariants: List of synthesized invariants
        contract_name: Name of target contract
        contract_path: Import path for contract
        project_dir: Foundry project directory
        setup_code: Custom setUp code

    Returns:
        ValidationReport with results
    """
    validator = InvariantValidator(project_dir=project_dir)

    try:
        return validator.validate(
            invariants=invariants,
            contract_name=contract_name,
            contract_path=contract_path,
            setup_code=setup_code,
        )
    finally:
        validator.cleanup()


def quick_validate(
    invariant: SynthesizedInvariant,
    contract_code: str,
    contract_name: str = "TestContract",
) -> bool:
    """
    Quickly validate a single invariant.

    Args:
        invariant: Invariant to validate
        contract_code: Source code of contract
        contract_name: Name of contract

    Returns:
        True if invariant passes
    """
    # Create temporary project
    with tempfile.TemporaryDirectory() as tmpdir:
        project_dir = Path(tmpdir)

        # Create minimal Foundry structure
        (project_dir / "src").mkdir()
        (project_dir / "test").mkdir()
        (project_dir / "lib").mkdir()

        # Write contract
        contract_file = project_dir / "src" / f"{contract_name}.sol"
        contract_file.write_text(contract_code)

        # Create foundry.toml
        foundry_toml = project_dir / "foundry.toml"
        foundry_toml.write_text("""
[profile.default]
src = "src"
out = "out"
libs = ["lib"]
""")

        # Validate
        validator = InvariantValidator(project_dir=project_dir)

        try:
            report = validator.validate(
                invariants=[invariant],
                contract_name=contract_name,
                contract_path=f"src/{contract_name}.sol",
            )
            return report.passed > 0
        except Exception as e:
            logger.error(f"Quick validation failed: {e}")
            return False
        finally:
            validator.cleanup()


# ============================================================================
# Export
# ============================================================================

__all__ = [
    "InvariantValidator",
    "InvariantTestGenerator",
    "InvariantTestResult",
    "ValidationReport",
    "validate_invariants",
    "quick_validate",
]
