"""Exercise the actual C block writer and application integration on both OSes."""
from dataclasses import replace
import json
import os
from pathlib import Path
import shutil
import struct
import subprocess
import sys
import tempfile
import unittest
from unittest.mock import patch

from app.simulation.native_codegen import build as builder, result_storage, runner
from app.simulation.native_codegen.runner import execute_native
from app.simulation.native_codegen.compiler import compile_native_program
from app.simulation.config import SolveIVPConfig
from tests.test_native_catalog import Circuit

ROOT = Path(__file__).resolve().parents[1]
HARNESS = r'''
#include "runtime.h"
#include <stdlib.h>
#include <string.h>
static int evaluations=0;
int model_eval(double t,const double *y,double *dy,double *w) {
    (void)t; (void)dy; evaluations++;w[0]=y[0];w[1]=y[1];return 1;
}
int main(int argc,char **argv) {
    if(argc!=4) return 64;
    NativeRun r={0};r.sample_path=argv[1];r.output_blocks_path=argv[2];
    for(int i=0;i<10003;i++) {
        double y[2]={(double)i,-0.0};
        if(!native_samples_append(&r,i,y)) return 2;
        /* Replace the last row while the four-row buffer is full. */
        if(i%4==3) {y[0]+=.5;if(!native_samples_append(&r,i,y)) return 3;}
        if(r.buffered>4 || r.capacity!=4 || evaluations) return 4;
    }
    if(!native_samples_flush(&r) || !native_samples_outputs(&r)) return 5;
    FILE *f=fopen(argv[3],"wb");if(!f) return 6;
    fputs("[",f);
    for(int j=0;j<3;j++) {
        if(j) fputc(',',f);
        if(!native_samples_json(&r,f,j)) return 7;
    }
    fputs("]",f);
    if(fclose(f) || fclose(r.sample_file) || fclose(r.output_file)) return 8;
    free(r.times);free(r.states);
    printf("%zu %d\n",r.count,evaluations);return 0;
}
'''


class NativeSampleBufferTests(unittest.TestCase):
    def test_bounded_buffer_same_time_replacement_and_exact_output_bits(self):
        with tempfile.TemporaryDirectory(prefix='sample buffer ') as temporary:
            root = Path(temporary)
            compiler = os.environ.get('SIMULATION_NATIVE_CC') or shutil.which('gcc')
            if not compiler:
                if os.environ.get('SIMULATION_NATIVE_REQUIRE_TOOLCHAIN'):
                    self.fail('C compiler is required')
                self.skipTest('C compiler not available')
            (root / 'model.h').write_text('#define NSTATES 2\n#define NOUTPUTS 2\n'
                'int model_eval(double,const double*,double*,double*);\n')
            (root / 'test.c').write_text(HARNESS)
            executable = root / ('test.exe' if os.name == 'nt' else 'test')
            flags = ['-std=c11', '-O2', '-Wall', '-Wextra', '-Werror',
                     '-DNATIVE_SAMPLE_BUFFER_BYTES=96', '-D_POSIX_C_SOURCE=200809L']
            if os.name == 'nt':
                flags += ['-D__USE_MINGW_ANSI_STDIO=1']
            builder._command([compiler, *flags, '-I', str(root), '-I', str(ROOT / 'native/include'),
                str(root / 'test.c'), str(ROOT / 'native/runtime/sample_storage.c'),
                str(ROOT / 'native/runtime/json_numbers.c'), str(ROOT / 'native/encoding/ryu/d2s.c'),
                '-o', str(executable)], log=[])
            result = subprocess.run([str(executable), str(root / 'states.bin'), str(root / 'outputs.bin'),
                                     str(root / 'result.json')], capture_output=True, text=True, timeout=30)
            self.assertEqual(result.returncode, 0, result.stderr)
            self.assertEqual(result.stdout.strip(), '10003 10003')
            values = json.loads((root / 'result.json').read_bytes())
            self.assertEqual(values[0], list(range(10003)))
            self.assertEqual(values[1], [i + (.5 if i % 4 == 3 else 0) for i in range(10003)])
            self.assertTrue(all(struct.pack('<d', x) == struct.pack('<d', -0.) for x in values[2]))
            summary = result_storage.scan_blocks(root / 'states.bin', expected_columns=3)
            self.assertEqual(summary['blocks'], 2501)
            self.assertEqual(summary['samples'], 10003)
            self.assertFalse(summary['incompleteTail'] or summary['corrupt'])


