use std::collections::HashSet;
use std::fmt::{self, Display, Formatter};
use std::ops::RangeInclusive;
use std::rc::Rc;
use std::sync::Arc;

use derive_more::From;

use ordermap::OrderSet;

use super::command::CommandKind;
use super::expression::{ErrorExpression, Expression};
use super::identifier::{Identifier, ParsedIdentifier, ParsedSimpleIdentifier};
use super::node::{Span, Spannable};
use super::pipeline::Pipeline;
use crate::antlr::hamelinparser::{
    DefRhsContextAll, DefRhsExpressionContextAttrs, DefRhsPipelineContextAttrs,
    DefStatementContextAll, DefStatementContextAttrs, ExpressionQueryMainContextAttrs,
    PipelineQueryMainContextAttrs, QueryContextAll, QueryEOFContextAttrs, QueryMainContextAll,
    QueryWithDefsContextAttrs, QueryWithoutDefsContextAttrs, SimpleIdentifierContextAll,
};
use crate::antlr::interval;
use crate::antlr::resilient_parse_query;
use crate::err::{TranslationError, TranslationErrors};
use crate::tree::ast::context::{FromCst, ParseContext};
use crate::tree::ast::{ParseWithContext, ParseWithErrors, TypeCheck};
use crate::tree::typed_ast::context::StatementTranslationContext;
use crate::tree::typed_ast::query::TypedStatement;
use crate::write_utils::pad;
/// Top-level query node
#[derive(Debug, Clone)]
pub struct Query {
    pub span: Span,
    pub kind: QueryKind,
}

impl Query {
    pub fn valid_ref(&self) -> Result<&ValidQuery, Arc<TranslationError>> {
        match &self.kind {
            QueryKind::Valid(valid_query) => Ok(valid_query),
            QueryKind::Error(translation_error) => Err(translation_error.clone()),
        }
    }

    /// Returns all dataset (table) references in the query.
    ///
    /// This walks the AST to find all tables referenced in:
    /// - FROM clauses
    /// - JOIN/LOOKUP clauses
    /// - APPEND target tables
    /// - MATCH pattern tables
    ///
    /// Tabular `DEF` names (`DefBody::Pipeline`, user or compiler-inserted) are treated as local
    /// CTEs and excluded from this list. Scalar `DEF` names do not shadow table references and are
    /// not excluded.
    pub fn datasets(&self) -> Result<Vec<Identifier>, Arc<TranslationError>> {
        let valid = self.valid_ref()?;

        let mut cte_names: HashSet<Identifier> = HashSet::new();
        for d in &valid.defs {
            if let DefBody::Pipeline(_) = &d.body {
                if let Ok(id) = d.name.valid_ref() {
                    cte_names.insert(id.clone());
                }
            }
        }

        let mut datasets = OrderSet::new();

        // Collect all pipelines to process (main + DEF clauses)
        let mut pipelines: Vec<&Pipeline> = vec![&valid.main_pipeline];
        for d in &valid.defs {
            if let DefBody::Pipeline(p) = &d.body {
                pipelines.push(&p);
            }
        }
        for pipeline in pipelines {
            for cmd in &pipeline.commands {
                match &cmd.kind {
                    CommandKind::From(from_cmd) => {
                        for clause in &from_cmd.clauses {
                            let table_ref = clause.table_reference();
                            let identifier = table_ref.identifier.valid_ref()?;
                            if !cte_names.contains(identifier) {
                                datasets.insert(identifier.clone());
                            }
                        }
                    }
                    CommandKind::Union(union_cmd) => {
                        for clause in &union_cmd.clauses {
                            let table_ref = clause.table_reference();
                            let identifier = table_ref.identifier.valid_ref()?;
                            if !cte_names.contains(identifier) {
                                datasets.insert(identifier.clone());
                            }
                        }
                    }
                    CommandKind::Join(join_cmd) => {
                        let table_ref = join_cmd.other.table_reference();
                        let identifier = table_ref.identifier.valid_ref()?;
                        if !cte_names.contains(identifier) {
                            datasets.insert(identifier.clone());
                        }
                    }
                    CommandKind::Lookup(lookup_cmd) => {
                        let table_ref = lookup_cmd.other.table_reference();
                        let identifier = table_ref.identifier.valid_ref()?;
                        if !cte_names.contains(identifier) {
                            datasets.insert(identifier.clone());
                        }
                    }
                    CommandKind::Append(append_cmd) => {
                        let identifier = append_cmd.table.identifier.valid_ref()?;
                        if !cte_names.contains(identifier) {
                            datasets.insert(identifier.clone());
                        }
                    }
                    CommandKind::Match(match_cmd) => {
                        let cte_refs: HashSet<&Identifier> = cte_names.iter().collect();
                        for pattern in &match_cmd.pattern {
                            pattern.extract_datasets(&cte_refs, &mut datasets)?;
                        }
                    }
                    _ => {}
                }
            }
        }

        Ok(datasets.into_iter().collect())
    }
}

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 Spannable for Query {
    fn span(&self) -> Option<RangeInclusive<usize>> {
        self.span.to_range()
    }
}

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

