"""Exercise the native build cache with real, tiny C translation units.

These tests isolate their sources, libraries and caches in a temporary directory.
They execute the linked programs to check behavior, and observe compiler commands
to distinguish actual object reuse from a reported cache hit.
"""
from __future__ import annotations

from concurrent.futures import ThreadPoolExecutor
from contextlib import ExitStack
from dataclasses import replace
from hashlib import sha256
import json
import os
from pathlib import Path
import shutil
import subprocess
import sys
import tempfile
import threading
import unittest
from unittest.mock import patch

from app.simulation.native_codegen import build as native_build
from app.simulation.native_codegen.compiler import NativeProgram
from app.simulation.native_codegen.cache_storage import prune_cache
from app.simulation.native_codegen.runner import execute_native
from app.simulation.config import SolveIVPConfig


@unittest.skipUnless(sys.platform.startswith("linux") and shutil.which("gcc"),
                     "Small real-compiler fixtures currently require Linux gcc")
class NativeBuildCacheTests(unittest.TestCase):
    def setUp(self):
        self.temporary = tempfile.TemporaryDirectory(prefix="native-build-cache-tests-")
        self.addCleanup(self.temporary.cleanup)
        self.root = Path(self.temporary.name)
        self.native = self.root / "native"
        self.cache = self.root / "cache"
        self.sundials = self.root / "sundials"
        self.compiler = str(Path(shutil.which("gcc")).resolve())
        self.real_run = subprocess.run
        self.compiler_version = self.real_run(
            [self.compiler, "--version"], capture_output=True, text=True,
            check=True, timeout=10,
        ).stdout.splitlines()[0]
        self.commands: list[tuple[str, ...]] = []
        self.command_lock = threading.Lock()
        self._write("native/THIRD_PARTY_NOTICES.txt", "Fixture code belongs to this test.\n")
        self._write("native/include/values.h", '#include "nested/value.h"\n')
        self._write("native/include/nested/value.h", "#define COMMON_VALUE 3\n")
        self._write("native/include/signal_config.h", "#define SIGNAL_VALUE 5\n")
        self._write("native/runtime/main.c", '''#include "model.h"
#include <stdio.h>
int common_value(void);
int main(void) {
    printf("%.17g\\n", model_value() + common_value() + HEADER_OFFSET);
    return 0;
}
''')
        self._write("native/runtime/common.c", '''#include "values.h"
int common_value(void) { return COMMON_VALUE; }
''')
        for name in ("rk45", "cvode_solver", "json_numbers"):
            self._write(f"native/runtime/{name}.c", "/* unused runtime fixture */\n")
        self._write("native/encoding/ryu/d2s.c", "/* unused encoder fixture */\n")
        self._write("native/components/modules/signal.c", '''#include "signal_config.h"
double native_signal(void) { return SIGNAL_VALUE; }
''')
        for name in ("properties", "pipe", "orifice", "mechanics"):
            self._write(f"native/components/modules/{name}.c", "/* unneeded component */\n")
        self._write("sundials/include/cvode/cvode.h", "/* fixture */\n")
        self._write("sundials/include/sundials/sundials_config.h", "/* fixture */\n")
        # An empty archive is sufficient: the fixture does not call SUNDIALS.
        for name in native_build.LIBRARIES:
            self._write(f"sundials/lib/libsundials_{name}.a", b"!<arch>\n")
        self.program = NativeProgram(
            source='#include "model.h"\ndouble model_value(void) { return native_signal() + 1; }\n',
            header="#define HEADER_OFFSET 0\ndouble model_value(void);\ndouble native_signal(void);\n",
            state_keys=(), variables=(), component_types=(),
        )
        self.patches = ExitStack()
        self.addCleanup(self.patches.close)
        self.patches.enter_context(patch.object(native_build, "ROOT", self.root))
        self.patches.enter_context(patch.object(native_build, "NATIVE", self.native))
        self.patches.enter_context(patch.object(
            native_build, "toolchain", return_value=(self.compiler, self.sundials, self.compiler_version),
        ))
        self.patches.enter_context(patch.object(native_build.subprocess, "run", side_effect=self._run))

    def _write(self, relative: str, content: str | bytes) -> Path:
        path = self.root / relative
        path.parent.mkdir(parents=True, exist_ok=True)
        if isinstance(content, bytes):
            path.write_bytes(content)
        else:
            path.write_text(content, encoding="utf-8")
        return path

    def _run(self, command, *args, **kwargs):
        with self.command_lock:
            self.commands.append(tuple(map(str, command)))
        return self.real_run(command, *args, **kwargs)

    def _build(self, program=None):
        return native_build.build_native(program or self.program, cache_dir=self.cache)

    def _compiled_count(self) -> int:
        return sum("-c" in command for command in self.commands)

    def _linked_count(self) -> int:
        return sum("-o" in command and "-c" not in command and "-E" not in command
                   for command in self.commands)

    def _output(self, build) -> float:
        completed = self.real_run([str(build.executable)], capture_output=True, text=True,
                                  check=True, timeout=10)
        return float(completed.stdout.strip())

    def _parameter_program(self, value: int):
        return replace(self.program, source=self.program.source.replace(" + 1;", f" + {value};"))

    def _objects(self):
        return sorted(path for path in (self.cache / "objects").iterdir()
                      if len(path.name) == 64 and (path / "manifest.json").is_file())

    def _object_for(self, suffix: str) -> Path:
        matches = [path for path in self._objects()
                   if json.loads((path / "manifest.json").read_text())["sourceName"].endswith(suffix)]
        self.assertEqual(len(matches), 1, suffix)
        return matches[0]

    def test_complete_hit_executes_same_program_without_compilation_or_linking(self):
        first = self._build()
        self.assertFalse(first.cache_hit)
        self.assertEqual(self._output(first), 9)
        self.assertGreater(self._compiled_count(), 1)
        self.commands.clear()
        second = self._build()
        self.assertTrue(second.cache_hit)
        self.assertEqual(first.manifest["buildKey"], second.manifest["buildKey"])
        self.assertEqual(self._output(second), 9)
        self.assertEqual(self._compiled_count(), 0)
        self.assertEqual(self._linked_count(), 0)

    def test_parameter_edit_reuses_shared_objects_and_recompiles_only_model(self):
        first = self._build()
        self.commands.clear()
        second = self._build(self._parameter_program(11))
        self.assertFalse(second.cache_hit)
        self.assertNotEqual(first.manifest["buildKey"], second.manifest["buildKey"])
        self.assertEqual(self._output(second), 19)
        self.assertEqual(self._compiled_count(), 1)
        self.assertEqual(self._linked_count(), 1)
        self.commands.clear()
        self.assertTrue(self._build().cache_hit)
        self.assertEqual(self._compiled_count(), 0)

    def test_model_header_edit_invalidates_only_translation_units_that_include_it(self):
        self._build()
        self.commands.clear()
        changed = replace(self.program, header=self.program.header.replace("OFFSET 0", "OFFSET 7"))
        result = self._build(changed)
        self.assertEqual(self._output(result), 16)
        # Both include model.h, but the model TU's preprocessed tokens are
        # unchanged because it does not use HEADER_OFFSET. Only main must rebuild.
        self.assertEqual(self._compiled_count(), 1)

    def test_model_header_declaration_edit_rebuilds_all_actual_dependents(self):
        self._build()
        self.commands.clear()
        changed = replace(self.program, header=self.program.header + "extern int unused_model_state;\n")
        result = self._build(changed)
        self.assertEqual(self._output(result), 9)
        self.assertEqual(self._compiled_count(), 2)

    def test_indirect_header_change_invalidates_the_actual_runtime_object(self):
        first = self._build()
        self.commands.clear()
        self._write("native/include/nested/value.h", "#define COMMON_VALUE 13\n")
        second = self._build()
        self.assertFalse(second.cache_hit)
        self.assertNotEqual(first.manifest["buildKey"], second.manifest["buildKey"])
        self.assertEqual(self._output(second), 19)
        self.assertEqual(self._compiled_count(), 1)

    def test_unneeded_modules_and_headers_do_not_invalidate_or_enter_the_build(self):
        first = self._build()
        self.commands.clear()
        self._write("native/components/modules/pipe.c", "this is deliberately invalid C\n")
        self._write("native/components/modules/new_unused.c", "another deliberately invalid unit\n")
        self._write("native/include/unused_new_header.h", "#error should never be included\n")
        second = self._build()
        self.assertTrue(second.cache_hit)
        self.assertEqual(first.manifest["buildKey"], second.manifest["buildKey"])
        self.assertEqual(self._output(second), 9)
        self.assertEqual(self._compiled_count(), 0)

    def test_used_module_change_rebuilds_one_object_and_changes_executable(self):
        first = self._build()
        self.commands.clear()
        self._write("native/include/signal_config.h", "#define SIGNAL_VALUE 15\n")
        second = self._build()
        self.assertNotEqual(first.manifest["buildKey"], second.manifest["buildKey"])
        self.assertEqual(self._output(second), 19)
        self.assertEqual(self._compiled_count(), 1)

    def test_compile_flags_and_compiler_identity_cannot_reuse_old_objects(self):
        first = self._build()
        original_count = self._compiled_count()
        self.commands.clear()
        with patch.object(native_build, "COMPILER_FLAGS", (*native_build.COMPILER_FLAGS, "-DTEST_CACHE_VARIANT=1")):
            flags_build = self._build()
        self.assertNotEqual(first.manifest["buildKey"], flags_build.manifest["buildKey"])
        self.assertEqual(self._compiled_count(), original_count)
        self.assertEqual(self._output(flags_build), 9)
        self.commands.clear()
        with patch.object(native_build, "toolchain", return_value=(self.compiler, self.sundials, self.compiler_version + " test revision")):
            compiler_build = self._build()
        self.assertNotEqual(first.manifest["buildKey"], compiler_build.manifest["buildKey"])
        self.assertEqual(self._compiled_count(), original_count)
        self.assertEqual(self._output(compiler_build), 9)

    def test_library_change_relinks_but_reuses_compiled_objects(self):
        first = self._build()
        self.commands.clear()
        # Valid ar symbol-table member with no entries, changing bytes only.
        empty_object = self.root / "library-member.c"
        empty_object.write_text("int library_fixture_symbol(void) { return 1; }\n")
        object_path = self.root / "library-member.o"
        self.real_run([self.compiler, "-c", str(empty_object), "-o", str(object_path)], check=True)
        self.real_run(["ar", "r", str(self.sundials / "lib/libsundials_core.a"), str(object_path)],
                      capture_output=True, check=True)
        second = self._build()
        self.assertNotEqual(first.manifest["buildKey"], second.manifest["buildKey"])
        self.assertEqual(self._compiled_count(), 0)
        self.assertEqual(self._linked_count(), 1)
        self.assertEqual(self._output(second), 9)

    def test_model_manifest_requires_identity_and_complete_safe_artifacts(self):
        built = self._build()
        manifest_path = built.executable.parent / "manifest.json"
        original = manifest_path.read_text()
        manifest = json.loads(original)
        mutations = {
            "empty artifacts": lambda value: value.update(artifacts={}),
            "missing executable": lambda value: value["artifacts"].pop(built.executable.name),
            "missing generated header": lambda value: value["artifacts"].pop("model.h"),
            "wrong build key": lambda value: value.update(buildKey="0" * 64),
            "wrong cache version": lambda value: value.update(cacheVersion=-1),
            "wrong object identities": lambda value: value.update(objectKeys=["0" * 64]),
            "wrong result contract": lambda value: value.update(stateKeys=["unexpected.state"]),
            "parent traversal": lambda value: value["artifacts"].update({"../outside": "0" * 64}),
            "absolute path": lambda value: value["artifacts"].update({str(self.root / "outside"): "0" * 64}),
        }
        for label, mutate in mutations.items():
            with self.subTest(label=label):
                modified = json.loads(original)
                mutate(modified)
                manifest_path.write_text(json.dumps(modified))
                try:
                    with self.assertRaises(RuntimeError):
                        self._build()
                finally:
                    manifest_path.write_text(original)
        self.assertEqual(manifest["cacheVersion"], 2)
        self.assertTrue(self._build().cache_hit)

    def test_model_source_cannot_be_replaced_by_updating_its_artifact_digest(self):
        built = self._build()
        source = built.executable.parent / "model.c"
        source.write_text(self._parameter_program(12).source)
        manifest_path = built.executable.parent / "manifest.json"
        manifest = json.loads(manifest_path.read_text())
        manifest["artifacts"]["model.c"] = sha256(source.read_bytes()).hexdigest()
        manifest_path.write_text(json.dumps(manifest))
        with self.assertRaises(RuntimeError):
            self._build()

    def test_changed_executable_and_artifact_symlink_are_rejected(self):
        built = self._build()
        executable = built.executable
        original = executable.read_bytes()
        executable.write_bytes(original + b"corrupt")
        with self.assertRaises(RuntimeError):
            self._build()
        executable.write_bytes(original)
        target = self.root / "outside-executable"
        target.write_bytes(original)
        executable.unlink()
        executable.symlink_to(target)
        with self.assertRaises(RuntimeError):
            self._build()

    def test_object_bytes_and_object_manifest_are_checked_when_reused(self):
        first = self._build()
        object_dir = self._object_for("signal.c")
        unit = object_dir / "unit.o"
        manifest_path = object_dir / "manifest.json"
        original_object = unit.read_bytes()
        original_manifest = manifest_path.read_text()
        changed_program = self._parameter_program(2)
        unit.write_bytes(original_object + b"corrupt")
        # Complete executable reuse need not inspect disposable intermediate objects.
        self.assertTrue(self._build().cache_hit)
        with self.assertRaises(RuntimeError):
            self._build(changed_program)
        unit.write_bytes(original_object)
        for label, mutate in {
            "empty artifacts": lambda value: value.update(artifacts={}),
            "wrong key": lambda value: value.update(objectKey="0" * 64),
            "wrong version": lambda value: value.update(cacheVersion=-1),
            "wrong source": lambda value: value.update(sourceName="native/another.c"),
            "wrong preprocessing": lambda value: value.update(preprocessedSha256="0" * 64),
            "wrong compiler": lambda value: value.update(compiler={}),
            "parent traversal": lambda value: value["artifacts"].update({"../unit.o": "0" * 64}),
        }.items():
            with self.subTest(label=label):
                modified = json.loads(original_manifest)
                mutate(modified)
                manifest_path.write_text(json.dumps(modified))
                try:
                    with self.assertRaises(RuntimeError):
                        self._build(changed_program)
                finally:
                    manifest_path.write_text(original_manifest)
        target = self.root / "outside-unit.o"
        target.write_bytes(original_object)
        unit.unlink()
        unit.symlink_to(target)
        with self.assertRaises(RuntimeError):
            self._build(changed_program)
        self.assertEqual(self._output(first), 9)

    def test_library_change_during_link_is_rejected_before_publication(self):
        first = self._build()
        before = {path.name for path in (self.cache / "models").iterdir() if len(path.name) == 64}
        source = self._write("replacement-library.c", "int replacement_library(void) { return 7; }\n")
        object_path = self.root / "replacement-library.o"
        alternate = self.root / "alternate.a"
        self.real_run([self.compiler, "-c", str(source), "-o", str(object_path)], check=True)
        self.real_run(["ar", "rc", str(alternate), str(object_path)], capture_output=True, check=True)
        library = self.sundials / "lib/libsundials_core.a"
        original = library.read_bytes()
        linked_successfully = []

        def replace_library_at_link(command, *args, **kwargs):
            is_link = "-o" in command and "-c" not in command and "-E" not in command
            if is_link:
                library.write_bytes(alternate.read_bytes())
            result = self._run(command, *args, **kwargs)
            if is_link:
                linked_successfully.append(result.returncode == 0)
            return result

        try:
            with patch.object(native_build.subprocess, "run", side_effect=replace_library_at_link):
                with self.assertRaises(RuntimeError):
                    self._build(self._parameter_program(2))
            self.assertEqual(linked_successfully, [True])
            published = {path.name for path in (self.cache / "models").iterdir() if len(path.name) == 64}
            self.assertEqual(published, before)
            self.assertEqual(self._output(first), 9)
        finally:
            library.write_bytes(original)
        self.assertTrue(self._build().cache_hit)

    def test_failed_compile_or_link_does_not_publish_a_model_or_break_old_cache(self):
        first = self._build()
        before = {path.name for path in (self.cache / "models").iterdir() if len(path.name) == 64}
        for source in (
            self.program.source + "this is invalid C;\n",
            self.program.source + "int unresolved(void); int force_link_failure(void) { return unresolved(); }\n",
        ):
            with self.subTest(source=source):
                with self.assertRaises(RuntimeError):
                    self._build(replace(self.program, source=source))
                published = {path.name for path in (self.cache / "models").iterdir() if len(path.name) == 64}
                self.assertEqual(published, before)
                self.assertTrue(self._build().cache_hit)
                self.assertEqual(self._output(first), 9)

    def test_concurrent_cold_builds_publish_complete_models_and_share_objects(self):
        gate = threading.Barrier(2)

        def build_together(program):
            gate.wait(timeout=10)
            return self._build(program)

        with ThreadPoolExecutor(max_workers=2) as pool:
            futures = [pool.submit(build_together, self.program) for _ in range(2)]
            first, second = [future.result(timeout=60) for future in futures]
        self.assertEqual(first.executable, second.executable)
        self.assertEqual(self._output(first), 9)
        self.assertEqual(self._output(second), 9)
        self.assertTrue(self._build().cache_hit)
        gate = threading.Barrier(2)
        with ThreadPoolExecutor(max_workers=2) as pool:
            futures = [pool.submit(build_together, self._parameter_program(value)) for value in (2, 3)]
            changed = [future.result(timeout=60) for future in futures]
        self.assertEqual([self._output(item) for item in changed], [10, 11])
        self.assertNotEqual(changed[0].manifest["buildKey"], changed[1].manifest["buildKey"])
        for path in self._objects():
            self.assertTrue((path / "unit.o").is_file())
        for item in (first, *changed):
            manifest = json.loads((item.executable.parent / "manifest.json").read_text())
            self.assertEqual(manifest["buildKey"], item.manifest["buildKey"])
            self.assertTrue(all((item.executable.parent / name).is_file() for name in manifest["artifacts"]))


    def test_generated_sources_keep_exact_lf_under_windows_default_translation(self):
        write_text = Path.write_text

        def windows_write_text(path, data, encoding=None, errors=None, newline=None):
            # Python text output on Windows expands LF when newline is omitted.
            # Emulate that default on Linux while honoring explicit newline="\n".
            if newline is None:
                data = data.replace("\n", "\r\n")
            return write_text(path, data, encoding=encoding, errors=errors, newline="\n")

        with patch.object(Path, "write_text", windows_write_text):
            probe = self._write("platform-newline-probe.txt", "first\nsecond\n")
            self.assertEqual(probe.read_bytes(), b"first\r\nsecond\r\n")
            first = self._build()
            self.assertEqual((first.executable.parent / "model.c").read_bytes(), self.program.source.encode())
            self.assertEqual((first.executable.parent / "model.h").read_bytes(), self.program.header.encode())
            self.assertEqual(self._output(first), 9)
            self.commands.clear()
            second = self._build()
            self.assertTrue(second.cache_hit)
            self.assertEqual(second.manifest["buildKey"], first.manifest["buildKey"])
            self.assertEqual(self._output(second), 9)
            self.assertEqual(self._compiled_count(), 0)
            self.assertEqual(self._linked_count(), 0)

    def test_capacity_sweep_protects_build_leases_and_close_allows_eviction(self):
        older = self._build()
        newer = self._build(self._parameter_program(2))
        try:
            report = prune_cache(self.cache, model_limit_bytes=0)
            self.assertGreaterEqual(report["models"]["skippedInUse"], 2)
            self.assertTrue(older.executable.is_file())
            self.assertTrue(newer.executable.is_file())
            self.assertEqual(self._output(older), 9)
            # Release the newer lease first under the normal budget, avoiding
            # any assumption that two model keys occupy distinct lock shards.
            newer.close()
            with patch.dict(os.environ, {"SIMULATION_NATIVE_MODEL_CACHE_MB": "0"}):
                older.close()
            self.assertFalse(older.executable.exists())
            self.assertTrue(newer.executable.is_file())
            self.assertEqual(self._output(newer), 10)
            report = prune_cache(self.cache, model_limit_bytes=0)
            self.assertEqual(report["models"]["removedEntries"], 0)
            self.assertEqual(report["models"]["oversizedEntries"], 1)
            self.assertGreater(report["models"]["overLimitBytes"], 0)
        finally:
            older.close()
            newer.close()

    def test_executing_native_worker_keeps_its_model_during_capacity_sweep(self):
        self._write("native/runtime/main.c", r"""#include "model.h"
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <time.h>
int common_value(void);
int main(int argc, char **argv) {
    const char *output = NULL;
    const char *release = getenv("NATIVE_CACHE_TEST_RELEASE");
    for (int i = 1; i + 1 < argc; ++i)
        if (strcmp(argv[i], "--output") == 0) output = argv[i + 1];
    if (!output || !release) return 64;
    fprintf(stderr, "{\"phase\":\"integrating\",\"time\":0,\"nfev\":0,\"acceptedSteps\":0}\n");
    fflush(stderr);
    struct timespec pause = {0, 1000000};
    int released = 0;
    for (int i = 0; i < 10000; ++i) {
        FILE *gate = fopen(release, "r");
        if (gate) { fclose(gate); released = 1; break; }
        nanosleep(&pause, NULL);
    }
    if (!released) return 70;
    FILE *result = fopen(output, "w");
    if (!result) return 74;
    int written = fprintf(result,
        "{\"simulatedUntil\":0.1,\"nfev\":1,\"acceptedSteps\":1,\"value\":%.17g}",
        model_value() + common_value() + HEADER_OFFSET);
    return fclose(result) != 0 || written < 0;
}
""")
        active = self._build()
        other = self._build(self._parameter_program(2))
        release = self.root / "release-worker"
        worker_running = threading.Event()
        try:
            with patch.dict(os.environ, {"NATIVE_CACHE_TEST_RELEASE": str(release)}):
                with ThreadPoolExecutor(max_workers=1) as pool:
                    future = pool.submit(
                        execute_native, active, SolveIVPConfig(t_stop=0.1), 0.1,
                        run_dir=self.root / "worker-result", timeout=10,
                        progress_callback=lambda fraction, phase: worker_running.set(),
                    )
                    try:
                        self.assertTrue(worker_running.wait(timeout=5), "Native worker did not report readiness")
                        self.assertFalse(future.done())
                        report = prune_cache(self.cache, model_limit_bytes=0)
                        self.assertGreaterEqual(report["models"]["skippedInUse"], 2)
                        self.assertTrue(active.executable.is_file())
                        self.assertTrue((active.executable.parent / "manifest.json").is_file())
                    finally:
                        release.write_text("finish\n")
                    payload = future.result(timeout=15)
            self.assertEqual(payload["value"], 9)
            self.assertEqual(payload["buildKey"], active.manifest["buildKey"])
            other.close()
            with patch.dict(os.environ, {"SIMULATION_NATIVE_MODEL_CACHE_MB": "0"}):
                active.close()
            self.assertFalse(active.executable.exists())
            self.assertTrue(other.executable.is_file())
        finally:
            release.write_text("finish\n")
            active.close()
            other.close()


if __name__ == "__main__":
    unittest.main()