import json
import os
from typing import Tuple, Dict, Any
from models import Observation, Action, Reward, LogEntry, ActionType

class SOCEnvironment:
    def __init__(self):
        # Initialize the internal state variables
        self.current_time = "2026-03-31T00:00:00Z"
        self.active_alerts = []
        self.system_status = {"database_online": True, "backup_server_online": True}
        
        # Track agent actions
        self.blocked_ips = set()
        self.isolated_hosts = set()
        self.current_task_id = None
        self.step_count = 0
        
        # Load the synthetic BOTS data safely
        data_path = os.path.join(os.path.dirname(__file__), 'data', 'mock_bots.json')
        with open(data_path, 'r') as f:
            self.mock_data = json.load(f)

    def state(self) -> Observation:
        """Returns the current state of the environment perfectly formatted to our Pydantic model."""
        return Observation(
            current_time=self.current_time,
            active_alerts=self.active_alerts,
            system_status=self.system_status
        )

    def reset(self, task_id: str = "task_1_triage") -> Observation:
        """Resets the environment for a specific task and loads the appropriate logs."""
        self.current_task_id = task_id
        self.step_count = 0
        self.blocked_ips.clear()
        self.isolated_hosts.clear()
        self.system_status = {"database_online": True, "backup_server_online": True}
        
        # Map the task_id from openenv.yaml to our JSON keys
        task_mapping = {
            "task_1_triage": "scenario_1_triage",
            "task_2_false_positive": "scenario_2_false_positive",
            "task_3_kill_chain": "scenario_3_kill_chain"
        }
        
        scenario_key = task_mapping.get(task_id, "scenario_1_triage")
        raw_logs = self.mock_data.get(scenario_key, [])
        
        # Convert raw JSON dicts into our strict Pydantic LogEntry models
        self.active_alerts = [LogEntry(**log) for log in raw_logs]
        
        # Set the starting time based on the first log
        if self.active_alerts:
            self.current_time = self.active_alerts[0].timestamp

        return self.state()


    def step(self, action: Action) -> Tuple[Observation, Reward, bool, Dict[str, Any]]:
        """Executes the agent's action, updates the state, and returns the result."""
        self.step_count += 1
        score_delta = 0.0
        done = False
        info = {"action_taken": action.action_type, "target": action.target_ip, "msg": ""}

        # 1. Process the Action
        if action.action_type == ActionType.BLOCK_IP:
            if action.target_ip:
                self.blocked_ips.add(action.target_ip)
                info["msg"] = f"Successfully blocked IP: {action.target_ip}"
            else:
                score_delta = -0.1  # Penalty for missing target
                info["msg"] = "Failed: Target IP required for BLOCK_IP."

        elif action.action_type == ActionType.ISOLATE_HOST:
            if action.target_ip:
                self.isolated_hosts.add(action.target_ip)
                info["msg"] = f"Isolated Host: {action.target_ip}"
                
                # Hardcoded consequence: Isolating the core backup server takes it offline
                if action.target_ip == "10.0.0.250": 
                    self.system_status["backup_server_online"] = False
                    score_delta = -0.5 # Immediate penalty for breaking a critical system
            else:
                score_delta = -0.1

        elif action.action_type == ActionType.DISMISS_ALERT:
            info["msg"] = "Alerts dismissed."

        elif action.action_type == ActionType.INVESTIGATE:
            info["msg"] = f"Pulled deeper telemetry for {action.target_ip}."

        elif action.action_type == ActionType.ESCALATE_TO_HUMAN:
            info["msg"] = "Escalated to Tier 2 Human Analyst."
            done = True # Escalation immediately ends the agent's turn

        # 2. Check Episode Termination Conditions (Max 10 steps to prevent infinite loops)
        if self.step_count >= 10:
            done = True
            info["msg"] = "Max steps reached. Operation timed out."

        # Task-specific early termination
        if self.current_task_id == "task_1_triage" and action.action_type == ActionType.BLOCK_IP:
            done = True
        elif self.current_task_id == "task_2_false_positive" and action.action_type in [ActionType.DISMISS_ALERT, ActionType.ISOLATE_HOST]:
            done = True

        return self.state(), Reward(score_delta=score_delta), done, info