from __future__ import annotations

from copy import deepcopy
from dataclasses import replace
import json
import os
from pathlib import Path
import subprocess
import tempfile
import unittest
from unittest.mock import patch

import numpy as np

from app.main import compile_system_xml_network, run_system_xml_simulation
from app.simulation.backends import simulation_config
from app.simulation.native_codegen.build import build_native, toolchain
from app.simulation.native_codegen.compiler import NativeCapabilityError, compile_native_program
from app.simulation.native_codegen.input import arithmetic_value, load_input, project_xml
from app.simulation.native_codegen.runner import execute_native
from app.simulation.config import SolverActivityTracker
from app.system_xml import validate_system_xml_document

ROOT = Path(__file__).resolve().parents[1]
FIXTURE = ROOT / "tests/fixtures/native-skill-test.xml"


def network(xml=None):
    doc = validate_system_xml_document(xml or FIXTURE.read_bytes()).document
    return compile_system_xml_network(doc)


class NativeInputTests(unittest.TestCase):
    def test_arithmetic_is_bounded_and_does_not_execute_code(self):
        self.assertEqual(arithmetic_value("3.14*10**2/4"), 78.5)
        for bad in ("__import__('os')", "a.b", "[1][0]", "2**100000000", "1/0"):
            with self.assertRaises((ValueError, ZeroDivisionError)):
                arithmetic_value(bad)

    def test_generated_program_has_complete_state_and_output_mapping(self):
        program = compile_native_program(network())
        self.assertEqual(len(program.state_keys), 12)
        self.assertEqual(len(program.variables), 175)
        self.assertNotIn("PyObject", program.source)
        self.assertNotIn("amesim_mecmas21_5", program.source.split("int model_eval")[1])

    def test_unknown_component_is_rejected_before_build(self):
        n = network()
        original = n.components["amesim_mecmas21_2"]
        original.__class__ = type('CustomUnportedMass', (type(original),), {})
        with self.assertRaisesRegex(NativeCapabilityError, "no native contract"):
            compile_native_program(n)


class NativeExecutionTests(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        try:
            toolchain()
        except (OSError, RuntimeError, subprocess.SubprocessError) as exc:
            raise unittest.SkipTest(f"Native toolchain unavailable: {exc}")
        cls.temporary = tempfile.TemporaryDirectory(prefix="native-tests-")
        cls.root = Path(cls.temporary.name)
        cls.document = validate_system_xml_document(FIXTURE.read_bytes()).document
        cls.program = compile_native_program(network())
        cls.build = build_native(cls.program)

    @classmethod
    def tearDownClass(cls):
        cls.temporary.cleanup()





    def test_cache_and_standalone_executable_without_python_path(self):
        cached = build_native(self.program)
        self.assertTrue(cached.cache_hit)
        isolated_path = str(Path(os.environ["SystemRoot"]) / "System32") if os.name == "nt" else str(self.root / "no-path")
        env = {**os.environ, "PATH": isolated_path}
        r = subprocess.run([str(cached.executable), "--init"], env=env, cwd=self.root,
                           capture_output=True, text=True, check=True)
        self.assertEqual(len(json.loads(r.stdout)), 12)

    def test_both_integrators_complete_and_respect_max_step(self):
        for method in ("RK45", "BDF"):
            config = replace(simulation_config(self.document.simulation), method=method, t_stop=.1)
            data = execute_native(self.build, config, .02, run_dir=self.root / method)
            self.assertTrue(data["success"])
            self.assertEqual(data["simulatedUntil"], .1)
            self.assertEqual(data["jacobianMode"], "dense-difference" if method == "BDF" else "not-used")
            self.assertEqual(data["jacobianRhsCalls"], 0)  # Compact model has no grouping benefit.
            self.assertLessEqual(data["maxAcceptedStep"], config.max_step+1e-14)
            self.assertEqual(set(data["series"]), {"time", *(v.key for v in self.program.variables)})
            self.assertTrue(all(np.isfinite(v).all() for v in map(np.asarray, data["series"].values())))

    def test_retired_jacobian_selector_is_rejected(self):
        for policy in ("dense", "auto", "verify"):
            with self.subTest(policy=policy):
                result = subprocess.run([str(self.build.executable), "--jacobian", policy],
                                        cwd=self.root, capture_output=True, text=True, timeout=10)
                self.assertEqual(result.returncode, 64)

    def test_cancellation_returns_partial_accepted_state(self):
        config = replace(simulation_config(self.document.simulation), max_step=1e-6)
        tracker = SolverActivityTracker()
        data = execute_native(self.build, config, .02, run_dir=self.root / "cancel",
                              cancel_check=lambda: True, activity_tracker=tracker)
        self.assertEqual(data["status"], "cancelled")
        self.assertLess(data["simulatedUntil"], config.t_stop)
        self.assertEqual(data["series"]["time"][-1], data["simulatedUntil"])

    def test_xml_api_defaults_to_native_and_preserves_result_contract(self):
        xml = FIXTURE.read_bytes().replace(b'tStop="10"', b'tStop="0.1"')
        with patch.dict(os.environ):
            os.environ.pop("SIMULATION_NUMERIC_ENGINE", None)
            data = run_system_xml_simulation(xml)
        self.assertTrue(data["success"])
        self.assertEqual(data["diagnostics"]["backend"], "native-c")
        self.assertEqual(data["diagnostics"]["stateCount"], 12)
        self.assertEqual(data["diagnostics"]["sampleCount"], len(data["series"]["time"]))
        self.assertEqual(len(data["variables"]), 175)
        self.assertEqual(data["series"]["time"][-1], .1)


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