#!/usr/bin/env python3
"""Prepare mutmut artifacts in throwaway processes before execution."""

from __future__ import annotations

import argparse
import hashlib
import inspect
import json
import os
from pathlib import Path
import sys
from typing import Any

PROJECT_ROOT = Path(__file__).resolve().parents[1]
if str(PROJECT_ROOT) not in sys.path:
    sys.path.insert(0, str(PROJECT_ROOT))

from scripts.check_mutation_targets import exact_python_targets
from scripts.check_mutation_targets import load_mutmut_config
from scripts.check_mutation_targets import SUPPORTED_MUTMUT_VERSION


MANIFEST_PATH = Path("mutants/.covered-lines-prepass.json")
STATS_PATH = Path("mutants/mutmut-stats.json")
HOOK_HASH_PATHS = (
    Path("scripts/prepare_mutation_coverage.py"),
    Path("scripts/run_prepared_mutmut.py"),
    Path("scripts/mutmut_coverage_sitecustomize/sitecustomize.py"),
)
META_KEYS = {
    "exit_code_by_key",
    "hash_by_function_name",
    "type_check_error_by_key",
    "durations_by_key",
    "estimated_durations_by_key",
}


def _sha256(path: Path) -> str:
    return hashlib.sha256(path.read_bytes()).hexdigest()


def _validated_relative_file(raw_path: str, *, label: str) -> Path:
    path = Path(raw_path)
    if path.is_absolute() or ".." in path.parts:
        raise RuntimeError(f"{label} must be a project-relative file: {raw_path}")
    if not path.is_file():
        raise RuntimeError(f"{label} is missing: {raw_path}")
    return path


def _hash_relative_files(raw_paths: list[str], *, label: str) -> dict[str, str]:
    return {
        str(path): _sha256(path)
        for path in (_validated_relative_file(raw_path, label=label) for raw_path in raw_paths)
    }


def _validated_mutation_path(path: Path) -> Path:
    if path.is_absolute() or ".." in path.parts:
        raise RuntimeError(f"mutation path must be project-relative: {path}")
    return path


def _owned_mutants_root() -> Path:
    root = Path.cwd().resolve()
    expected = root / "mutants"
    if expected.is_symlink():
        raise RuntimeError(
            "mutants directory must be project-owned as the literal project-owned mutants directory",
        )
    return expected


def _owned_mutation_artifact(path: Path, suffix: str) -> Path:
    relative = _validated_mutation_path(path)
    root = _owned_mutants_root()
    artifact = Path("mutants") / f"{relative}{suffix}"
    resolved = artifact.resolve() if artifact.exists() else artifact.resolve(strict=False)
    try:
        resolved.relative_to(root)
    except ValueError as exc:
        raise RuntimeError(f"mutation artifact escapes project-owned mutants/: {artifact}") from exc
    return artifact


def _mutmut_version() -> str:
    import mutmut

    return getattr(mutmut, "__version__", "")


def assert_supported_mutmut() -> None:
    """Fail closed if the mutmut internals this script calls have changed."""
    from mutmut import __main__ as mutmut_main

    if _mutmut_version() != SUPPORTED_MUTMUT_VERSION:
        raise RuntimeError(
            f"unsupported mutmut version {_mutmut_version()!r}; expected {SUPPORTED_MUTMUT_VERSION}",
        )

    expected = {
        "copy_src_dir": ("() -> 'None'",),
        "copy_also_copy_files": ("() -> 'None'",),
        "setup_source_paths": ("() -> 'None'",),
        "create_mutants": ("(max_children: 'int') -> 'MutantGenerationStats'",),
        "collect_or_load_stats": (
            "(runner: 'TestRunner', *, mutants_caught_by_type_checker: 'dict[str, Any] | None' = None, "
            "apply_config_invalidation: 'bool' = False, invalidate_stale_callers: 'bool' = True) -> 'None'",
            "(runner, *, mutants_caught_by_type_checker=None, apply_config_invalidation=False, "
            "invalidate_stale_callers=True) -> 'None'",
        ),
        "load_stats": ("() -> 'bool'",),
    }
    for name, signatures in expected.items():
        actual = str(inspect.signature(getattr(mutmut_main, name)))
        if actual not in signatures:
            raise RuntimeError(f"mutmut internal {name} signature changed: {actual}")


