"""Java test result comparison.

This module provides functionality to compare test results between
original and optimized Java code using the codeflash-runtime Comparator.
"""

from __future__ import annotations

import json
import logging
import math
import os
import platform
import shutil
import subprocess
from functools import lru_cache
from pathlib import Path
from typing import TYPE_CHECKING

from codeflash.languages.java.build_tools import CODEFLASH_RUNTIME_JAR_NAME, CODEFLASH_RUNTIME_VERSION

if TYPE_CHECKING:
    from codeflash.models.models import TestDiff

_IS_DARWIN = platform.system() == "Darwin"

logger = logging.getLogger(__name__)


def _find_comparator_jar(project_root: Path | None = None) -> Path | None:
    """Find the codeflash-runtime JAR with the Comparator class.

    Args:
        project_root: Project root directory.

    Returns:
        Path to codeflash-runtime JAR if found, None otherwise.

    """
    search_dirs = []
    if project_root:
        search_dirs.append(project_root)
    search_dirs.append(Path.cwd())

    # Search for the JAR in common locations
    for base_dir in search_dirs:
        # Check in target directory (after Maven install)
        for jar_path in [
            base_dir / "target" / "dependency" / CODEFLASH_RUNTIME_JAR_NAME,
            base_dir / "target" / CODEFLASH_RUNTIME_JAR_NAME,
            base_dir / "lib" / CODEFLASH_RUNTIME_JAR_NAME,
            base_dir / ".codeflash" / CODEFLASH_RUNTIME_JAR_NAME,
        ]:
            if jar_path.exists():
                return jar_path

        # Check local Maven repository
        m2_jar = (
            Path.home()
            / ".m2"
            / "repository"
            / "com"
            / "codeflash"
            / "codeflash-runtime"
            / CODEFLASH_RUNTIME_VERSION
            / CODEFLASH_RUNTIME_JAR_NAME
        )
        if m2_jar.exists():
            return m2_jar

    # Check bundled JAR in package resources
    resources_jar = Path(__file__).parent / "resources" / CODEFLASH_RUNTIME_JAR_NAME
    if resources_jar.exists():
        return resources_jar

    return None


@lru_cache(maxsize=1)
def _find_java_executable() -> str | None:
    """Find the Java executable.

    Returns:
        Path to java executable, or None if not found.

    """
    # Check JAVA_HOME
    java_home = os.environ.get("JAVA_HOME")
    if java_home:
        java_path = Path(java_home) / "bin" / "java"
        if java_path.exists():
            return str(java_path)

    # On macOS, try to get JAVA_HOME from the system helper or Maven
    if _IS_DARWIN:
        # Try to extract Java home from Maven (which always finds it)
        try:
            result = subprocess.run(["mvn", "--version"], capture_output=True, text=True, timeout=10, check=False)
            if result.returncode == 0:
                for line in result.stdout.split("\n"):
                    if "runtime:" in line:
                        runtime_path = line.split("runtime:")[-1].strip()
                        java_path = Path(runtime_path) / "bin" / "java"
                        if java_path.exists():
                            return str(java_path)
        except (subprocess.TimeoutExpired, FileNotFoundError):
            pass

        # Check common Homebrew locations
        for homebrew_java in [
            "/opt/homebrew/opt/openjdk/bin/java",
            "/opt/homebrew/opt/openjdk@25/bin/java",
            "/opt/homebrew/opt/openjdk@21/bin/java",
            "/opt/homebrew/opt/openjdk@17/bin/java",
            "/usr/local/opt/openjdk/bin/java",
        ]:
            if Path(homebrew_java).exists():
                return homebrew_java

    # Check PATH (on macOS, /usr/bin/java may be a stub that fails)
    java_which = shutil.which("java")
    if java_which:
        # Verify it's a real Java, not a macOS stub
        try:
            result = subprocess.run([java_which, "--version"], capture_output=True, text=True, timeout=5, check=False)
            if result.returncode == 0:
                return java_which
        except (subprocess.TimeoutExpired, FileNotFoundError):
            pass

    return None


