from __future__ import annotations

import json
from pathlib import Path
from typing import Any, Dict, Optional

from .dataset_loader import DatasetLoader
from .injection_engine import InjectionEngine
from .policy_engine import PolicyEngine
from .scenario_types import HiddenEpisodeState


class ScenarioGenerator:
    def __init__(
        self,
        dataset_loader: DatasetLoader,
        injection_engine: InjectionEngine,
        policy_engine: PolicyEngine,
        config_dir: Optional[Path] = None,
    ) -> None:
        self._dataset_loader = dataset_loader
        self._injection_engine = injection_engine
        self._policy_engine = policy_engine
        self._config_dir = config_dir or Path(__file__).resolve().parent / "data" / "config"
        self._attacker_profiles = self._load_attacker_profiles()

    def generate_episode(self, difficulty: str, seed: Optional[int] = None) -> HiddenEpisodeState:
        scenario = self._dataset_loader.get_scenario(difficulty=difficulty, seed=seed)
        return self._build_hidden_state(scenario)

    def generate_episode_by_id(self, difficulty: str, scenario_id: str) -> HiddenEpisodeState:
        scenario = self._dataset_loader.get_scenario_by_id(difficulty=difficulty, scenario_id=scenario_id)
        return self._build_hidden_state(scenario)

    def _build_hidden_state(self, scenario) -> HiddenEpisodeState:
        attack_metadata = self._injection_engine.build_attack_metadata(scenario)
        attacker_profile = self._attacker_profiles.get("default_attacker", {})
        attack_metadata["attacker_profile"] = attacker_profile
        activation = dict(scenario.metadata.get("activation", {}))
        attack_metadata["activation"] = activation
        attack_active = activation.get("mode", "immediate") == "immediate"
        attack_triggered_step = 0 if attack_active else None
        attack_trigger_reason = "initial_immediate_attack" if attack_active else None

        return HiddenEpisodeState(
            scenario=scenario,
            policy=self._policy_engine.load_policy_rules(),
            attack_metadata=attack_metadata,
            attack_active=attack_active,
            attack_triggered_step=attack_triggered_step,
            attack_trigger_reason=attack_trigger_reason,
            last_action_result=f"Scenario {scenario.scenario_id} loaded. {scenario.benign_task}",
        )

    def _load_attacker_profiles(self) -> Dict[str, Dict[str, Any]]:
        path = self._config_dir / "attacker_profiles.json"
        with path.open("r", encoding="utf-8-sig") as handle:
            return json.load(handle)