//! Parsers for loop statements and loop condition clauses.

use std::collections::HashMap;

use crate::ast::Condition;

use super::super::error::{ErrorSpan, ParseError};
use super::super::role::parse_role_ref;
use super::super::statement::parse_block;
use super::super::types::Statement;
use super::super::Rule;
use super::{next_required, parse_loop_predicate};

fn parse_loop_condition(
    spec: pest::iterators::Pair<Rule>,
    declared_roles: &std::collections::HashSet<String>,
    input: &str,
) -> Result<Option<Condition>, ParseError> {
    match spec.as_rule() {
        Rule::loop_decide => {
            let span = spec.as_span();
            let role_pair = spec
                .into_inner()
                .find(|p| p.as_rule() == Rule::role_ref)
                .ok_or_else(|| ParseError::Syntax {
                    span: ErrorSpan::from_pest_span(span, input),
                    message: "loop decide by is missing a role".to_string(),
                })?;
            let role = parse_role_ref(role_pair, declared_roles, input)?;
            Ok(Some(Condition::RoleDecides(role)))
        }
        Rule::loop_repeat => {
            let span = spec.as_span();
            let mut repeat_inner = spec.into_inner();
            let count_pair = next_required(
                &mut repeat_inner,
                span,
                input,
                "loop repeat is missing value",
            )?;
            let count_str = count_pair.as_str();
            if let Ok(count) = count_str.parse::<usize>() {
                Ok(Some(Condition::Count(count)))
            } else {
                let token_stream =
                    syn::parse_str::<proc_macro2::TokenStream>(count_str).map_err(|e| {
                        ParseError::InvalidCondition {
                            message: format!("Invalid repeat value: {e}"),
                            span: ErrorSpan::from_pest_span(span, input),
                        }
                    })?;
                Ok(Some(Condition::Custom(token_stream)))
            }
        }
        Rule::loop_while => {
            let span = spec.as_span();
            let mut while_inner = spec.into_inner();
            let cond_pair = next_required(
                &mut while_inner,
                span,
                input,
                "loop while is missing condition string",
            )?;
            let cond_lit = syn::parse_str::<syn::LitStr>(cond_pair.as_str()).map_err(|e| {
                ParseError::InvalidCondition {
                    message: format!("Invalid loop condition string: {e}"),
                    span: ErrorSpan::from_pest_span(span, input),
                }
            })?;
            Ok(Some(Condition::Custom(parse_loop_predicate(
                &cond_lit.value(),
                span,
                input,
            )?)))
        }
        Rule::loop_forever => Ok(None),
        _ => Ok(None),
    }
}

/// Parse loop statement.
pub(crate) fn parse_loop_stmt(
    pair: pest::iterators::Pair<Rule>,
    declared_roles: &std::collections::HashSet<String>,
    input: &str,
    protocol_defs: &HashMap<String, Vec<Statement>>,
) -> std::result::Result<Statement, ParseError> {
    let mut condition = None;
    let mut body = Vec::new();

    for item in pair.into_inner() {
        match item.as_rule() {
            Rule::loop_spec => {
                for spec in item.into_inner() {
                    if matches!(
                        spec.as_rule(),
                        Rule::loop_decide
                            | Rule::loop_repeat
                            | Rule::loop_while
                            | Rule::loop_forever
                    ) {
                        condition = parse_loop_condition(spec, declared_roles, input)?;
                    }
                }
            }
            Rule::loop_decide | Rule::loop_repeat | Rule::loop_while | Rule::loop_forever => {
                condition = parse_loop_condition(item, declared_roles, input)?;
            }
            Rule::block => {
                body = parse_block(item, declared_roles, input, protocol_defs)?;
            }
            _ => {}
        }
    }

    Ok(Statement::Loop { condition, body })
}
