use std::{collections::HashMap, sync::Arc};

use anyhow::{Result, bail};
use fxprof_processed_profile::LibraryInfo;

use crate::emulator::{GameInfo, InlineStack};

pub struct ProgramState {
    info: Arc<GameInfo>,
    call_stacks: HashMap<u16, Vec<StackFrame>>,
    context_stack: Vec<u16>,
}

pub enum StackFrame {
    Address(u32),
    Label(Arc<String>),
}

pub const RESET_CODE: u16 = 0xfff0;
impl ProgramState {
    pub async fn new(info: Arc<GameInfo>) -> Self {
        let mut call_stacks = HashMap::new();
        call_stacks.insert(RESET_CODE, vec![StackFrame::Address(0xfffffff0)]);
        Self {
            info,
            call_stacks,
            context_stack: vec![RESET_CODE],
        }
    }

    pub fn name(&self) -> &str {
        self.info.name()
    }

    pub fn library_info(&self) -> &LibraryInfo {
        self.info.library_info()
    }

    pub fn current_stack(&self) -> Option<(u16, &[StackFrame])> {
        let code = self.context_stack.last()?;
        let call_stack = self.call_stacks.get(code)?;
        Some((*code, call_stack))
    }

    pub fn track_call(&mut self, address: u32) -> Result<()> {
        let Some(code) = self.context_stack.last() else {
            bail!("How did we call anything when we're halted?");
        };
        let Some(stack) = self.call_stacks.get_mut(code) else {
            bail!("missing stack {code:04x}");
        };
        stack.push(StackFrame::Address(address));
        Ok(())
    }

    pub fn track_return(&mut self) -> Result<()> {
        let Some(code) = self.context_stack.last() else {
            bail!("how did we return when we're halted?");
        };
        let Some(stack) = self.call_stacks.get_mut(code) else {
            bail!("missing stack {code:04x}");
        };
        // just popping the inline frames first
        while stack
            .pop_if(|f| matches!(f, StackFrame::Label(_)))
            .is_some()
        {}
        if stack.pop().is_none() {
            bail!("returned from {code:04x} but stack was empty");
        }
        if stack.is_empty() {
            bail!("returned to oblivion");
        }
        Ok(())
    }

    pub fn track_halt(&mut self) -> Result<()> {
        let Some(RESET_CODE) = self.context_stack.pop() else {
            bail!("halted when not in an interrupt");
        };
        Ok(())
    }

    pub fn track_interrupt(&mut self, code: u16, address: u32) -> Result<()> {
        // if the CPU was halted before, wake it up now
        if self.context_stack.is_empty() {
            self.context_stack.push(RESET_CODE);
        }

        self.context_stack.push(code);
        if self
            .call_stacks
            .insert(code, vec![StackFrame::Address(address)])
            .is_some()
        {
            bail!("{code:04x} fired twice");
        }
        Ok(())
    }

    pub fn track_reti(&mut self) -> Result<()> {
        let Some(code) = self.context_stack.pop() else {
            bail!("RETI when halted");
        };
        if code == RESET_CODE {
            bail!("RETI when not in interrupt");
        }
        if self.call_stacks.remove(&code).is_none() {
            bail!("{code:04x} popped but never called");
        }
        Ok(())
    }

    pub fn track_inline_stack(&mut self, inline_stack: InlineStack) {
        let Some(code) = self.context_stack.last() else {
            return;
        };
        let Some(call_stack) = self.call_stacks.get_mut(code) else {
            return;
        };
        while call_stack
            .pop_if(|f| matches!(f, StackFrame::Label(_)))
            .is_some()
        {}
        for label in inline_stack.iter() {
            call_stack.push(StackFrame::Label(label.clone()));
        }
    }
}