/// A valid query: optional `DEF` prefix (user and/or compiler-generated tabular or scalar
/// bindings), plus a main pipeline (or expression lowered to a pipeline).
#[derive(Debug, Clone)]
pub struct ValidQuery {
    pub span: Span,
    /// `DEF` bindings in source order (all precede the main body). Tabular `DEF`s are CTE-shaped
    /// (`DefBody::Pipeline`); scalars use `DefBody::Expression`.
    pub defs: Vec<DefStatement>,
    pub main_pipeline: Arc<Pipeline>,
}

/// `DEF name = … ;` — tabular (pipeline) or scalar (`expression`) binding.
#[derive(Debug, Clone, PartialEq)]
pub struct DefStatement {
    pub span: Span,
    pub name: ParsedIdentifier,
    pub body: DefBody,
}

#[derive(Debug, Clone, PartialEq)]
pub enum DefBody {
    Pipeline(Arc<Pipeline>),
    Expression(Arc<Expression>),
}

impl FromCst<Rc<DefRhsContextAll<'static>>> for DefBody {
    fn from_cst_with_context(rhs: Rc<DefRhsContextAll<'static>>, ctx: &mut ParseContext) -> Self {
        match rhs.as_ref() {
            DefRhsContextAll::DefRhsPipelineContext(pctx) => {
                let pipe = pctx
                    .pipeline()
                    .map(|p| Arc::new(Pipeline::from_cst_with_context(p, ctx)));
                pipe.map(DefBody::Pipeline).unwrap_or_else(|| {
                    DefBody::Expression(Arc::new(Expression::from_kind(ErrorExpression {
                        error: ctx.error("DEF pipeline missing").at(rhs.as_ref()).emit(),
                    })))
                })
            }
            DefRhsContextAll::DefRhsExpressionContext(ectx) => {
                let expr = ectx
                    .expression()
                    .map(|e| Arc::new(Expression::from_cst_with_context(e, ctx)));
                expr.map(DefBody::Expression).unwrap_or_else(|| {
                    DefBody::Expression(Arc::new(Expression::from_kind(ErrorExpression {
                        error: ctx.error("DEF expression missing").at(rhs.as_ref()).emit(),
                    })))
                })
            }
            DefRhsContextAll::Error(_) => {
                DefBody::Expression(Arc::new(Expression::from_kind(ErrorExpression {
                    error: ctx
                        .error("DEF right-hand side parse error")
                        .at(rhs.as_ref())
                        .emit(),
                })))
            }
        }
    }
}

impl DefStatement {
    fn from_name_and_rhs(
        name_cst: Rc<SimpleIdentifierContextAll<'static>>,
        rhs: Rc<DefRhsContextAll<'static>>,
        ctx: &mut ParseContext,
    ) -> Self {
        let name_iv = interval(name_cst.as_ref());
        let rhs_iv = interval(rhs.as_ref());
        let span = Span::from_range(
            *name_iv.start().min(rhs_iv.start())..=*name_iv.end().max(rhs_iv.end()),
        );
        let name = ParsedSimpleIdentifier::from_cst_with_context(name_cst.as_ref(), ctx).into();
        let body = DefBody::from_cst_with_context(rhs, ctx);
        DefStatement { span, name, body }
    }

    fn from_cst_with_context(
        def_cst: Rc<DefStatementContextAll<'static>>,
        ctx: &mut ParseContext,
    ) -> Option<Self> {
        match (def_cst.simpleIdentifier(), def_cst.defRhs()) {
            (Some(name), Some(rhs)) => Some(Self::from_name_and_rhs(name, rhs, ctx)),
            _ => None,
        }
    }

    pub fn fmt_indented(&self, f: &mut Formatter<'_>, indentation: usize) -> fmt::Result {
        pad(f, indentation)?;
        write!(f, "DEF {} = ", self.name)?;
        match &self.body {
            DefBody::Pipeline(p) => p.fmt_indented(f, indentation)?,
            DefBody::Expression(e) => write!(f, "{e}")?,
        }
        write!(f, ";")
    }
}

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

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

// ============================================================================
// Display / fmt_indented implementations
// ============================================================================