def compare_test_results(
    original_sqlite_path: Path,
    candidate_sqlite_path: Path,
    comparator_jar: Path | None = None,
    project_root: Path | None = None,
    project_classpath: str | None = None,
) -> tuple[bool, list]:
    """Compare Java test results using the codeflash-runtime Comparator.

    This function calls the Java Comparator CLI that:
    1. Reads serialized behavior data from both SQLite databases
    2. Deserializes using Kryo
    3. Compares results using deep equality (handles Maps, Lists, arrays, etc.)
    4. Returns comparison results as JSON

    Args:
        original_sqlite_path: Path to SQLite database with original code results.
        candidate_sqlite_path: Path to SQLite database with candidate code results.
        comparator_jar: Optional path to the codeflash-runtime JAR.
        project_root: Project root directory.
        project_classpath: Full project classpath from the build tool. When provided,
            the Comparator JVM uses this classpath so Kryo can resolve project-specific
            classes during deserialization. Without it, only the runtime JAR is on the
            classpath and any project class causes ClassNotFoundException.

    Returns:
        Tuple of (all_equivalent, list of TestDiff objects).

    """
    # Import lazily to avoid circular imports
    from codeflash.models.models import TestDiff, TestDiffScope

    java_exe = _find_java_executable()
    if not java_exe:
        logger.error("Java not found. Please install Java to compare test results.")
        return False, []

    jar_path = comparator_jar or _find_comparator_jar(project_root)
    if not jar_path or not jar_path.exists():
        logger.error(
            "codeflash-runtime JAR not found. Please ensure the codeflash-runtime is installed in your project."
        )
        return False, []

    if not original_sqlite_path.exists():
        logger.error("Original SQLite database not found: %s", original_sqlite_path)
        return False, []

    if not candidate_sqlite_path.exists():
        logger.error("Candidate SQLite database not found: %s", candidate_sqlite_path)
        return False, []

    cwd = project_root or Path.cwd()

    # Build classpath: runtime JAR + project classpath (if available).
    # The project classpath is needed so Kryo can resolve project-specific classes
    # during deserialization. Without it, Class.forName() fails for any type not
    # bundled in the runtime JAR.
    cp_separator = ";" if platform.system() == "Windows" else ":"
    if project_classpath:
        full_cp = f"{jar_path}{cp_separator}{project_classpath}"
    else:
        full_cp = str(jar_path)

    try:
        result = subprocess.run(
            [
                java_exe,
                # Java 16+ module system: Kryo needs reflective access to internal JDK classes
                "--add-opens",
                "java.base/java.util=ALL-UNNAMED",
                "--add-opens",
                "java.base/java.lang=ALL-UNNAMED",
                "--add-opens",
                "java.base/java.lang.reflect=ALL-UNNAMED",
                "--add-opens",
                "java.base/java.io=ALL-UNNAMED",
                "--add-opens",
                "java.base/java.math=ALL-UNNAMED",
                "--add-opens",
                "java.base/java.net=ALL-UNNAMED",
                "--add-opens",
                "java.base/java.util.zip=ALL-UNNAMED",
                "-cp",
                full_cp,
                "com.codeflash.Comparator",
                str(original_sqlite_path),
                str(candidate_sqlite_path),
            ],
            check=False,
            capture_output=True,
            text=True,
            timeout=60,
            cwd=str(cwd),
        )

        # Parse the JSON output
        try:
            if not result.stdout or not result.stdout.strip():
                logger.error("Java comparator returned empty output")
                if result.stderr:
                    logger.error("stderr: %s", result.stderr)
                return False, []

            comparison = json.loads(result.stdout)

            if result.stderr:
                logger.debug("Java comparator stderr: %s", result.stderr.strip())
        except json.JSONDecodeError as e:
            logger.exception("Failed to parse Java comparator output: %s", e)
            logger.exception("stdout: %s", result.stdout[:500] if result.stdout else "(empty)")
            if result.stderr:
                logger.exception("stderr: %s", result.stderr[:500])
            return False, []

        # Check for errors in the JSON response
        if comparison.get("error"):
            logger.error("Java comparator error: %s", comparison["error"])
            return False, []

        # Check for unexpected exit codes
        if result.returncode not in {0, 1}:
            logger.error("Java comparator failed with exit code %s", result.returncode)
            if result.stderr:
                logger.error("stderr: %s", result.stderr)
            return False, []

        # Convert diffs to TestDiff objects
        test_diffs: list[TestDiff] = []
        for diff in comparison.get("diffs", []):
            scope_str = diff.get("scope", "return_value")
            scope = TestDiffScope.RETURN_VALUE
            if scope_str in {"exception", "missing"}:
                scope = TestDiffScope.DID_PASS

            # Build test identifier
            method_id = diff.get("methodId", "unknown")
            call_id = diff.get("callId", 0)
            test_src_code = f"// Method: {method_id}\n// Call ID: {call_id}"

            test_diffs.append(
                TestDiff(
                    scope=scope,
                    original_value=diff.get("originalValue"),
                    candidate_value=diff.get("candidateValue"),
                    test_src_code=test_src_code,
                    candidate_pytest_error=diff.get("candidateError"),
                    original_pass=scope_str != "exception",
                    candidate_pass=scope_str not in ("missing", "exception"),
                    original_pytest_error=diff.get("originalError"),
                )
            )

            logger.debug(
                "Java test diff:\n  Method: %s\n  Call ID: %s\n  Scope: %s\n  Original: %s\n  Candidate: %s",
                method_id,
                call_id,
                scope_str,
                str(diff.get("originalValue", "N/A"))[:100],
                str(diff.get("candidateValue", "N/A"))[:100],
            )

        equivalent = comparison.get("equivalent", False)
        actual_comparisons = comparison.get("actualComparisons", -1)
        skipped_placeholders = comparison.get("skippedPlaceholders", 0)
        skipped_deser_errors = comparison.get("skippedDeserializationErrors", 0)

        if actual_comparisons == 0:
            logger.warning(
                "Java comparison: no actual comparisons performed "
                "(total=%s, skipped_placeholders=%s, skipped_deser_errors=%s). "
                "Treating as NOT equivalent.",
                comparison.get("totalInvocations", 0),
                skipped_placeholders,
                skipped_deser_errors,
            )
            return False, []

        logger.info(
            "Java comparison: %s (%s invocations, %s compared, %s placeholder skips, %s deser skips, %s diffs)",
            "equivalent" if equivalent else "DIFFERENT",
            comparison.get("totalInvocations", 0),
            actual_comparisons,
            skipped_placeholders,
            skipped_deser_errors,
            len(test_diffs),
        )

        return equivalent, test_diffs

    except subprocess.TimeoutExpired:
        logger.exception("Java comparator timed out")
        return False, []
    except FileNotFoundError:
        logger.exception("Java not found. Please install Java to compare test results.")
        return False, []
    except Exception as e:
        logger.exception("Error running Java comparator: %s", e)
        return False, []


