use super::node::Span;
use crate::antlr::hamelinparser::{
    ExpressionQueryContextAttrs, QueryContextAll, QueryEOFContextAttrs,
    StandaloneQueryContextAttrs, WithQueryContextAttrs,
};
use crate::antlr::interval;
use crate::err::{TranslationError, TranslationErrors};
use crate::resilient_parse_query;
use crate::tree::ast::context::{FromCst, ParseContext};
use crate::tree::ast::{
    IntoParsed, IntoTyped, ParseOptions, ParseWithContext, ParseWithErrors,
    QueryTranslationOptions, TypeCheck,
};
use crate::tree::typed_ast::context::StatementTranslationContext;
use crate::tree::typed_ast::query::TypedStatement;
use derive_more::From;
use std::ops::RangeInclusive;
use std::rc::Rc;

use super::expression::Expression;
use super::identifier::{ParsedIdentifier, ParsedSimpleIdentifier};
use super::node::Spannable;
use super::pipeline::Pipeline;

/// Top-level query node
#[derive(Debug, Clone)]
pub struct Query {
    pub span: Span,
    pub kind: QueryKind,
}

impl PartialEq for Query {
    fn eq(&self, other: &Self) -> bool {
        // Compare only the semantic content (kind), ignoring CST
        self.kind == other.kind
    }
}

impl ParseWithContext for Query {
    fn parse_with_context(input: String, mut ctx: ParseContext) -> (Self, TranslationErrors) {
        match resilient_parse_query(input) {
            Ok((csteof, parse_errors)) => match csteof.query() {
                Some(cst) => {
                    let ast = Self::from_cst_with_context(cst, &mut ctx);
                    let mut errors = ctx.take_errors();
                    errors.extend(parse_errors);
                    (ast, errors)
                }
                None => {
                    let err = ctx.error("parse did not resolve a query tree").emit();
                    let mut errors = ctx.take_errors();
                    errors.extend(parse_errors);
                    (
                        Self {
                            span: Span::NONE,
                            kind: QueryKind::Error(err),
                        },
                        errors,
                    )
                }
            },
            Err(e) => {
                let error = ctx
                    .error("fatal error initializing parser")
                    .with_source_boxed(e.into())
                    .emit();
                (
                    Self {
                        span: Span::NONE,
                        kind: QueryKind::Error(error.clone()),
                    },
                    (*error).clone().single(),
                )
            }
        }
    }
}

impl ParseWithErrors for Query {
    fn parse_with_errors(input: impl Into<String>) -> (Self, TranslationErrors) {
        Self::parse_with_context(input.into(), ParseContext::new())
    }
}

impl IntoParsed for Query {
    type Options = ParseOptions<Query>;

    fn parse_with(input: impl Into<String>) -> ParseOptions<Query> {
        ParseOptions::new(input)
    }
}

impl Spannable for Query {
    fn span(&self) -> Option<RangeInclusive<usize>> {
        self.span.to_range()
    }
}

#[derive(Debug, Clone, PartialEq, From)]
pub enum QueryKind {
    Valid(ValidQuery),
    Error(Rc<TranslationError>),
}

/// A valid query with optional WITH clauses and a main pipeline
#[derive(Debug, Clone)]
pub struct ValidQuery {
    pub span: Span,
    pub with_clauses: Vec<WithClause>,
    pub main_pipeline: Rc<Pipeline>,
}

impl PartialEq for ValidQuery {
    fn eq(&self, other: &Self) -> bool {
        self.with_clauses == other.with_clauses && self.main_pipeline == other.main_pipeline
    }
}

impl Spannable for ValidQuery {
    fn span(&self) -> Option<RangeInclusive<usize>> {
        self.span.to_range()
    }
}

/// A single WITH clause binding
#[derive(Debug, Clone)]
pub struct WithClause {
    pub span: Span,
    pub name: ParsedIdentifier,
    pub pipeline: Rc<Pipeline>,
}

impl PartialEq for WithClause {
    fn eq(&self, other: &Self) -> bool {
        self.name == other.name && self.pipeline == other.pipeline
    }
}

impl Spannable for WithClause {
    fn span(&self) -> Option<RangeInclusive<usize>> {
        self.span.to_range()
    }
}

