import argparse
import json
import shutil
import tempfile
import unittest
from pathlib import Path

from scripts import generate_training_data as bridge


class GenerateTrainingDataTests(unittest.TestCase):
    def setUp(self):
        self.repo = Path(__file__).resolve().parents[1]
        self.temp = tempfile.TemporaryDirectory()
        self.root = Path(self.temp.name)
        self.dictionary_dir = self.root / "dictionary"
        shutil.copytree(self.repo / "dictionary", self.dictionary_dir)
        self.output_dir = self.root / "workspace_Data" / "data"
        self.output_dir.mkdir(parents=True)
        self.manifest = self.root / "workspace_Data" / "manifest.llm.json"
        self.manifest.write_text(json.dumps({"manifest_version": "0.0.0", "datasets": {}}))

    def tearDown(self):
        self.temp.cleanup()

    def args(self, **overrides):
        values = {
            "dictionary_dir": str(self.dictionary_dir),
            "output_dir": str(self.output_dir),
            "manifest": str(self.manifest),
            "dry_run": False,
            "sdg_host": None,
            "sdg_user": None,
            "sdg_key": None,
            "num_records": 1000,
            "timestamp": 1234567890,
        }
        values.update(overrides)
        return argparse.Namespace(**values)

    def test_seed_generation_covers_current_dictionary(self):
        artifacts = bridge.load_dictionary_artifacts(self.dictionary_dir)
        records = bridge.build_seed_records(artifacts, generated_at=123)
        actions = {record["metadata"]["action_id"] for record in records}
        self.assertEqual(len(actions), 13)
        self.assertIn("calculate", actions)
        self.assertIn("delete_file", actions)
        self.assertEqual(len(records), 260)

    def test_compact_output_uses_codebook_codes(self):
        artifacts = bridge.load_dictionary_artifacts(self.dictionary_dir)
        compact = bridge.compact_target(
            artifacts,
            "hash_string",
            {"text": "hello", "algorithm": "sha256"},
        )
        self.assertEqual(compact["a"], artifacts["action_to_code"]["hash_string"])
        self.assertIn(artifacts["field_to_code"]["text"], compact["p"])
        self.assertIn(artifacts["field_to_code"]["algorithm"], compact["p"])

    def test_empty_payload_action(self):
        artifacts = bridge.load_dictionary_artifacts(self.dictionary_dir)
        records = bridge.build_seed_records(artifacts, generated_at=123)
        uuid_records = [
            record for record in records if record["metadata"]["action_id"] == "generate_uuid"
        ]
        self.assertTrue(uuid_records)
        output = json.loads(uuid_records[0]["output"])
        compact = json.loads(uuid_records[0]["metadata"]["compact"])
        self.assertEqual(output["payload"], {})
        self.assertEqual(compact["p"], {})

    def test_split_includes_every_action(self):
        artifacts = bridge.load_dictionary_artifacts(self.dictionary_dir)
        records = bridge.build_seed_records(artifacts, generated_at=123)
        train, validation = bridge.split_records_by_action(records)
        train_actions = {record["metadata"]["action_id"] for record in train}
        validation_actions = {record["metadata"]["action_id"] for record in validation}
        self.assertEqual(train_actions, validation_actions)
        self.assertEqual(len(train_actions), 13)

    def test_dry_run_writes_nothing(self):
        summary = bridge.run_bridge(self.args(dry_run=True))
        self.assertTrue(summary["dry_run"])
        self.assertFalse((self.output_dir / "module_controller_intents_1234567890").exists())
        manifest = json.loads(self.manifest.read_text())
        self.assertEqual(manifest["datasets"], {})

    def test_manifest_patch_adds_dataset_without_damaging_existing_keys(self):
        bridge.run_bridge(self.args())
        manifest = json.loads(self.manifest.read_text())
        self.assertEqual(manifest["manifest_version"], "0.0.0")
        dataset = manifest["datasets"]["module_controller_intents_1234567890"]
        self.assertEqual(dataset["schema_profile"], "instruction")
        self.assertEqual(dataset["codebook_version"], "1.0.0")
        self.assertIn("codebook_checksum", dataset)
        self.assertTrue((self.output_dir / "module_controller_intents_1234567890" / "train.jsonl").exists())


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