def configured_paths() -> tuple[list[Path], list[Path]]:
    """Return all mutatable paths and exact required paths from the loaded config."""
    from mutmut import __main__ as mutmut_main
    from mutmut.configuration import Config

    Config.ensure_loaded()
    if not Config.get().mutate_only_covered_lines:
        raise RuntimeError("[tool.mutmut].mutate_only_covered_lines must be true")

    mutatable = sorted(mutmut_main.walk_mutatable_files(), key=str)
    required = exact_python_targets(load_mutmut_config(Path("pyproject.toml")))
    missing = [path for path in required if path not in mutatable]
    if missing:
        joined = ", ".join(str(path) for path in missing)
        raise RuntimeError(f"required mutation target is not mutatable by mutmut config: {joined}")
    if not mutatable:
        raise RuntimeError("mutmut configuration produced no mutatable files")
    return mutatable, required


def remove_mutation_artifacts(paths: list[Path]) -> None:
    """Remove generated files whose freshness matters for covered-line reuse."""
    for path in paths:
        for suffix in ("", ".meta", ".spans"):
            artifact = _owned_mutation_artifact(path, suffix)
            if artifact.exists():
                artifact.unlink()
    if MANIFEST_PATH.exists():
        MANIFEST_PATH.unlink()


def _read_meta(path: Path) -> dict[str, Any]:
    meta_path = Path("mutants") / f"{path}.meta"
    try:
        metadata = json.loads(meta_path.read_text(encoding="utf-8"))
    except FileNotFoundError as exc:
        raise RuntimeError(f"missing mutmut metadata for {path}") from exc
    except json.JSONDecodeError as exc:
        raise RuntimeError(f"malformed mutmut metadata for {path}: {exc}") from exc
    if not isinstance(metadata, dict) or set(metadata) != META_KEYS:
        raise RuntimeError(f"unexpected mutmut metadata schema for {path}: {sorted(metadata)}")
    if not isinstance(metadata["exit_code_by_key"], dict):
        raise RuntimeError(f"mutmut metadata exit_code_by_key is not an object for {path}")
    return metadata


def _validate_spans(path: Path, metadata: dict[str, Any]) -> None:
    from mutmut.mutation.data import MutantLineSpans

    spans = MutantLineSpans.load(path)
    if spans is None:
        raise RuntimeError(f"missing or unsupported mutmut spans index for {path}")
    generated_names = {key.rpartition(".")[2] for key in metadata["exit_code_by_key"]}
    missing = sorted(generated_names - set(spans.span_by_function_name))
    if missing:
        raise RuntimeError(f"mutmut spans index for {path} is missing mutants: {missing[:5]}")


def validate_prepared_artifacts(
    *,
    mutatable: list[Path],
    required: list[Path],
    covered_lines: dict[str, set[int]] | None = None,
) -> dict[str, Any]:
    """Validate generated artifacts and return a fresh-source manifest."""
    from mutmut.configuration import Config

    source_hashes: dict[str, str] = {}
    selected_test_hashes: dict[str, str] = {}
    hook_hashes: dict[str, str] = {}
    covered_line_counts: dict[str, int] = {}
    mutant_counts: dict[str, int] = {}

    for path in mutatable:
        source = _validated_mutation_path(Path(path))
        mutant = _owned_mutation_artifact(path, "")
        if not source.is_file():
            raise RuntimeError(f"configured mutation source does not exist: {path}")
        if not mutant.is_file() or source.stat().st_mtime >= mutant.stat().st_mtime:
            raise RuntimeError(f"mutmut artifact for {path} is missing or stale")
        metadata = _read_meta(path)
        _validate_spans(path, metadata)
        source_hashes[str(path)] = _sha256(source)
        mutant_counts[str(path)] = len(metadata["exit_code_by_key"])
        if covered_lines is not None:
            key = str((Path("mutants") / path).absolute())
            covered_line_counts[str(path)] = len(covered_lines.get(key, set()))

    for path in required:
        if mutant_counts.get(str(path), 0) == 0:
            raise RuntimeError(f"required mutation target produced zero covered mutants: {path}")
        if covered_lines is not None and covered_line_counts.get(str(path), 0) == 0:
            raise RuntimeError(f"required mutation target had no covered lines: {path}")

    config = Config.get()
    selected_test_hashes = _hash_relative_files(
        list(config.pytest_add_cli_args_test_selection),
        label="selected mutation test file",
    )
    hook_hashes = _hash_relative_files([str(path) for path in HOOK_HASH_PATHS], label="mutation hook")
    return {
        "version": 1,
        "mutmut_version": _mutmut_version(),
        "mutate_only_covered_lines": config.mutate_only_covered_lines,
        "source_paths": [str(path) for path in config.source_paths],
        "only_mutate": list(config.only_mutate),
        "pytest_add_cli_args": list(config.pytest_add_cli_args),
        "pytest_add_cli_args_test_selection": list(config.pytest_add_cli_args_test_selection),
        "required_mutation_targets": [str(path) for path in required],
        "mutatable_paths": [str(path) for path in mutatable],
        "source_hashes": source_hashes,
        "selected_test_hashes": selected_test_hashes,
        "hook_hashes": hook_hashes,
        "covered_line_counts": covered_line_counts,
        "mutant_counts": mutant_counts,
    }


