use crate::graph::Graph;
use intuicio_data::lifetime::{ValueReadAccess, ValueWriteAccess};
use intuicio_framework_arena::AnyIndex;
use std::marker::PhantomData;

pub struct QueryIter<'a, Fetch: QueryFetch<'a>> {
    access: Fetch::Access,
}

impl<'a, Fetch: QueryFetch<'a>> QueryIter<'a, Fetch> {
    pub fn new(graph: &'a Graph, index: AnyIndex) -> Self {
        Self {
            access: Fetch::access(graph, index),
        }
    }
}

impl<'a, Fetch: QueryFetch<'a>> Iterator for QueryIter<'a, Fetch> {
    type Item = Fetch::Value;

    fn next(&mut self) -> Option<Self::Item> {
        Fetch::fetch(&mut self.access)
    }
}

pub trait QueryFetch<'a> {
    type Value;
    type Access;

    fn access(graph: &'a Graph, index: AnyIndex) -> Self::Access;
    fn fetch(access: &mut Self::Access) -> Option<Self::Value>;
}

pub trait QueryTransform<'a> {
    type Input;
    type Output;

    fn transform(graph: &'a Graph, input: Self::Input) -> impl Iterator<Item = Self::Output>;
}

impl<'a> QueryFetch<'a> for () {
    type Value = ();
    type Access = ();

    fn access(_: &'a Graph, _: AnyIndex) -> Self::Access {}

    fn fetch(_: &mut Self::Access) -> Option<Self::Value> {
        Some(())
    }
}

pub struct Related<'a, T, Transform: QueryTransform<'a, Input = AnyIndex>>(
    PhantomData<fn() -> &'a (T, Transform)>,
);

impl<'a, T, Transform: QueryTransform<'a, Input = AnyIndex>> QueryFetch<'a>
    for Related<'a, T, Transform>
{
    type Value = Transform::Output;
    type Access = Box<dyn Iterator<Item = Self::Value> + 'a>;

    fn access(graph: &'a Graph, index: AnyIndex) -> Self::Access {
        Box::new(
            graph
                .relations_outgoing::<T>(index)
                .flat_map(|index| Transform::transform(graph, index)),
        )
    }

    fn fetch(access: &mut Self::Access) -> Option<Self::Value> {
        access.next()
    }
}

pub struct Traverse<'a, T, Transform: QueryTransform<'a, Input = AnyIndex>>(
    PhantomData<fn() -> &'a (T, Transform)>,
);

impl<'a, T, Transform: QueryTransform<'a, Input = AnyIndex>> QueryFetch<'a>
    for Traverse<'a, T, Transform>
{
    type Value = Transform::Output;
    type Access = Box<dyn Iterator<Item = Self::Value> + 'a>;

    fn access(graph: &'a Graph, index: AnyIndex) -> Self::Access {
        Box::new(
            graph
                .relations_traverse::<T>(index)
                .flat_map(|index| Transform::transform(graph, index)),
        )
    }

    fn fetch(access: &mut Self::Access) -> Option<Self::Value> {
        access.next()
    }
}

pub struct Node<T>(PhantomData<fn() -> T>);

impl<T> QueryTransform<'_> for Node<T> {
    type Input = AnyIndex;
    type Output = AnyIndex;

    fn transform(graph: &Graph, input: Self::Input) -> impl Iterator<Item = Self::Output> {
        graph.is::<T>(input).then_some(input).into_iter()
    }
}

pub struct Is<T>(PhantomData<fn() -> T>);

impl<T> QueryTransform<'_> for Is<T> {
    type Input = AnyIndex;
    type Output = ();

    fn transform(graph: &Graph, input: Self::Input) -> impl Iterator<Item = Self::Output> {
        graph.is::<T>(input).then_some(()).into_iter()
    }
}

pub struct IsNot<T>(PhantomData<fn() -> T>);

impl<T> QueryTransform<'_> for IsNot<T> {
    type Input = AnyIndex;
    type Output = ();

    fn transform(graph: &Graph, input: Self::Input) -> impl Iterator<Item = Self::Output> {
        (!graph.is::<T>(input)).then_some(()).into_iter()
    }
}

