//! 推导完整遍历、提前停止及归约中的失败处理义务。
//! 未知长度归约合并初始与迭代摘要直到稳定，保留空输入的可能性。

use super::{State, flow};
use crate::compiler::model::*;
use crate::syntax::diagnostic::Diagnostic;
use std::collections::{BTreeMap, BTreeSet};

impl State<'_> {
    pub(super) fn collection(&mut self, collection: &Collection) -> Result<Flow, Diagnostic> {
        if collection
            .sequence
            .value_type
            .database_stream(self.definitions)
            && let Some(definition) = self.database_error
        {
            let Shape::Choice(variants) = &self.definitions[definition].shape else {
                unreachable!()
            };
            self.failures
                .extend((0..variants.len()).map(|variant| Failure {
                    definition,
                    variant,
                }));
        }
        let sequence = self.expression(&collection.sequence)?;
        if let Some(limit) = &collection.limit {
            self.expression(limit)?;
        }
        let extra = collection
            .extra
            .as_ref()
            .map(|value| self.expression(value))
            .transpose()?;
        let signature = &self.signatures[collection.handler];
        self.failures.extend(&signature.failures);
        let before = self.pending.clone();
        // Stream 的失败元素尚未产生；只在实际消费合同中引入义务，关闭不等于丢弃已读错误。
        let shape =
            if let Type::Stream(item) | Type::AsyncStream(item) = &collection.sequence.value_type {
                match &sequence.sequence {
                    // 转发流可能没有现存错误 token，但未来拉取仍会产生错误义务。
                    Some(Sequence::Unknown(value)) if value.tokens.is_empty() => {
                        Sequence::Unknown(Box::new(self.unknown(item, 0)?))
                    }
                    Some(sequence) => sequence.clone(),
                    None => Sequence::Unknown(Box::new(self.unknown(item, 0)?)),
                }
            } else {
                sequence
                    .sequence
                    .clone()
                    .unwrap_or_else(|| Sequence::Unknown(Box::default()))
            };
        let mut produced = Vec::new();
        let mut state = extra.clone().unwrap_or_default();
        let mut fresh = BTreeMap::new();
        match &shape {
            Sequence::Known(items) => {
                for item in items {
                    self.budget(0)?;
                    let value =
                        self.apply_handler(collection, item, &state, &mut BTreeMap::new())?;
                    if collection.operation.reduction() {
                        state = value;
                        produced.push(state.clone());
                    } else {
                        produced.push(value);
                    }
                }
            }
            Sequence::Unknown(item) => {
                if collection.operation.reduction() {
                    // 初始状态参与合并，因为未知集合可以完全不执行 handler。
                    loop {
                        self.budget(0)?;
                        let value = self.apply_handler(collection, item, &state, &mut fresh)?;
                        let previous = state.clone();
                        state.merge(&value);
                        if previous == state {
                            break;
                        }
                    }
                    produced.push(state.clone());
                } else {
                    produced.push(self.apply_handler(collection, item, &state, &mut fresh)?);
                }
            }
        }
        let empty = matches!(&shape, Sequence::Known(items) if items.is_empty());
        let nonempty = matches!(&shape, Sequence::Known(items) if !items.is_empty());
        let partial = matches!(
            collection.operation,
            Operation::Find | Operation::ReduceUntil
        );
        self.pending = before;
        // 提前结束只保证访问首项，不能用选中的输出消费剩余元素的义务。
        if partial {
            if let Sequence::Known(items) = &shape
                && let Some(first) = items.first()
            {
                self.pending.retain(|token| !first.tokens.contains(token));
            }
        } else {
            self.pending
                .retain(|token| !sequence.tokens.contains(token));
        }
        if let Some(extra) = &extra
            && (nonempty || collection.operation.reduction())
        {
            self.pending.retain(|token| !extra.tokens.contains(token));
        }
        let mut output = match collection.operation {
            Operation::ParallelEach => Flow::default(),
            Operation::Each if signature.outputs.is_empty() => Flow::default(),
            Operation::Each => Flow::sequence(match shape {
                Sequence::Known(_) => Sequence::Known(produced),
                Sequence::Unknown(_) => Sequence::Unknown(Box::new(produced.pop().unwrap())),
            }),
            Operation::Filter if empty => Flow::sequence(Sequence::Known(Vec::new())),
            Operation::Filter => Flow::sequence(Sequence::Unknown(Box::new(shape.element()))),
            Operation::Find => {
                let item = if empty {
                    return Ok(Flow::variants([(0, Flow::default())]));
                } else if matches!(&shape, Sequence::Known(items) if items.len() == 1) {
                    shape.element()
                } else {
                    self.unknown(&signature.inputs[0].1, 0)?
                };
                Flow::variants([(0, Flow::default()), (1, item)])
            }
            Operation::Sum => Flow::default(),
            Operation::Reduce => state,
            Operation::ReduceUntil => {
                if let Some((first, rest)) = produced.split_first() {
                    let mut output = first.clone();
                    for value in rest {
                        output.merge(value);
                    }
                    output
                } else {
                    state
                }
            }
        };
        let retained: BTreeSet<_> = self.pending.intersection(&output.tokens).copied().collect();
        if !retained.is_empty() {
            let mut mapping = BTreeMap::new();
            for token in &output.tokens {
                let target = if retained.contains(token) {
                    self.token()?
                } else {
                    *token
                };
                mapping.insert(*token, BTreeSet::from([target]));
            }
            output = self.instantiate(&output, &mut mapping, 0)?;
        }
        self.pending.extend(&output.tokens);
        self.closed_streams.extend(sequence.streams);
        self.refresh_streams(&mut output);
        Ok(output)
    }

    fn apply_handler(
        &mut self,
        collection: &Collection,
        item: &Flow,
        extra: &Flow,
        fresh: &mut BTreeMap<usize, BTreeSet<usize>>,
    ) -> Result<Flow, Diagnostic> {
        let signature = &self.signatures[collection.handler];
        let mut mapping = fresh.clone();
        for template in &signature.input_flows {
            for token in &template.tokens {
                mapping.remove(token);
            }
        }
        for (template, actual) in signature.input_flows.iter().zip([item, extra]) {
            flow::bind(template, actual, &mut mapping);
        }
        let output = if let Some(value) = &signature.returns {
            self.instantiate(value, &mut mapping, 0)?
        } else {
            Flow::default()
        };
        *fresh = mapping;
        if collection.operation == Operation::ReduceUntil {
            Ok(output.fields.get(&0).cloned().unwrap_or_default())
        } else {
            Ok(output)
        }
    }
}