def prepare(max_children: int) -> dict[str, Any]:
    """Collect coverage, generate covered-only mutants, validate them, and write a manifest."""
    import mutmut
    from mutmut import __main__ as mutmut_main
    from mutmut.code_coverage import gather_coverage
    from mutmut.state import state

    assert_supported_mutmut()
    mutatable, required = configured_paths()
    Path("mutants").mkdir(exist_ok=True)
    remove_mutation_artifacts(mutatable)

    mutmut_main.copy_src_dir()
    mutmut_main.copy_also_copy_files()
    mutmut_main.setup_source_paths()
    source_files = list(mutmut_main.walk_source_files())
    covered_lines = gather_coverage(mutmut_main.PytestRunner(), source_files)
    mutmut._covered_lines = covered_lines
    stats = mutmut_main.create_mutants(max_children)
    state().current_function_hashes.clear()

    manifest = validate_prepared_artifacts(
        mutatable=mutatable,
        required=required,
        covered_lines=covered_lines,
    )
    manifest["generation_stats"] = {
        "mutated": stats.mutated,
        "ignored": stats.ignored,
        "unmodified": stats.unmodified,
    }
    MANIFEST_PATH.write_text(json.dumps(manifest, indent=2, sort_keys=True) + "\n", encoding="utf-8")
    return manifest


def validate_coverage_manifest() -> dict[str, Any]:
    """Validate that a prior prepass still matches the current sources and config."""
    assert_supported_mutmut()
    mutatable, required = configured_paths()
    try:
        manifest = json.loads(MANIFEST_PATH.read_text(encoding="utf-8"))
    except FileNotFoundError as exc:
        raise RuntimeError(f"missing covered-line prepass manifest: {MANIFEST_PATH}") from exc
    except json.JSONDecodeError as exc:
        raise RuntimeError(f"malformed covered-line prepass manifest: {exc}") from exc

    expected = validate_prepared_artifacts(mutatable=mutatable, required=required)
    comparable_keys = (
        "version",
        "mutmut_version",
        "mutate_only_covered_lines",
        "source_paths",
        "only_mutate",
        "pytest_add_cli_args",
        "pytest_add_cli_args_test_selection",
        "required_mutation_targets",
        "mutatable_paths",
        "source_hashes",
        "selected_test_hashes",
        "hook_hashes",
    )
    for key in comparable_keys:
        if manifest.get(key) != expected[key]:
            raise RuntimeError(f"covered-line prepass manifest is stale for {key}")
    return manifest


def validate_manifest() -> dict[str, Any]:
    """Backward-compatible name for covered-line manifest validation."""
    return validate_coverage_manifest()


def _read_stats() -> dict[str, Any]:
    try:
        stats = json.loads(STATS_PATH.read_text(encoding="utf-8"))
    except FileNotFoundError as exc:
        raise RuntimeError(f"missing prepared mutmut stats: {STATS_PATH}") from exc
    except json.JSONDecodeError as exc:
        raise RuntimeError(f"malformed prepared mutmut stats: {exc}") from exc
    if not isinstance(stats, dict):
        raise RuntimeError("prepared mutmut stats root must be an object")
    return stats