impl Query {
    pub fn fmt_indented(&self, f: &mut Formatter<'_>, indentation: usize) -> fmt::Result {
        self.kind.fmt_indented(f, indentation)
    }
}

impl Display for Query {
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        self.fmt_indented(f, 0)
    }
}

impl QueryKind {
    fn from_query_main_with_defs(
        main_cst: Rc<QueryMainContextAll<'static>>,
        query_span: Span,
        defs: Vec<DefStatement>,
        ctx: &mut ParseContext,
    ) -> Self {
        match main_cst.as_ref() {
            QueryMainContextAll::PipelineQueryMainContext(pctx) => match pctx.pipeline() {
                Some(pipeline_cst) => QueryKind::Valid(ValidQuery {
                    span: query_span,
                    defs,
                    main_pipeline: Arc::new(Pipeline::from_cst_with_context(pipeline_cst, ctx)),
                }),
                None => QueryKind::Error(ctx.error("Query missing its main pipeline").emit()),
            },
            QueryMainContextAll::ExpressionQueryMainContext(ectx) => match ectx.expression() {
                Some(expr_cst) => {
                    let expr = Expression::from_cst_with_context(expr_cst, ctx);
                    let pipeline = Pipeline::from(expr);
                    QueryKind::Valid(ValidQuery {
                        span: query_span,
                        defs,
                        main_pipeline: Arc::new(pipeline),
                    })
                }
                None => QueryKind::Error(ctx.error("Expression did not parse").emit()),
            },
            QueryMainContextAll::Error(_) => {
                let error = ctx
                    .error("query main body parse error")
                    .at(main_cst.as_ref())
                    .emit();
                if defs.is_empty() {
                    QueryKind::Error(error)
                } else {
                    let fallback_main = Pipeline::from(Expression::from_kind(ErrorExpression {
                        error: error.clone(),
                    }));
                    QueryKind::Valid(ValidQuery {
                        span: query_span,
                        defs,
                        main_pipeline: Arc::new(fallback_main),
                    })
                }
            }
        }
    }

    pub fn fmt_indented(&self, f: &mut Formatter<'_>, indentation: usize) -> fmt::Result {
        match self {
            QueryKind::Valid(v) => v.fmt_indented(f, indentation),
            QueryKind::Error(_) => write!(f, "?!"),
        }
    }
}

impl ValidQuery {
    pub fn fmt_indented(&self, f: &mut Formatter<'_>, indentation: usize) -> fmt::Result {
        for d in &self.defs {
            d.fmt_indented(f, indentation)?;
            writeln!(f)?;
            writeln!(f)?;
            pad(f, indentation)?;
        }
        self.main_pipeline.fmt_indented(f, indentation)
    }
}

impl Display for DefStatement {
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        write!(f, "DEF {} = ", self.name)?;
        match &self.body {
            DefBody::Pipeline(p) => write!(f, "{p}")?,
            DefBody::Expression(e) => write!(f, "{e}")?,
        }
        write!(f, ";")
    }
}

impl Display for ValidQuery {
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        self.fmt_indented(f, 0)
    }
}

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::QueryWithDefsContext(sctx) => {
                let defs: Vec<DefStatement> = sctx
                    .defStatement_all()
                    .into_iter()
                    .filter_map(|def_cst| DefStatement::from_cst_with_context(def_cst, ctx))
                    .collect();

                match sctx.queryMain() {
                    Some(main_cst) => QueryKind::from_query_main_with_defs(
                        main_cst,
                        query_span.clone().into(),
                        defs,
                        ctx,
                    ),
                    None => {
                        let error = ctx
                            .error("Query with DEFs missing its main body")
                            .at(cst.as_ref())
                            .emit();
                        let fallback_main =
                            Pipeline::from(Expression::from_kind(ErrorExpression {
                                error: error.clone(),
                            }));
                        QueryKind::Valid(ValidQuery {
                            span: query_span.clone().into(),
                            defs,
                            main_pipeline: Arc::new(fallback_main),
                        })
                    }
                }
            }
            QueryContextAll::QueryWithoutDefsContext(sctx) => match sctx.queryMain() {
                Some(main_cst) => QueryKind::from_query_main_with_defs(
                    main_cst,
                    query_span.clone().into(),
                    vec![],
                    ctx,
                ),
                None => QueryKind::Error(
                    ctx.error("Query missing its main body")
                        .at(cst.as_ref())
                        .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: Arc<Self>,
        ctx: &mut StatementTranslationContext,
    ) -> Self::Output {
        TypedStatement::from_ast_with_context(ast, ctx)
    }
}
