//! 由已检查表达式计算事实，同时验证每个构造和调用边界。
//! 不支持推导的数值操作只失去精度，不允许绕过合同。

use super::{Fact, Proof, number, require};
use crate::compiler::model::{
    numeric::{Number, Range},
    *,
};
use crate::syntax::{ast::Literal, diagnostic::Diagnostic, source::Span};
use std::{collections::BTreeMap, rc::Rc};

impl Proof<'_> {
    pub(super) fn expression(&mut self, value: &Expression) -> Result<Rc<Fact>, Diagnostic> {
        self.remaining = self.remaining.checked_sub(1).ok_or_else(|| {
            Diagnostic::new("C015", "range proof exceeds its work limit", value.span)
        })?;
        let unknown = || Fact::unknown(&value.value_type);
        Ok(match &value.kind {
            Node::Literal(Literal::Number(text))
                if matches!(value.value_type, Type::Int | Type::Decimal) =>
            {
                Fact::scalar(Some(Range::point(
                    Number::parse(text, &value.value_type).expect("checked literal"),
                )))
            }
            Node::Literal(_) => unknown(),
            Node::TestSecret(inner)
            | Node::TestClock(inner)
            | Node::TestDrain { limit: inner, .. } => {
                self.expression(inner)?;
                unknown()
            }
            Node::Variable(index) => self.variables[*index].clone(),
            Node::Field(base, index) => {
                self.expression(base)?
                    .field(*index, &base.value_type, self.definitions)
            }
            Node::Call(index, arguments) => {
                self.arguments(*index, arguments)?;
                self.outputs(*index)
            }
            Node::Record(index, values) => {
                let mut fields = BTreeMap::new();
                for (position, value) in values {
                    let fact = self.expression(value)?;
                    self.require_value(value, &fact, self.definitions[*index].ranges[*position])?;
                    fields.insert(*position, fact);
                }
                Rc::new(Fact {
                    fields,
                    ..Fact::default()
                })
            }
            Node::Variant(index, variant, values) => {
                let Shape::Choice(variants) = &self.definitions[*index].shape else {
                    unreachable!()
                };
                for (value, &expected) in values.iter().zip(&variants[*variant].ranges) {
                    let fact = self.expression(value)?;
                    self.require_value(value, &fact, expected)?;
                }
                unknown()
            }
            Node::Capture(capture) => {
                let fact = self.expression(&capture.call)?;
                let Shape::Choice(variants) = &self.definitions[capture.definition].shape else {
                    unreachable!()
                };
                if let Node::Call(function, _) = capture.call.kind {
                    let signature = &self.signatures[function];
                    if signature.returns.is_some() {
                        for (((_, kind), &actual), &expected) in signature
                            .outputs
                            .iter()
                            .zip(&signature.output_ranges)
                            .zip(&variants[capture.success].ranges)
                        {
                            require(
                                actual.or_else(|| super::numeric_range(kind)),
                                expected,
                                value.span,
                            )?;
                        }
                    }
                } else {
                    require(
                        fact.range
                            .or_else(|| super::numeric_range(&capture.call.value_type)),
                        variants[capture.success].ranges[0],
                        value.span,
                    )?;
                }
                unknown()
            }
            Node::Unary(operator, base) => {
                let fact = self.expression(base)?;
                if operator == "-" {
                    fact.range
                        .and_then(number::negate)
                        .map_or_else(unknown, |range| Fact::scalar(Some(range)))
                } else {
                    unknown()
                }
            }
            Node::Binary(operator, left, right) => {
                let left = self.expression(left)?;
                let right = self.expression(right)?;
                if matches!(value.value_type, Type::Int | Type::Decimal) {
                    let range = left
                        .range
                        .zip(right.range)
                        .and_then(|(left, right)| number::binary(operator, left, right));
                    if value.value_type == Type::Int
                        && matches!(operator.as_str(), "+" | "-" | "*")
                        && range.is_some()
                    {
                        let span = value.span;
                        self.arithmetic.insert((span.source, span.start, span.end));
                    }
                    range.map_or_else(unknown, |range| Fact::scalar(Some(range)))
                } else {
                    unknown()
                }
            }
            Node::Decimal(base) => {
                let fact = self.expression(base)?;
                Fact::scalar(fact.range.map(number::decimal))
            }
            Node::Some(base) => self.expression(base)?,
            Node::Length(base) => {
                self.expression(base)?;
                Fact::scalar(Some(Range {
                    minimum: Number::Int(0),
                    maximum: Number::Int(i64::MAX),
                }))
            }
            Node::List(values) => {
                let mut range = None;
                let mut numeric = true;
                for value in values {
                    let fact = self.expression(value)?;
                    numeric &= fact.range.is_some();
                    if let Some(actual) = fact.range {
                        range = Some(range.map_or(actual, |previous: Range| previous.hull(actual)));
                    }
                }
                Rc::new(Fact {
                    item: numeric.then(|| Fact::scalar(range)),
                    ..Fact::default()
                })
            }
            Node::Map(entries) => {
                for (key, value) in entries {
                    self.expression(key)?;
                    self.expression(value)?;
                }
                unknown()
            }
            Node::Collection(collection) => self.collection(collection, value.span)?,
            Node::Access(operation, values) => {
                let facts = values
                    .iter()
                    .map(|value| self.expression(value))
                    .collect::<Result<Vec<_>, _>>()?;
                match operation {
                    Access::First => facts[0].item.clone().unwrap_or_else(unknown),
                    _ => unknown(),
                }
            }
            Node::Crypto(_, values)
            | Node::Assert(values)
            | Node::Calendar(_, values)
            | Node::Standard(_, values)
            | Node::Bytes(_, values)
            | Node::Net(_, values)
            | Node::Http(_, values)
            | Node::Tls(_, values)
            | Node::Database(_, values)
            | Node::Auth(_, values)
            | Node::Authorization(_, values)
            | Node::Enqueue(_, values)
            | Node::Io(_, values)
            | Node::Upload(_, values)
            | Node::Json(_, values) => {
                for value in values {
                    self.expression(value)?;
                }
                unknown()
            }
            Node::Schedule(operation, target, values) => {
                self.arguments(*target, values)?;
                if *operation == Schedule::Run {
                    unknown()
                } else {
                    self.outputs(*target)
                }
            }
            Node::Join { task, .. } => {
                self.expression(task)?;
                if let Type::Task(target) = task.value_type {
                    self.outputs(target)
                } else {
                    unknown()
                }
            }
            Node::Timeout {
                task,
                millis,
                fallback,
            } => {
                self.expression(task)?;
                self.expression(millis)?;
                let Type::Task(target) = task.value_type else {
                    unreachable!("checked Task");
                };
                self.alternatives([target, *fallback], value)?
            }
            Node::Race(tasks) => {
                let mut targets = Vec::new();
                for task in tasks {
                    self.expression(task)?;
                    let Type::Task(target) = task.value_type else {
                        unreachable!("checked Task");
                    };
                    targets.push(target);
                }
                self.alternatives(targets, value)?
            }
            Node::Stream(base) => self.expression(base)?,
            Node::Channel(limit)
            | Node::Receive(limit)
            | Node::Group(limit)
            | Node::Ticks(limit) => {
                self.expression(limit)?;
                unknown()
            }
            Node::Send(channel, sent) => {
                self.expression(channel)?;
                self.expression(sent)?;
                unknown()
            }
            Node::GroupRun {
                group,
                target,
                arguments,
            } => {
                self.expression(group)?;
                self.arguments(*target, arguments)?;
                unknown()
            }
            Node::Timer(_, base) => {
                self.expression(base)?;
                unknown()
            }
            Node::Id(base) | Node::IdText(base) | Node::Sum(base) | Node::Close(base) => {
                self.expression(base)?;
                unknown()
            }
        })
    }

    fn arguments(&mut self, function: usize, arguments: &[Expression]) -> Result<(), Diagnostic> {
        for (argument, &expected) in arguments
            .iter()
            .zip(&self.signatures[function].input_ranges)
        {
            let fact = self.expression(argument)?;
            self.require_value(argument, &fact, expected)?;
        }
        Ok(())
    }

    fn alternatives(
        &mut self,
        targets: impl IntoIterator<Item = usize>,
        source: &Expression,
    ) -> Result<Rc<Fact>, Diagnostic> {
        let mut fact: Option<Rc<Fact>> = None;
        for target in targets {
            if self.signatures[target].returns.is_none() {
                continue;
            }
            self.remaining = self
                .remaining
                .checked_sub(self.signatures[target].outputs.len().max(1))
                .ok_or_else(|| {
                    Diagnostic::new("C015", "range proof exceeds its work limit", source.span)
                })?;
            let value = self.outputs(target);
            fact = Some(match fact {
                Some(previous) => Fact::union(&previous, &value),
                None => value,
            });
        }
        Ok(fact.unwrap_or_else(|| Fact::unknown(&source.value_type)))
    }

    fn outputs(&self, function: usize) -> Rc<Fact> {
        let ranges = &self.signatures[function].output_ranges;
        let outputs = &self.signatures[function].outputs;
        if let [range] = ranges.as_slice() {
            return Fact::scalar(range.or_else(|| super::numeric_range(&outputs[0].1)));
        }
        Rc::new(Fact {
            fields: ranges
                .iter()
                .enumerate()
                .map(|(index, range)| {
                    (
                        index,
                        Fact::scalar(range.or_else(|| super::numeric_range(&outputs[index].1))),
                    )
                })
                .collect(),
            ..Fact::default()
        })
    }

    fn collection(&mut self, collection: &Collection, span: Span) -> Result<Rc<Fact>, Diagnostic> {
        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];
        let item = if collection.sequence.value_type == Type::Bytes {
            Fact::scalar(Some(Range {
                minimum: Number::Int(0),
                maximum: Number::Int(255),
            }))
        } else {
            sequence
                .item
                .clone()
                .unwrap_or_else(|| Fact::unknown(&signature.inputs[0].1))
        };
        self.require_value(&collection.sequence, &item, signature.input_ranges[0])?;
        if let (Some(extra), Some(value)) = (&extra, &collection.extra) {
            self.require_value(value, extra, signature.input_ranges[1])?;
            if collection.operation.reduction() && signature.returns.is_some() {
                require(signature.output_ranges[0], signature.input_ranges[1], span)?;
            }
        }
        Ok(match collection.operation {
            Operation::ParallelEach => Fact::unknown(&Type::Unit),
            Operation::Each => Rc::new(Fact {
                item: Some(self.outputs(collection.handler)),
                ..Fact::default()
            }),
            Operation::Filter => sequence,
            Operation::Find => item,
            Operation::Reduce | Operation::ReduceUntil => {
                let initial = extra.expect("reduction state").range;
                let range = if signature.returns.is_none() {
                    initial
                } else {
                    initial
                        .zip(signature.output_ranges[0])
                        .map(|(initial, next)| initial.hull(next))
                };
                Fact::scalar(range)
            }
            Operation::Sum => Fact::unknown(&signature.outputs[0].1),
        })
    }
}