pub struct Copied<T: Copy>(PhantomData<fn() -> T>);

impl<T: Copy> QueryTransform<'_> for Copied<T> {
    type Input = AnyIndex;
    type Output = T;

    fn transform(graph: &Graph, input: Self::Input) -> impl Iterator<Item = Self::Output> {
        graph.read(input).ok().map(|v| *v).into_iter()
    }
}

pub struct Cloned<T: Clone>(PhantomData<fn() -> T>);

impl<T: Clone> QueryTransform<'_> for Cloned<T> {
    type Input = AnyIndex;
    type Output = T;

    fn transform(graph: &Graph, input: Self::Input) -> impl Iterator<Item = Self::Output> {
        graph
            .read(input)
            .ok()
            .map(|v: ValueReadAccess<'_, T>| v.clone())
            .into_iter()
    }
}

impl<'a> QueryFetch<'a> for AnyIndex {
    type Value = AnyIndex;
    type Access = Option<AnyIndex>;

    fn access(_: &'a Graph, index: AnyIndex) -> Self::Access {
        Some(index)
    }

    fn fetch(access: &mut Self::Access) -> Option<Self::Value> {
        access.take()
    }
}

impl QueryTransform<'_> for AnyIndex {
    type Input = AnyIndex;
    type Output = AnyIndex;

    fn transform(_: &Graph, input: Self::Input) -> impl Iterator<Item = Self::Output> {
        std::iter::once(input)
    }
}

impl<'a, T> QueryFetch<'a> for &'a T {
    type Value = ValueReadAccess<'a, T>;
    type Access = Option<ValueReadAccess<'a, T>>;

    fn access(graph: &'a Graph, index: AnyIndex) -> Self::Access {
        graph.read(index).ok()
    }

    fn fetch(access: &mut Self::Access) -> Option<Self::Value> {
        access.take()
    }
}

impl<'a, T> QueryTransform<'a> for &'a T {
    type Input = AnyIndex;
    type Output = ValueReadAccess<'a, T>;

    fn transform(graph: &'a Graph, input: Self::Input) -> impl Iterator<Item = Self::Output> {
        graph.read(input).ok().into_iter()
    }
}

impl<'a, T> QueryFetch<'a> for &'a mut T {
    type Value = ValueWriteAccess<'a, T>;
    type Access = Option<ValueWriteAccess<'a, T>>;

    fn access(graph: &'a Graph, index: AnyIndex) -> Self::Access {
        graph.write(index).ok()
    }

    fn fetch(access: &mut Self::Access) -> Option<Self::Value> {
        access.take()
    }
}

impl<'a, T> QueryTransform<'a> for &'a mut T {
    type Input = AnyIndex;
    type Output = ValueWriteAccess<'a, T>;

    fn transform(graph: &'a Graph, input: Self::Input) -> impl Iterator<Item = Self::Output> {
        graph.write(input).ok().into_iter()
    }
}

impl<'a, T> QueryFetch<'a> for Option<&'a T> {
    type Value = Option<ValueReadAccess<'a, T>>;
    type Access = Option<Option<ValueReadAccess<'a, T>>>;

    fn access(graph: &'a Graph, index: AnyIndex) -> Self::Access {
        Some(graph.read(index).ok())
    }

    fn fetch(access: &mut Self::Access) -> Option<Self::Value> {
        access.take()
    }
}

impl<'a, T> QueryTransform<'a> for Option<&'a T> {
    type Input = AnyIndex;
    type Output = Option<ValueReadAccess<'a, T>>;

    fn transform(graph: &'a Graph, input: Self::Input) -> impl Iterator<Item = Self::Output> {
        std::iter::once(graph.read(input).ok())
    }
}

impl<'a, T> QueryFetch<'a> for Option<&'a mut T> {
    type Value = Option<ValueWriteAccess<'a, T>>;
    type Access = Option<Option<ValueWriteAccess<'a, T>>>;

    fn access(graph: &'a Graph, index: AnyIndex) -> Self::Access {
        Some(graph.write(index).ok())
    }

    fn fetch(access: &mut Self::Access) -> Option<Self::Value> {
        access.take()
    }
}

