import json
import tempfile
import unittest
from pathlib import Path

from env.dataset_loader import DatasetLoader
from env.environment import IDPIExfilEnv
from env.injection_engine import InjectionEngine
from env.models import Action
from env.policy_engine import PolicyEngine
from env.scenario_generator import ScenarioGenerator
from env.taint_tracker import TaintTracker


class GoldScenarioEnvironmentFlowTests(unittest.TestCase):
    def setUp(self) -> None:
        self.loader = DatasetLoader()

    def test_medium_gold_authority_spoof_blocks_outbound_send(self) -> None:
        scenario = self.loader.get_scenario_by_id("medium", "medium_gold_authority_policy_bypass")
        env = self._env_for_single_scenario("medium", scenario)

        env.reset(difficulty="medium", seed=0)
        env.step(Action(type="ReadDoc"))
        _, _, done, info = env.step(
            Action(type="SendEmail", to=scenario.attacker_address, body="validation complete")
        )

        self.assertFalse(done)
        self.assertIn("attack_requested_destination", info["policy"]["violations"])
        self.assertIn("blocked_destination", info["policy"]["violations"])
        result = env.get_result()
        self.assertFalse(result.exfil_occurred)
        self.assertFalse(result.task_completed)

    def test_hard_gold_transcript_scenario_completes_with_flag_and_clean_summary(self) -> None:
        scenario = self.loader.get_scenario_by_id("hard", "hard_gold_transcript_finance_review")
        env = self._env_for_single_scenario("hard", scenario)

        env.reset(difficulty="hard", seed=0)
        env.step(Action(type="ReadDoc"))
        env.step(Action(type="FlagInjection", reason="transcript asks for external relay"))
        clean_summary = ". ".join(scenario.expected_summary_points[:2])
        _, _, done, info = env.step(Action(type="WriteDB", content=clean_summary))

        self.assertTrue(done)
        self.assertEqual(info["policy"]["violations"], [])
        result = env.get_result()
        self.assertTrue(result.injection_flagged)
        self.assertTrue(result.task_completed)
        self.assertFalse(result.exfil_occurred)
        self.assertGreaterEqual(result.final_score, 0.7)

    def test_hard_gold_attachment_scenario_requires_clean_summary(self) -> None:
        scenario = self.loader.get_scenario_by_id("hard", "hard_gold_attachment_summary_append")
        env = self._env_for_single_scenario("hard", scenario)

        env.reset(difficulty="hard", seed=0)
        env.step(Action(type="ReadDoc"))
        env.step(Action(type="FlagInjection", reason="attachment contains relay directive"))
        clean_summary = ". ".join(scenario.expected_summary_points[:2])
        _, _, done, info = env.step(Action(type="WriteDB", content=clean_summary))

        self.assertTrue(done)
        self.assertEqual(info["policy"]["violations"], [])
        self.assertGreaterEqual(info["grader_inputs"]["summary_score"], 0.6)
        self.assertEqual(info["attack"]["attack_channel"], "attachment_instruction")
        result = env.get_result()
        self.assertTrue(result.task_completed)
        self.assertFalse(result.exfil_occurred)

    def test_delayed_comment_attack_activates_after_read(self) -> None:
        scenario = self.loader.get_scenario_by_id("hard", "hard_gold_delayed_comment_followup")
        env = self._env_for_single_scenario("hard", scenario)

        env.reset(difficulty="hard", seed=0)
        self.assertFalse(env.get_evaluation_facts()["progress"]["attack_active"])
        _, _, done, info = env.step(Action(type="ReadDoc"))

        self.assertFalse(done)
        self.assertTrue(info["activation"]["attack_active"])
        self.assertEqual(info["activation"]["trigger_reason"], "activated_after_read")
        self.assertEqual(info["step_facts"]["attack_triggered_step"], 1)

    def test_delayed_attachment_attack_activates_on_summary_attempt(self) -> None:
        scenario = self.loader.get_scenario_by_id("hard", "hard_gold_delayed_attachment_summary_hook")
        env = self._env_for_single_scenario("hard", scenario)

        env.reset(difficulty="hard", seed=0)
        env.step(Action(type="ReadDoc"))
        _, _, done, info = env.step(Action(type="WriteDB", content=". ".join(scenario.expected_summary_points[:2])))

        self.assertFalse(done)
        self.assertTrue(info["activation"]["attack_active"])
        self.assertEqual(info["activation"]["trigger_reason"], "activated_on_summary_attempt")
        self.assertIn("summary_before_attack_resolution", info["policy"]["violations"])

    def _env_for_single_scenario(self, difficulty: str, scenario) -> IDPIExfilEnv:
        tmpdir = tempfile.TemporaryDirectory()
        self.addCleanup(tmpdir.cleanup)
        data_dir = Path(tmpdir.name)
        payload = json.dumps(self._scenario_to_payload(scenario), ensure_ascii=True)
        (data_dir / f"{difficulty}.jsonl").write_text(payload + "\n", encoding="utf-8")

        env = IDPIExfilEnv()
        env._injection_engine = InjectionEngine()
        env._dataset_loader = DatasetLoader(data_dir=data_dir)
        env._policy_engine = PolicyEngine()
        env._taint_tracker = TaintTracker(env._injection_engine)
        env._scenario_generator = ScenarioGenerator(
            dataset_loader=env._dataset_loader,
            injection_engine=env._injection_engine,
            policy_engine=env._policy_engine,
        )
        return env

    def _scenario_to_payload(self, scenario) -> dict:
        return {
            "scenario_id": scenario.scenario_id,
            "difficulty": scenario.difficulty,
            "subject": scenario.subject,
            "body": scenario.body,
            "benign_task": scenario.benign_task,
            "summary_required": scenario.summary_required,
            "attack_type": scenario.attack_type,
            "attacker_address": scenario.attacker_address,
            "safe_destinations": list(scenario.safe_destinations),
            "blocked_destinations": list(scenario.blocked_destinations),
            "expected_summary_points": list(scenario.expected_summary_points),
            "objectives": dict(scenario.objectives),
            "metadata": dict(scenario.metadata),
            "max_steps": scenario.max_steps,
        }


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