"""CPU checks for leakage guards, whole-project token masks, and streamed LoRA math."""
import importlib.util
import json
import tempfile
import unittest
from pathlib import Path
from unittest import mock

ROOT = Path(__file__).resolve().parents[1]

def module(name):
    spec = importlib.util.spec_from_file_location(name, ROOT / 'scripts' / (name + '.py'))
    obj = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(obj)
    return obj

train = module('train')
export = module('export_model')

class FakeTokenizer:
    def apply_chat_template(self, messages, tokenize, add_generation_prompt):
        prefix = [1, 2, 3]
        return prefix if add_generation_prompt else prefix + [4, 5, 6, 7]

class TrainingTests(unittest.TestCase):
    def row(self, instruction='one', marker='one'):
        return {'instruction': instruction, 'output': json.dumps({'files': {'project.godot': marker, 'main.gd': 'extends Node'}})}

    def test_assistant_mask_and_whole_project_rejection(self):
        rows, rejected = train.tokenize_rows([self.row()], FakeTokenizer(), 7, 'train')
        self.assertEqual(rows[0]['labels'], [-100, -100, -100, 4, 5, 6, 7])
        self.assertEqual(rejected, [])
        with self.assertRaisesRegex(ValueError, 'No complete'):
            train.tokenize_rows([self.row()], FakeTokenizer(), 6, 'train')

    def test_blind_and_output_leakage(self):
        blind = [{'instruction': f'blind {i}'} for i in range(500)]
        train.check_splits([self.row('training')], [self.row('development', 'two')], blind)
        with self.assertRaisesRegex(ValueError, 'contamination'):
            train.check_splits([self.row('blind 1')], [self.row('development', 'two')], blind)
        with self.assertRaisesRegex(ValueError, 'Identical project'):
            train.check_splits([self.row('training')], [self.row('development')], blind)
        with self.assertRaisesRegex(ValueError, '500'):
            train.check_splits([self.row()], [self.row('dev', 'other')], blind[:-1])

    def test_only_intact_complete_checkpoint_resumes(self):
        with tempfile.TemporaryDirectory() as d:
            root = Path(d)
            old = root / 'checkpoint-1'; old.mkdir()
            new = root / 'checkpoint-2'; new.mkdir()
            (old / 'adapter').write_text('complete')
            train.atomic_json(old / 'COMPLETE.json', {'sha256': {'adapter': train.digest(old / 'adapter')}})
            self.assertEqual(train.latest_complete_checkpoint(root), old)
            (new / 'adapter').write_text('complete')
            train.atomic_json(new / 'COMPLETE.json', {'sha256': {'adapter': train.digest(new / 'adapter')}})
            (new / 'adapter').write_text('interrupted')
            self.assertEqual(train.latest_complete_checkpoint(root), old)

class ExportTests(unittest.TestCase):
    def test_converter_installs_only_into_overlay_and_preserves_torch(self):
        versions = {'python': '3.12.12', 'packages': {'torch': '2.11.0+cu130'}}
        with tempfile.TemporaryDirectory() as d:
            workspace = Path(d)
            with mock.patch('importlib.metadata.version', return_value='2.11.0+cu130'), \
                 mock.patch.object(export, 'run') as run, \
                 mock.patch.object(export.subprocess, 'check_output', return_value=json.dumps(versions)):
                python, env, actual = export.converter_environment(workspace / 'llama.cpp', workspace)
            commands = [call.args[0] for call in run.call_args_list]
            self.assertIn('--system-site-packages', commands[0])
            self.assertIn('--without-pip', commands[0])
            self.assertEqual(commands[1][:2], [python, '-c'])
            self.assertIn('sys.prefix != sys.base_prefix', commands[1][2])
            self.assertIn('sysconfig.get_path', commands[1][2])
            self.assertIn("find_spec('pip')", commands[1][2])
            self.assertEqual(commands[2][:4], [python, '-m', 'pip', 'install'])
            self.assertEqual(python, workspace / 'converter-venv/bin/python')
            self.assertEqual((workspace / 'converter-constraints.txt').read_text(), 'torch===2.11.0+cu130\n')
            self.assertEqual(env['PIP_REQUIRE_VIRTUALENV'], 'true')
            self.assertEqual(env['OMP_NUM_THREADS'], '1')
            self.assertEqual(actual, versions)
        with mock.patch('importlib.metadata.version', return_value='2.10.0'), \
             mock.patch.object(export, 'run') as run:
            with self.assertRaisesRegex(RuntimeError, 'requires installed torch'):
                export.converter_environment(Path('/unused'), Path('/unused'))
            run.assert_not_called()

    def test_row_chunked_merge_matches_dense_math(self):
        import torch
        torch.manual_seed(3)
        w = torch.randn(513, 257, dtype=torch.bfloat16)
        a = torch.randn(32, 257) * .01
        b = torch.randn(513, 32) * .01
        merged = export.merge_tensor(w, {'A': a, 'B': b}, 2, row_chunk=41)
        expected = (w.float() + (b @ a) * 2).half()
        torch.testing.assert_close(merged, expected, rtol=.001, atol=.001)
        self.assertEqual(merged.dtype, torch.float16)

    def test_adapter_names_and_nonfinite_are_rejected(self):
        import torch
        from safetensors.torch import save_file
        with tempfile.TemporaryDirectory() as d:
            path = Path(d) / 'adapter.safetensors'
            save_file({'base_model.model.model.layers.0.self_attn.q_proj.lora_A.weight': torch.zeros(32, 3),
                       'base_model.model.model.layers.0.self_attn.q_proj.lora_B.weight': torch.zeros(4, 32)}, path)
            self.assertIn('model.layers.0.self_attn.q_proj.weight', export.lora_pairs(path))
            save_file({'unknown': torch.zeros(2)}, path)
            with self.assertRaisesRegex(ValueError, 'Unsupported'):
                export.lora_pairs(path)
        with self.assertRaisesRegex(ValueError, 'Nonfinite'):
            export.merge_tensor(torch.zeros(3, 3), {'A': torch.ones(1, 3), 'B': torch.full((3, 1), float('inf'))}, 2)

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