use std::cell::OnceCell;

use krilla::geom as kg;
use krilla::tagging::{BBox, Identifier, Node, TagKind};
use typst_library::layout::{Abs, Point, Rect};

use crate::convert::FrameContext;
use crate::tags::context::figure::build_figure;
use crate::tags::context::grid::build_grid;
use crate::tags::context::table::build_table;
use crate::tags::groups::GroupKind;
use crate::tags::tree::ResolvedTextAttrs;
use crate::tags::tree::Tree;
use crate::tags::util::{Id, IdVec};
use crate::util::AbsExt;

pub use crate::tags::context::figure::FigureCtx;
pub use crate::tags::context::grid::GridCtx;
pub use crate::tags::context::list::ListCtx;
pub use crate::tags::context::outline::OutlineCtx;
pub use crate::tags::context::table::TableCtx;

mod figure;
mod grid;
mod list;
mod outline;
mod table;

pub type TableId = Id<TableCtx>;
pub type GridId = Id<GridCtx>;
pub type FigureId = Id<FigureCtx>;
pub type ListId = Id<ListCtx>;
pub type OutlineId = Id<OutlineCtx>;
pub type BBoxId = Id<BBoxCtx>;
pub type TagId = Id<TagKind>;
pub type AnnotationId = Id<krilla::annotation::Annotation>;

pub struct Tags {
    pub in_tiling: bool,
    pub tree: Tree,
    /// A list of placeholders for annotations in the tag tree.
    pub annotations: Annotations,
}

impl Tags {
    pub fn new(tree: Tree) -> Self {
        Self {
            in_tiling: false,
            tree,
            annotations: Annotations::new(),
        }
    }

    pub fn push_leaf(&mut self, id: Identifier) {
        let group = self.tree.groups.get_mut(self.tree.current());
        group.push_leaf(id);
    }

    pub fn push_text(&mut self, new_attrs: ResolvedTextAttrs, text_id: Identifier) {
        let group = self.tree.groups.get_mut(self.tree.current());
        group.push_text(new_attrs, text_id);
    }
}

pub struct Ctx {
    pub tables: IdVec<TableCtx>,
    pub grids: IdVec<GridCtx>,
    pub figures: IdVec<FigureCtx>,
    pub lists: IdVec<ListCtx>,
    pub outlines: IdVec<OutlineCtx>,
    pub bboxes: IdVec<BBoxCtx>,
}

impl Ctx {
    pub const fn new() -> Self {
        Self {
            tables: IdVec::new(),
            grids: IdVec::new(),
            figures: IdVec::new(),
            lists: IdVec::new(),
            outlines: IdVec::new(),
            bboxes: IdVec::new(),
        }
    }

    pub fn new_bbox(&mut self) -> BBoxId {
        self.bboxes.push(BBoxCtx::new())
    }

    pub fn bbox(&self, kind: &GroupKind) -> Option<&BBoxCtx> {
        Some(self.bboxes.get(kind.bbox()?))
    }
}

pub struct Annotations(Vec<OnceCell<Identifier>>);

impl Annotations {
    pub const fn new() -> Self {
        Self(Vec::new())
    }

    pub fn reserve(&mut self) -> AnnotationId {
        let id = AnnotationId::new(self.0.len() as u32);
        self.0.push(OnceCell::new());
        id
    }

    pub fn init(&mut self, id: AnnotationId, annot: Identifier) {
        self.0[id.idx()]
            .set(annot)
            .map_err(|_| ())
            .expect("annotation to be uninitialized");
    }

    pub fn take(&mut self, id: AnnotationId) -> Node {
        let annot = self.0[id.idx()].take().expect("initialized annotation node");
        Node::Leaf(annot)
    }
}

#[derive(Debug, Clone, PartialEq)]
pub struct BBoxCtx {
    pub rect: Option<(usize, Rect)>,
    pub multi_page: bool,
}

impl BBoxCtx {
    pub fn new() -> Self {
        Self { rect: None, multi_page: false }
    }

    /// Expand the bounding box with a `rect` relative to the current frame
    /// context transform.
    pub fn expand_frame(
        &mut self,
        fc: &FrameContext,
        compute_rect: impl FnOnce() -> Rect,
    ) {
        let Some(page_idx) = fc.page_idx else { return };
        if self.multi_page {
            return;
        }
        let (idx, bbox) = self.rect.get_or_insert((
            page_idx,
            Rect::new(Point::splat(Abs::inf()), Point::splat(-Abs::inf())),
        ));
        if *idx != page_idx {
            self.multi_page = true;
            self.rect = None;
            return;
        }

        let rect = compute_rect();
        let size = rect.size();
        for point in [
            rect.min,
            rect.min + Point::with_x(size.x),
            rect.min + Point::with_y(size.y),
            rect.max,
        ] {
            let p = point.transform(fc.state().transform());
            bbox.min = bbox.min.min(p);
            bbox.max = bbox.max.max(p);
        }
    }

    /// Expand the bounding box with a rectangle that's already transformed into
    /// page coordinates.
    pub fn expand_page(&mut self, inner: &BBoxCtx) {
        self.multi_page |= inner.multi_page;
        if self.multi_page {
            return;
        }

        let Some((page_idx, rect)) = inner.rect else { return };
        let (idx, bbox) = self.rect.get_or_insert((
            page_idx,
            Rect::new(Point::splat(Abs::inf()), Point::splat(-Abs::inf())),
        ));
        if *idx != page_idx {
            self.multi_page = true;
            self.rect = None;
            return;
        }

        bbox.min = bbox.min.min(rect.min);
        bbox.max = bbox.max.max(rect.max);
    }

    pub fn to_krilla(&self) -> Option<BBox> {
        let (page_idx, rect) = self.rect?;
        let rect = kg::Rect::from_ltrb(
            rect.min.x.to_f32(),
            rect.min.y.to_f32(),
            rect.max.x.to_f32(),
            rect.max.y.to_f32(),
        )
        .unwrap();
        Some(BBox::new(page_idx, rect))
    }
}

pub fn finish(tree: &mut Tree) {
    for figure_id in tree.ctx.figures.ids() {
        build_figure(tree, figure_id);
    }
    for table_id in tree.ctx.tables.ids() {
        build_table(tree, table_id);
    }
    for grid_id in tree.ctx.grids.ids() {
        build_grid(tree, grid_id);
    }
}