def validate_stats_artifacts() -> dict[str, Any]:
    """Validate prepared mutmut test-selection stats and return manifest fields."""
    stats = _read_stats()
    expected_keys = {
        "tests_by_mangled_function_name",
        "duration_by_test",
        "stats_time",
        "function_hashes",
        "function_dependencies",
        "config_fingerprint",
        "watched_file_hashes",
        "git_commit",
    }
    if set(stats) != expected_keys:
        raise RuntimeError(f"unexpected mutmut stats schema: {sorted(stats)}")

    tests_by_function = stats["tests_by_mangled_function_name"]
    duration_by_test = stats["duration_by_test"]
    function_hashes = stats["function_hashes"]
    if not isinstance(tests_by_function, dict) or not tests_by_function:
        raise RuntimeError("prepared mutmut stats have no test-to-function mapping")
    if not isinstance(duration_by_test, dict) or not duration_by_test:
        raise RuntimeError("prepared mutmut stats have no test duration mapping")
    if not isinstance(function_hashes, dict) or not function_hashes:
        raise RuntimeError("prepared mutmut stats have no function hash baseline")
    associated_tests = {
        test
        for tests in tests_by_function.values()
        if isinstance(tests, list)
        for test in tests
    }
    if not associated_tests:
        raise RuntimeError("prepared mutmut stats do not associate any tests with mutants")
    missing_durations = sorted(associated_tests - set(duration_by_test))
    if missing_durations:
        raise RuntimeError(
            "prepared mutmut stats reference tests without durations: "
            f"{missing_durations[:5]}",
        )
    return {
        "stats_hash": _sha256(STATS_PATH),
        "stats_function_count": len(function_hashes),
        "stats_test_count": len(duration_by_test),
        "stats_association_count": sum(
            len(tests) for tests in tests_by_function.values() if isinstance(tests, list)
        ),
    }


def prepare_stats(max_children: int) -> dict[str, Any]:
    """Build mutmut test-selection stats in a fresh process, then exit."""
    from mutmut import __main__ as mutmut_main
    from mutmut.state import state

    manifest = validate_coverage_manifest()
    if STATS_PATH.exists():
        STATS_PATH.unlink()

    mutmut_main.copy_src_dir()
    mutmut_main.copy_also_copy_files()
    mutmut_main.setup_source_paths()
    mutmut_main.create_mutants(max_children)
    runner = mutmut_main.PytestRunner()
    runner.prepare_main_test_run()
    mutmut_main.collect_or_load_stats(
        runner,
        mutants_caught_by_type_checker={},
        apply_config_invalidation=False,
        invalidate_stale_callers=False,
    )
    if not state().current_function_hashes:
        raise RuntimeError("mutmut stats prepass did not populate function hashes")

    manifest.update(validate_stats_artifacts())
    MANIFEST_PATH.write_text(json.dumps(manifest, indent=2, sort_keys=True) + "\n", encoding="utf-8")
    return manifest


def validate_execution_manifest() -> dict[str, Any]:
    """Validate coverage artifacts plus prepared test-selection stats."""
    manifest = validate_coverage_manifest()
    stats_fields = validate_stats_artifacts()
    for key, value in stats_fields.items():
        if manifest.get(key) != value:
            raise RuntimeError(f"prepared mutmut execution manifest is stale for {key}")
    return manifest


def _parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "command",
        choices=("prepare", "prepare-coverage", "prepare-stats", "validate", "validate-execution"),
    )
    parser.add_argument("--max-children", type=int, default=4)
    return parser


def main() -> int:
    args = _parser().parse_args()
    try:
        if args.command in {"prepare", "prepare-coverage"}:
            manifest = prepare(args.max_children)
            total = sum(manifest["mutant_counts"].values())
            print(f"prepared covered-line mutation artifacts for {total} mutants")
        elif args.command == "prepare-stats":
            manifest = prepare_stats(args.max_children)
            print(
                "prepared mutmut test-selection stats for "
                f"{manifest['stats_association_count']} function/test associations",
            )
        elif args.command == "validate-execution":
            validate_execution_manifest()
            print("prepared mutation execution artifacts are fresh")
        else:
            validate_coverage_manifest()
            print("covered-line mutation artifacts are fresh")
    except Exception as exc:
        print(f"mutation coverage preparation failed: {exc}")
        return 1
    return 0


if __name__ == "__main__":
    exit_code = main()
    sys.stdout.flush()
    sys.stderr.flush()
    os._exit(exit_code)