def values_equal(orig: str | None, cand: str | None) -> bool:
    """Compare two serialized values with numeric-aware equality.

    Handles boxing mismatches where Integer(0) and Long(0) serialize to different strings
    (e.g., "0" vs "0.0") but represent the same numeric value.
    """
    if orig == cand:
        return True
    if orig is None or cand is None:
        return False
    try:
        orig_num = float(orig)
        cand_num = float(cand)
        if math.isnan(orig_num) and math.isnan(cand_num):
            return True
        return orig_num == cand_num or math.isclose(orig_num, cand_num, rel_tol=1e-9)
    except (ValueError, TypeError):
        return False


def compare_invocations_directly(original_results: dict, candidate_results: dict) -> tuple[bool, list]:
    """Compare test invocations directly from Python dictionaries.

    This is a fallback when the Java comparator is not available.
    It performs basic equality comparison on serialized JSON values.

    Args:
        original_results: Dict mapping call_id to result data from original code.
        candidate_results: Dict mapping call_id to result data from candidate code.

    Returns:
        Tuple of (all_equivalent, list of TestDiff objects).

    """
    # Import lazily to avoid circular imports
    from codeflash.models.models import TestDiff, TestDiffScope

    test_diffs: list[TestDiff] = []

    # Get all call IDs
    all_call_ids = set(original_results.keys()) | set(candidate_results.keys())

    for call_id in all_call_ids:
        original = original_results.get(call_id)
        candidate = candidate_results.get(call_id)

        if original is None and candidate is not None:
            # Candidate has extra invocation
            test_diffs.append(
                TestDiff(
                    scope=TestDiffScope.DID_PASS,
                    original_value=None,
                    candidate_value=candidate.get("result_json"),
                    test_src_code=f"// Call ID: {call_id}",
                    candidate_pytest_error=None,
                    original_pass=True,
                    candidate_pass=True,
                    original_pytest_error=None,
                )
            )
        elif original is not None and candidate is None:
            # Candidate missing invocation
            test_diffs.append(
                TestDiff(
                    scope=TestDiffScope.DID_PASS,
                    original_value=original.get("result_json"),
                    candidate_value=None,
                    test_src_code=f"// Call ID: {call_id}",
                    candidate_pytest_error="Missing invocation in candidate",
                    original_pass=True,
                    candidate_pass=False,
                    original_pytest_error=None,
                )
            )
        elif original is not None and candidate is not None:
            # Both have invocations - compare results
            orig_result = original.get("result_json")
            cand_result = candidate.get("result_json")
            orig_error = original.get("error_json")
            cand_error = candidate.get("error_json")

            # Check for exception differences
            if orig_error != cand_error:
                test_diffs.append(
                    TestDiff(
                        scope=TestDiffScope.DID_PASS,
                        original_value=orig_error,
                        candidate_value=cand_error,
                        test_src_code=f"// Call ID: {call_id}",
                        candidate_pytest_error=cand_error,
                        original_pass=orig_error is None,
                        candidate_pass=cand_error is None,
                        original_pytest_error=orig_error,
                    )
                )
            elif not values_equal(orig_result, cand_result):
                # Results differ
                test_diffs.append(
                    TestDiff(
                        scope=TestDiffScope.RETURN_VALUE,
                        original_value=orig_result,
                        candidate_value=cand_result,
                        test_src_code=f"// Call ID: {call_id}",
                        candidate_pytest_error=None,
                        original_pass=True,
                        candidate_pass=True,
                        original_pytest_error=None,
                    )
                )

    equivalent = len(test_diffs) == 0

    logger.info(
        "Python comparison: %s (%s invocations, %s diffs)",
        "equivalent" if equivalent else "DIFFERENT",
        len(all_call_ids),
        len(test_diffs),
    )

    return equivalent, test_diffs