impl<'a, T> QueryTransform<'a> for Option<&'a mut T> {
    type Input = AnyIndex;
    type Output = Option<ValueWriteAccess<'a, T>>;

    fn transform(graph: &'a Graph, input: Self::Input) -> impl Iterator<Item = Self::Output> {
        std::iter::once(graph.write(input).ok())
    }
}

pub struct Limit<'a, const COUNT: usize, Transform: QueryTransform<'a>>(
    PhantomData<fn() -> &'a Transform>,
);

impl<'a, const COUNT: usize, Transform: QueryTransform<'a>> QueryTransform<'a>
    for Limit<'a, COUNT, Transform>
{
    type Input = Transform::Input;
    type Output = Transform::Output;

    fn transform(graph: &'a Graph, input: Self::Input) -> impl Iterator<Item = Self::Output> {
        Transform::transform(graph, input).take(COUNT)
    }
}

pub struct Single<'a, Transform: QueryTransform<'a>>(PhantomData<fn() -> &'a Transform>);

impl<'a, Transform: QueryTransform<'a>> QueryTransform<'a> for Single<'a, Transform> {
    type Input = Transform::Input;
    type Output = Transform::Output;

    fn transform(graph: &'a Graph, input: Self::Input) -> impl Iterator<Item = Self::Output> {
        Transform::transform(graph, input).take(1)
    }
}

pub struct Query<'a, Transform, Fetch>(PhantomData<fn() -> &'a (Transform, Fetch)>)
where
    Transform: QueryTransform<'a, Input = AnyIndex, Output = AnyIndex>,
    Fetch: QueryFetch<'a>;

impl<'a, Transform, Fetch> QueryTransform<'a> for Query<'a, Transform, Fetch>
where
    Transform: QueryTransform<'a, Input = AnyIndex, Output = AnyIndex>,
    Fetch: QueryFetch<'a>,
{
    type Input = AnyIndex;
    type Output = Fetch::Value;

    fn transform(graph: &'a Graph, input: Self::Input) -> impl Iterator<Item = Self::Output> {
        Transform::transform(graph, input).flat_map(|index| graph.query::<Fetch>(index))
    }
}

macro_rules! impl_fetch_tuple {
    ($($type:ident),+) => {
        impl<'a, $($type: QueryFetch<'a>),+> QueryFetch<'a> for ($($type,)+) {
            type Value = ($($type::Value,)+);
            type Access = ($($type::Access,)+);

            fn access(graph: &'a Graph, index: AnyIndex) -> Self::Access {
                ($($type::access(graph, index),)+)
            }

            fn fetch(access: &mut Self::Access) -> Option<Self::Value> {
                #[allow(non_snake_case)]
                let ($($type,)+) = access;
                Some(($($type::fetch($type)?,)+))
            }
        }
    };
}

impl_fetch_tuple!(A);
impl_fetch_tuple!(A, B);
impl_fetch_tuple!(A, B, C);
impl_fetch_tuple!(A, B, C, D);
impl_fetch_tuple!(A, B, C, D, E);
impl_fetch_tuple!(A, B, C, D, E, F);
impl_fetch_tuple!(A, B, C, D, E, F, G);
impl_fetch_tuple!(A, B, C, D, E, F, G, H);
impl_fetch_tuple!(A, B, C, D, E, F, G, H, I);
impl_fetch_tuple!(A, B, C, D, E, F, G, H, I, J);
impl_fetch_tuple!(A, B, C, D, E, F, G, H, I, J, K);
impl_fetch_tuple!(A, B, C, D, E, F, G, H, I, J, K, L);
impl_fetch_tuple!(A, B, C, D, E, F, G, H, I, J, K, L, M);
impl_fetch_tuple!(A, B, C, D, E, F, G, H, I, J, K, L, M, O);
impl_fetch_tuple!(A, B, C, D, E, F, G, H, I, J, K, L, M, O, P);
