use smallvec::SmallVec;

use rustc_middle::mir::interpret::{AllocId, AllocRange};
use rustc_span::{Span, SpanData};
use rustc_target::abi::Size;

use crate::helpers::CurrentSpan;
use crate::stacked_borrows::{err_sb_ub, AccessKind};
use crate::*;

use rustc_middle::mir::interpret::InterpError;

#[derive(Clone, Debug)]
pub struct AllocHistory {
    // The time tags can be compressed down to one bit per event, by just storing a Vec<u8>
    // where each bit is set to indicate if the event was a creation or a retag
    current_time: usize,
    creations: smallvec::SmallVec<[Event; 2]>,
    invalidations: smallvec::SmallVec<[Event; 1]>,
    protectors: smallvec::SmallVec<[Protection; 1]>,
}

#[derive(Clone, Debug)]
struct Protection {
    orig_tag: SbTag,
    tag: SbTag,
    span: Span,
}

#[derive(Clone, Debug)]
struct Event {
    parent: Option<SbTag>,
    tag: SbTag,
    range: AllocRange,
    span: Span,
}

pub enum TagHistory {
    Tagged {
        tag: SbTag,
        created: (AllocRange, SpanData),
        invalidated: Option<(AllocRange, SpanData)>,
        protected: Option<(SbTag, SpanData, SpanData)>,
    },
}

impl AllocHistory {
    pub fn new() -> Self {
        Self {
            current_time: 0,
            creations: SmallVec::new(),
            invalidations: SmallVec::new(),
            protectors: SmallVec::new(),
        }
    }

    pub fn log_creation(
        &mut self,
        parent: Option<SbTag>,
        tag: SbTag,
        range: AllocRange,
        current_span: &mut CurrentSpan<'_, '_, '_>,
    ) {
        let span = current_span.get();
        self.creations.push(Event { parent, tag, range, span });
        self.current_time += 1;
    }

    pub fn log_invalidation(
        &mut self,
        tag: SbTag,
        range: AllocRange,
        current_span: &mut CurrentSpan<'_, '_, '_>,
    ) {
        let span = current_span.get();
        self.invalidations.push(Event { parent: None, tag, range, span });
        self.current_time += 1;
    }

    pub fn log_protector(
        &mut self,
        orig_tag: SbTag,
        tag: SbTag,
        current_span: &mut CurrentSpan<'_, '_, '_>,
    ) {
        let span = current_span.get();
        self.protectors.push(Protection { orig_tag, tag, span });
        self.current_time += 1;
    }

    pub fn get_logs_relevant_to(
        &self,
        tag: SbTag,
        protector_tag: Option<SbTag>,
    ) -> Option<TagHistory> {
        let protected = protector_tag
            .and_then(|protector| {
                self.protectors.iter().find_map(|protection| {
                    if protection.tag == protector {
                        Some((protection.orig_tag, protection.span.data()))
                    } else {
                        None
                    }
                })
            })
            .and_then(|(tag, call_span)| {
                self.creations.iter().rev().find_map(|event| {
                    if event.tag == tag {
                        Some((event.parent?, event.span.data(), call_span))
                    } else {
                        None
                    }
                })
            });

        let get_matching = |events: &[Event]| {
            events.iter().rev().find_map(|event| {
                if event.tag == tag { Some((event.range, event.span.data())) } else { None }
            })
        };
        Some(TagHistory::Tagged {
            tag,
            created: get_matching(&self.creations)?,
            invalidated: get_matching(&self.invalidations),
            protected,
        })
    }

    /// Report a descriptive error when `new` could not be granted from `derived_from`.
    pub fn grant_error<'tcx>(
        &self,
        derived_from: ProvenanceExtra,
        new: Item,
        alloc_id: AllocId,
        alloc_range: AllocRange,
        error_offset: Size,
        stack: &Stack,
    ) -> InterpError<'tcx> {
        let action = format!(
            "trying to reborrow from {derived_from:?} for {new_perm:?} permission at {alloc_id:?}[{offset:#x}]",
            new_perm = new.perm(),
            offset = error_offset.bytes(),
        );
        err_sb_ub(
            format!("{}{}", action, error_cause(stack, derived_from)),
            Some(operation_summary("a reborrow", alloc_id, alloc_range)),
            derived_from.and_then(|derived_from| self.get_logs_relevant_to(derived_from, None)),
        )
    }

    /// Report a descriptive error when `access` is not permitted based on `tag`.
    pub fn access_error<'tcx>(
        &self,
        access: AccessKind,
        tag: ProvenanceExtra,
        alloc_id: AllocId,
        alloc_range: AllocRange,
        error_offset: Size,
        stack: &Stack,
    ) -> InterpError<'tcx> {
        let action = format!(
            "attempting a {access} using {tag:?} at {alloc_id:?}[{offset:#x}]",
            offset = error_offset.bytes(),
        );
        err_sb_ub(
            format!("{}{}", action, error_cause(stack, tag)),
            Some(operation_summary("an access", alloc_id, alloc_range)),
            tag.and_then(|tag| self.get_logs_relevant_to(tag, None)),
        )
    }
}

fn operation_summary(
    operation: &'static str,
    alloc_id: AllocId,
    alloc_range: AllocRange,
) -> String {
    format!("this error occurs as part of {operation} at {alloc_id:?}{alloc_range:?}")
}

fn error_cause(stack: &Stack, prov_extra: ProvenanceExtra) -> &'static str {
    if let ProvenanceExtra::Concrete(tag) = prov_extra {
        if (0..stack.len())
            .map(|i| stack.get(i).unwrap())
            .any(|item| item.tag() == tag && item.perm() != Permission::Disabled)
        {
            ", but that tag only grants SharedReadOnly permission for this location"
        } else {
            ", but that tag does not exist in the borrow stack for this location"
        }
    } else {
        ", but no exposed tags have suitable permission in the borrow stack for this location"
    }
}