use crate::VirtualMachineRuntime;
use crate::fs::FsState;
use crate::process::{Process, ProcessId, ProcessKind};
use crate::task::{Task, TaskId, TaskKind};
use prometeu_hal::app_mode::AppMode;
use prometeu_hal::telemetry::CertificationConfig;
use prometeu_vm::VirtualMachine;
use std::collections::HashMap;
use std::sync::Arc;
use std::sync::atomic::AtomicU32;

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct VmSessionId(pub TaskId);

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum VmSessionError {
    SessionAlreadyExists(TaskId),
    TaskProcessMismatch { task_id: TaskId, task_process: ProcessId, process_id: ProcessId },
    TaskKindMismatch { task_id: TaskId, task_kind: TaskKind, process_kind: ProcessKind },
    ProcessIsNotVmBacked(ProcessId),
}

pub struct VmSession {
    pub id: VmSessionId,
    pub task_id: TaskId,
    pub process_id: ProcessId,
    pub app_id: u32,
    pub title: String,
    pub app_version: String,
    pub app_mode: AppMode,
    pub vm: VirtualMachine,
    pub runtime: VirtualMachineRuntime,
    pub fs_state: FsState,
    pub open_files: HashMap<u32, String>,
    pub next_handle: u32,
}

impl VmSession {
    fn new(
        task: &Task,
        process: &Process,
        app_mode: AppMode,
        app_version: impl Into<String>,
        cap_config: Option<CertificationConfig>,
        logs_count: Arc<AtomicU32>,
    ) -> Self {
        Self {
            id: VmSessionId(task.id),
            task_id: task.id,
            process_id: process.id,
            app_id: task.app_id,
            title: task.title.clone(),
            app_version: app_version.into(),
            app_mode,
            vm: VirtualMachine::default(),
            runtime: VirtualMachineRuntime::new_with_log_counter(cap_config, logs_count),
            fs_state: FsState::Unmounted,
            open_files: HashMap::new(),
            next_handle: 1,
        }
    }

    pub fn clear_cartridge_service_state(&mut self) {
        self.open_files.clear();
        self.next_handle = 1;
    }
}

#[derive(Default)]
pub struct VmSessionRegistry {
    sessions: HashMap<VmSessionId, VmSession>,
}

impl VmSessionRegistry {
    pub fn new() -> Self {
        Self { sessions: HashMap::new() }
    }

    pub fn create_for_task(
        &mut self,
        task: &Task,
        process: &Process,
        app_version: impl Into<String>,
        cap_config: Option<CertificationConfig>,
        logs_count: Arc<AtomicU32>,
    ) -> Result<VmSessionId, VmSessionError> {
        if task.process_id != process.id {
            return Err(VmSessionError::TaskProcessMismatch {
                task_id: task.id,
                task_process: task.process_id,
                process_id: process.id,
            });
        }

        let app_mode = app_mode_for_task_and_process(task, process)?;
        let id = VmSessionId(task.id);
        if self.sessions.contains_key(&id) {
            return Err(VmSessionError::SessionAlreadyExists(task.id));
        }

        self.sessions.insert(
            id,
            VmSession::new(task, process, app_mode, app_version, cap_config, logs_count),
        );
        Ok(id)
    }

    pub fn get(&self, id: VmSessionId) -> Option<&VmSession> {
        self.sessions.get(&id)
    }

    pub fn get_mut(&mut self, id: VmSessionId) -> Option<&mut VmSession> {
        self.sessions.get_mut(&id)
    }

    pub fn for_task(&self, task_id: TaskId) -> Option<&VmSession> {
        self.get(VmSessionId(task_id))
    }

    pub fn for_task_mut(&mut self, task_id: TaskId) -> Option<&mut VmSession> {
        self.get_mut(VmSessionId(task_id))
    }

    pub fn resident_game(&self) -> Option<&VmSession> {
        self.sessions.values().find(|session| session.app_mode == AppMode::Game)
    }

    pub fn resident_game_for_app(&self, app_id: u32) -> Option<&VmSession> {
        self.sessions
            .values()
            .find(|session| session.app_mode == AppMode::Game && session.app_id == app_id)
    }

    pub fn len(&self) -> usize {
        self.sessions.len()
    }

    pub fn is_empty(&self) -> bool {
        self.sessions.is_empty()
    }

    pub fn clear(&mut self) {
        self.sessions.clear();
    }
}