impl FromCst<Rc<QueryContextAll<'static>>> for Query {
    fn from_cst_with_context(cst: Rc<QueryContextAll<'static>>, ctx: &mut ParseContext) -> Self {
        let query_span = Some(interval(cst.as_ref()));

        let kind = match cst.as_ref() {
            QueryContextAll::WithQueryContext(wctx) => {
                // WITH clause query: (WITH simpleIdentifier ASSIGN pipeline)+ pipeline
                let simple_idents = wctx.simpleIdentifier_all();
                let pipelines = wctx.pipeline_all();

                // The grammar guarantees at least one WITH clause and a main pipeline
                // simple_idents pairs with the first N pipelines (the WITH clauses)
                // The last pipeline is the main pipeline
                if pipelines.is_empty() {
                    QueryKind::Error(ctx.error("WITH query missing pipelines").at(wctx).emit())
                } else {
                    // Zip naturally stops when simple_idents runs out
                    let with_clauses: Vec<WithClause> = simple_idents
                        .into_iter()
                        .zip(pipelines.iter())
                        .map(|(ident_cst, pipeline_cst)| {
                            // Compute span from ident to end of pipeline
                            let ident_span = interval(ident_cst.as_ref());
                            let pipeline_span = interval(pipeline_cst.as_ref());
                            let clause_span = *ident_span.start()..=*pipeline_span.end();

                            WithClause {
                                span: clause_span.into(),
                                name: ParsedSimpleIdentifier::from_cst_with_context(
                                    ident_cst.as_ref(),
                                    ctx,
                                )
                                .into(),
                                pipeline: Rc::new(Pipeline::from_cst_with_context(
                                    pipeline_cst.clone(),
                                    ctx,
                                )),
                            }
                        })
                        .collect();

                    // Last pipeline is the main pipeline
                    // We already checked pipelines is not empty, so this is safe
                    let main_pipeline = match pipelines.last() {
                        Some(pipeline_cst) => {
                            Rc::new(Pipeline::from_cst_with_context(pipeline_cst.clone(), ctx))
                        }
                        None => {
                            // Should never happen due to check above, but be defensive
                            return Query {
                                span: query_span.into(),
                                kind: QueryKind::Error(
                                    ctx.error("WITH query missing main pipeline")
                                        .at(wctx)
                                        .emit(),
                                ),
                            };
                        }
                    };

                    QueryKind::Valid(ValidQuery {
                        span: query_span.clone().into(),
                        with_clauses,
                        main_pipeline,
                    })
                }
            }
            QueryContextAll::StandaloneQueryContext(sctx) => {
                // Just a pipeline, no WITH clauses
                match sctx.pipeline() {
                    Some(pipeline_cst) => QueryKind::Valid(ValidQuery {
                        span: query_span.clone().into(),
                        with_clauses: vec![],
                        main_pipeline: Rc::new(Pipeline::from_cst_with_context(
                            pipeline_cst.clone(),
                            ctx,
                        )),
                    }),
                    None => QueryKind::Error(ctx.error("Query missing its main pipeline").emit()),
                }
            }
            QueryContextAll::ExpressionQueryContext(ectx) => {
                // Just an expression, convert to a pipeline with a SELECT command
                match ectx.expression() {
                    Some(expr_cst) => {
                        let expr = Expression::from_cst_with_context(expr_cst.clone(), ctx);
                        let pipeline = Pipeline::from(expr);
                        QueryKind::Valid(ValidQuery {
                            span: query_span.clone().into(),
                            with_clauses: vec![],
                            main_pipeline: Rc::new(pipeline),
                        })
                    }
                    None => QueryKind::Error(ctx.error("Expression did not parse").emit()),
                }
            }
            QueryContextAll::Error(_) => ctx
                .error("query parse error")
                .at(cst.as_ref())
                .emit()
                .into(),
        };

        Query {
            span: query_span.into(),
            kind,
        }
    }
}

impl TypeCheck for Query {
    type Output = TypedStatement;

    fn type_check_with_context(
        ast: Rc<Self>,
        ctx: &mut StatementTranslationContext,
    ) -> Self::Output {
        TypedStatement::from_ast_with_context(ast, ctx)
    }
}

impl IntoTyped for Query {
    type Output = TypedStatement;
    type Options = QueryTranslationOptions<Query>;

    fn typed_with(self) -> QueryTranslationOptions<Query> {
        QueryTranslationOptions::from_ast(Rc::new(self))
    }
}

impl IntoTyped for Rc<Query> {
    type Output = TypedStatement;
    type Options = QueryTranslationOptions<Query>;

    fn typed_with(self) -> QueryTranslationOptions<Query> {
        QueryTranslationOptions::from_ast(self)
    }
}