class NativeSampleIntegrationTests(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        cls.temp = tempfile.TemporaryDirectory(prefix='native sample integration ')
        cls.addClassCleanup(cls.temp.cleanup)
        cls.root = Path(cls.temp.name)
        circuit = Circuit()
        mass = circuit.add('amesim_mecmas21', 'mass', mass=1, v0=1, x0=0, useFriction=1, stoptype=4)
        for i in (1, 2):
            free = circuit.add('amesim_f000', 'free' + str(i))
            circuit.connect(mass, 'port_' + str(i), free, 'port_1')
        cls.build = builder.build_native(compile_native_program(circuit.net))
        cls.addClassCleanup(cls.build.close)

    def setUp(self):
        self.case = Path(tempfile.mkdtemp(dir=self.root))
        patcher = patch.object(result_storage, 'RESULT_ROOT', self.case / 'simresults')
        patcher.start()
        self.addCleanup(patcher.stop)

    def test_complete_multiple_blocks_and_partial_tail_for_both_solvers(self):
        for method in ('RK45', 'BDF'):
            result = execute_native(self.build, SolveIVPConfig(t_stop=.25, method=method), .0001,
                                    run_dir=self.case / method)
            self.assertTrue(result['success'])
            summary = result['resultStorage']
            self.assertGreater(summary['states']['blocks'], 1)
            self.assertEqual(summary['states']['samples'], len(result['series']['time']))
            self.assertEqual(summary['outputs']['samples'], len(result['series']['time']))
            self.assertEqual(summary['states']['storedUntil'], .25)
            self.assertFalse(summary['states']['corrupt'] or summary['outputs']['corrupt'])
            # Inspect every persisted output against the existing JSON contract.
            columns = [[] for _ in result['series']]
            with (result_storage.result_archive_path(summary['id']) / 'outputs.bin').open('rb') as stream:
                while header := stream.read(40):
                    _, _, rows, width, _ = struct.unpack('<8sQQQQ', header)
                    for j in range(width):
                        columns[j].extend(struct.unpack('<' + 'd' * rows, stream.read(rows * 8)))
                    self.assertEqual(stream.read(8), b'COMMIT01')
            for stored, returned in zip(columns, result['series'].values(), strict=True):
                self.assertEqual(b''.join(struct.pack('<d', x) for x in stored),
                                 b''.join(struct.pack('<d', x) for x in returned))

    def test_cancel_flushes_short_block_and_solve_only_creates_no_archive(self):
        result = execute_native(self.build, SolveIVPConfig(t_stop=1), .001,
                                run_dir=self.case / 'cancel', cancel_check=lambda: True)
        self.assertEqual(result['status'], 'cancelled')
        self.assertEqual(result['resultStorage']['outputs']['samples'], 1)
        self.assertEqual(result['series']['time'], [0])
        count = len(list((self.case / 'simresults').glob('run-*')))
        execute_native(self.build, SolveIVPConfig(t_stop=.01), .001,
                       run_dir=self.case / 'solve-only', record_samples=False)
        self.assertEqual(len(list((self.case / 'simresults').glob('run-*'))), count)

    def test_exhausted_quota_stops_with_valid_partial_result(self):
        with patch.dict(os.environ, {'SIMULATION_RESULT_STORAGE_MB': '1'}):
            result = execute_native(self.build, SolveIVPConfig(t_stop=1), .00001,
                                    run_dir=self.case / 'full')
        self.assertFalse(result['success'])
        self.assertEqual(result['solverControl']['reason'], 'storage-quota')
        self.assertGreater(len(result['series']['time']), 0)
        self.assertLess(result['series']['time'][-1], 1)
        summary = result['resultStorage']
        self.assertEqual(summary['outputs']['samples'], len(result['series']['time']))
        self.assertLessEqual(sum(p.stat().st_size for p in (self.case / 'simresults').rglob('*') if p.is_file()),
                             1024**2)

    def test_killed_worker_retains_committed_states_without_inventing_outputs(self):
        helper = self.case / 'worker.py'
        width = len(self.build.manifest['stateKeys']) + 1
        helper.write_text('import sys,time,struct,zlib,json\nfrom pathlib import Path\n'
            'p=Path(sys.argv[1]); width=int(sys.argv[2])\n'
            'data=struct.pack("<"+"d"*(2*width),0,.1,*([1.]*(2*(width-1))))\n'
            'print(json.dumps({"phase":"storage-reserve","bytes":len(data)+48}),file=sys.stderr,flush=True)\n'
            'assert sys.stdin.readline()=="ok\\n"\n'
            'p.write_bytes(struct.pack("<8sQQQQ",b"SIMBLK01",0,2,width,zlib.crc32(data))+data+b"COMMIT01")\n'
            'print(json.dumps({"phase":"integrating","time":.1,"nfev":7,"acceptedSteps":2}),file=sys.stderr,flush=True)\n'
            'time.sleep(30)\n', encoding='utf-8')
        real_popen = subprocess.Popen
        def launch(command, **kwargs):
            samples = command[command.index('--sample-file') + 1]
            return real_popen([sys.executable, str(helper), samples, str(width)], **kwargs)
        progressed = []
        with patch.object(runner.subprocess, 'Popen', side_effect=launch), \
                patch.object(runner, '_WORKER_SHUTDOWN_GRACE_SECONDS', .3):
            result = execute_native(self.build, SolveIVPConfig(t_stop=1), .01,
                run_dir=self.case / 'killed', cancel_check=lambda: bool(progressed),
                progress_callback=lambda *args: progressed.append(args))
        self.assertEqual(result['status'], 'cancelled')
        self.assertEqual(result['series'], {})
        self.assertEqual(result['resultStorage']['states']['samples'], 2)
        self.assertEqual(result['resultStorage']['states']['storedUntil'], .1)
        self.assertEqual(result['resultStorage']['outputs']['samples'], 0)

    def test_storage_io_error_is_distinct_from_quota_exhaustion(self):
        with patch.object(result_storage.ResultArchive, 'reserve', side_effect=OSError('write denied')):
            result = execute_native(self.build, SolveIVPConfig(t_stop=.01), .001,
                                    run_dir=self.case / 'io-error')
        self.assertFalse(result['success'])
        self.assertEqual(result['solverControl']['reason'], 'storage-io')
        self.assertEqual(result['resultStorage']['states']['samples'], 0)

    def test_algebraic_model_private_state_has_explicit_column_metadata(self):
        circuit = Circuit()
        circuit.add('amesim_step0', 'signal')
        build = builder.build_native(compile_native_program(circuit.net))
        try:
            self.assertEqual(len(build.manifest['stateKeys']), 0)
            result = execute_native(build, SolveIVPConfig(t_stop=.02), .001,
                                    run_dir=self.case / 'algebraic')
        finally:
            build.close()
        self.assertTrue(result['success'])
        self.assertFalse(result['resultStorage']['states']['corrupt'])
        self.assertEqual(result['resultStorage']['states']['samples'], len(result['series']['time']))
        manifest = json.loads((result_storage.result_archive_path(result['resultStorage']['id']) / 'manifest.json').read_bytes())
        self.assertEqual(manifest['metadata']['stateColumns'], ['time', None])


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