fn app_mode_for_task_and_process(
    task: &Task,
    process: &Process,
) -> Result<AppMode, VmSessionError> {
    match (task.kind, process.kind) {
        (TaskKind::Game, ProcessKind::VmGame) => Ok(AppMode::Game),
        (TaskKind::Shell, ProcessKind::VmShell) => Ok(AppMode::Shell),
        (_, ProcessKind::NativeShell) => Err(VmSessionError::ProcessIsNotVmBacked(process.id)),
        _ => Err(VmSessionError::TaskKindMismatch {
            task_id: task.id,
            task_kind: task.kind,
            process_kind: process.kind,
        }),
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::process::Process;
    use crate::task::Task;

    #[test]
    fn creates_vm_game_session_for_vm_game_task() {
        let process = Process::new(ProcessId(1), 42, "Stress", ProcessKind::VmGame);
        let task = Task::new(TaskId(7), process.id, 42, "Stress", TaskKind::Game);
        let mut registry = VmSessionRegistry::new();

        let id = registry
            .create_for_task(&task, &process, "", None, Arc::new(AtomicU32::new(0)))
            .expect("session should create");
        let session = registry.get(id).expect("session should exist");

        assert_eq!(session.id, VmSessionId(task.id));
        assert_eq!(session.task_id, task.id);
        assert_eq!(session.process_id, process.id);
        assert_eq!(session.app_id, 42);
        assert_eq!(session.title, "Stress");
        assert_eq!(session.app_mode, AppMode::Game);
        assert_eq!(session.runtime.current_app_id, 0);
        assert_eq!(session.fs_state, FsState::Unmounted);
        assert_eq!(session.next_handle, 1);
    }

    #[test]
    fn creates_vm_shell_session_for_vm_shell_task() {
        let process = Process::new(ProcessId(1), 7, "Shell", ProcessKind::VmShell);
        let task = Task::new(TaskId(3), process.id, 7, "Shell", TaskKind::Shell);
        let mut registry = VmSessionRegistry::new();

        let id = registry
            .create_for_task(&task, &process, "1.0", None, Arc::new(AtomicU32::new(0)))
            .expect("session should create");
        let session = registry.get(id).expect("session should exist");

        assert_eq!(session.app_mode, AppMode::Shell);
        assert_eq!(session.app_version, "1.0");
    }

    #[test]
    fn rejects_native_shell_process() {
        let process = Process::new(ProcessId(1), 7, "Shell", ProcessKind::NativeShell);
        let task = Task::new(TaskId(3), process.id, 7, "Shell", TaskKind::Shell);
        let mut registry = VmSessionRegistry::new();

        assert_eq!(
            registry.create_for_task(&task, &process, "", None, Arc::new(AtomicU32::new(0))),
            Err(VmSessionError::ProcessIsNotVmBacked(process.id))
        );
    }

    #[test]
    fn finds_resident_game_by_app_id() {
        let process = Process::new(ProcessId(1), 42, "Stress", ProcessKind::VmGame);
        let task = Task::new(TaskId(7), process.id, 42, "Stress", TaskKind::Game);
        let mut registry = VmSessionRegistry::new();
        registry
            .create_for_task(&task, &process, "", None, Arc::new(AtomicU32::new(0)))
            .expect("session should create");

        assert_eq!(
            registry.resident_game_for_app(42).map(|session| session.task_id),
            Some(task.id)
        );
        assert_eq!(registry.resident_game_for_app(43).map(|session| session.task_id), None);
    }

    #[test]
    fn session_runtime_state_is_independent_per_session() {
        let game_process = Process::new(ProcessId(1), 42, "Stress", ProcessKind::VmGame);
        let game_task = Task::new(TaskId(7), game_process.id, 42, "Stress", TaskKind::Game);
        let shell_process = Process::new(ProcessId(2), 9, "Shell", ProcessKind::VmShell);
        let shell_task = Task::new(TaskId(8), shell_process.id, 9, "Shell", TaskKind::Shell);
        let mut registry = VmSessionRegistry::new();

        registry
            .create_for_task(&game_task, &game_process, "", None, Arc::new(AtomicU32::new(0)))
            .expect("game session should create");
        registry
            .create_for_task(&shell_task, &shell_process, "", None, Arc::new(AtomicU32::new(0)))
            .expect("shell session should create");

        registry
            .for_task_mut(game_task.id)
            .expect("game session should exist")
            .runtime
            .tick_index = 11;
        registry
            .for_task_mut(shell_task.id)
            .expect("shell session should exist")
            .runtime
            .tick_index = 3;
        registry
            .for_task_mut(game_task.id)
            .expect("game session should exist")
            .open_files
            .insert(1, "/game.save".to_string());

        assert_eq!(
            registry.for_task(game_task.id).map(|session| session.runtime.tick_index),
            Some(11)
        );
        assert_eq!(
            registry.for_task(shell_task.id).map(|session| session.runtime.tick_index),
            Some(3)
        );
        assert!(
            registry
                .for_task(shell_task.id)
                .expect("shell session should exist")
                .open_files
                .is_empty()
        );
    }
}