//! Query builder for constructing Query AST nodes
//!
//! This module provides a fluent API for building queries with WITH clauses.
//!
//! The builder uses type-state to ensure the main pipeline is set exactly once:
//! - `QueryBuilder<NoMain>`: Can add WITH clauses, must call `.main()` before `.build()`
//! - `QueryBuilder<HasMain>`: Can add WITH clauses, can call `.build()`

use std::marker::PhantomData;
use std::rc::Rc;

use crate::tree::ast::identifier::{Identifier, ParsedIdentifier};
use crate::tree::ast::node::Span;
use crate::tree::ast::pipeline::Pipeline;
use crate::tree::ast::query::{Query, ValidQuery, WithClause};
use crate::tree::typed_ast::pipeline::TypedPipeline;

use super::pipeline::PipelineBuilder;

/// Trait for types that can be converted to `Rc<Pipeline>` for use in query building.
pub trait IntoPipelineRc {
    fn into_pipeline_rc(self) -> Rc<Pipeline>;
}

impl IntoPipelineRc for Pipeline {
    fn into_pipeline_rc(self) -> Rc<Pipeline> {
        Rc::new(self)
    }
}

impl IntoPipelineRc for Rc<Pipeline> {
    fn into_pipeline_rc(self) -> Rc<Pipeline> {
        self
    }
}

impl IntoPipelineRc for PipelineBuilder {
    fn into_pipeline_rc(self) -> Rc<Pipeline> {
        Rc::new(self.build())
    }
}

impl IntoPipelineRc for Rc<TypedPipeline> {
    fn into_pipeline_rc(self) -> Rc<Pipeline> {
        self.ast.clone()
    }
}

impl IntoPipelineRc for &Rc<TypedPipeline> {
    fn into_pipeline_rc(self) -> Rc<Pipeline> {
        self.ast.clone()
    }
}

/// Marker type: main pipeline not yet set
#[derive(Debug)]
pub struct NoMain;

/// Marker type: main pipeline has been set
#[derive(Debug)]
pub struct HasMain;

/// Builder for Query with type-state tracking
///
/// Use `query()` to start building without a main pipeline,
/// or `query_with_main(pipeline)` to start with a main pipeline already set.
#[derive(Debug)]
pub struct QueryBuilder<S = NoMain> {
    with_clauses: Vec<WithClause>,
    main_pipeline: Option<Rc<Pipeline>>,
    _marker: PhantomData<S>,
}

impl QueryBuilder<NoMain> {
    /// Create a new query builder without a main pipeline
    pub fn new() -> Self {
        Self {
            with_clauses: Vec::new(),
            main_pipeline: None,
            _marker: PhantomData,
        }
    }

    /// Set the main pipeline
    pub fn main(self, pipeline: impl IntoPipelineRc) -> QueryBuilder<HasMain> {
        QueryBuilder {
            with_clauses: self.with_clauses,
            main_pipeline: Some(pipeline.into_pipeline_rc()),
            _marker: PhantomData,
        }
    }

    /// Add a WITH clause
    pub fn with(mut self, name: impl Into<Identifier>, pipeline: impl IntoPipelineRc) -> Self {
        self.with_clauses.push(WithClause {
            span: Span::NONE,
            name: ParsedIdentifier::Valid(name.into()),
            pipeline: pipeline.into_pipeline_rc(),
        });
        self
    }
}

impl Default for QueryBuilder<NoMain> {
    fn default() -> Self {
        Self::new()
    }
}

impl QueryBuilder<HasMain> {
    /// Add a WITH clause
    pub fn with(mut self, name: impl Into<Identifier>, pipeline: impl IntoPipelineRc) -> Self {
        self.with_clauses.push(WithClause {
            span: Span::NONE,
            name: ParsedIdentifier::Valid(name.into()),
            pipeline: pipeline.into_pipeline_rc(),
        });
        self
    }

    pub fn build(self) -> Query {
        Query {
            span: Span::NONE,
            kind: ValidQuery {
                span: Span::NONE,
                with_clauses: self.with_clauses,
                main_pipeline: self
                    .main_pipeline
                    .expect("HasMain state guarantees main_pipeline is Some"),
            }
            .into(),
        }
    }
}

impl QueryBuilder<NoMain> {
    /// Merge a Query's CTEs into this builder, then add its main as a new WITH clause.
    ///
    /// Use this when transforming existing WITH clauses that may have generated
    /// additional CTEs during transformation.
    pub fn merge_as_cte(mut self, query: Query, name: impl Into<Identifier>) -> Self {
        if let crate::tree::ast::query::QueryKind::Valid(valid) = query.kind {
            // Add all CTEs from the query
            self.with_clauses.extend(valid.with_clauses);
            // Add the query's main pipeline as a new CTE with the given name
            self.with_clauses.push(WithClause {
                span: Span::NONE,
                name: ParsedIdentifier::Valid(name.into()),
                pipeline: valid.main_pipeline,
            });
        }
        self
    }

    /// Merge a Query's CTEs into this builder, using its main as this builder's main.
    ///
    /// Use this for the final main pipeline transformation.
    pub fn merge_as_main(mut self, query: Query) -> Query {
        if let crate::tree::ast::query::QueryKind::Valid(valid) = query.kind {
            // Add all CTEs from the query
            self.with_clauses.extend(valid.with_clauses);
            // Use the query's main as our main
            Query {
                span: Span::NONE,
                kind: ValidQuery {
                    span: Span::NONE,
                    with_clauses: self.with_clauses,
                    main_pipeline: valid.main_pipeline,
                }
                .into(),
            }
        } else {
            // Error query - just return it as-is
            query
        }
    }
}

/// Create a query builder
///
/// Use `.with()` to add CTEs, then `.main()` to set the main pipeline,
/// then `.build()` to create the Query.
///
/// ```ignore
/// query()
///     .with("cte1", pipeline1)
///     .with("cte2", pipeline2)
///     .main(main_pipeline)
///     .build()
/// ```
pub fn query() -> QueryBuilder<NoMain> {
    QueryBuilder::new()
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::tree::builder::pipeline;

    #[test]
    fn test_query_builder_main_then_build() {
        let q = query()
            .main(pipeline().from(|f| f.table_reference("events")))
            .build();

        match &q.kind {
            crate::tree::ast::query::QueryKind::Valid(valid) => {
                assert!(valid.with_clauses.is_empty());
            }
            _ => panic!("Expected valid query"),
        }
    }

    #[test]
    fn test_query_builder_with_then_main() {
        let q = query()
            .with("cte1", pipeline().from(|f| f.table_reference("events")))
            .with("cte2", pipeline().from(|f| f.table_reference("logs")))
            .main(pipeline().from(|f| f.table_reference("cte1")))
            .build();

        match &q.kind {
            crate::tree::ast::query::QueryKind::Valid(valid) => {
                assert_eq!(valid.with_clauses.len(), 2);
            }
            _ => panic!("Expected valid query"),
        }
    }

    #[test]
    fn test_query_builder_main_then_with() {
        let q = query()
            .main(pipeline().from(|f| f.table_reference("cte1")))
            .with("cte1", pipeline().from(|f| f.table_reference("events")))
            .build();

        match &q.kind {
            crate::tree::ast::query::QueryKind::Valid(valid) => {
                assert_eq!(valid.with_clauses.len(), 1);
            }
            _ => panic!("Expected valid query"),
        }
    }

    #[test]
    fn test_query_with_main_shorthand() {
        let q = query()
            .main(pipeline().from(|f| f.table_reference("events")))
            .with("cte1", pipeline().from(|f| f.table_reference("logs")))
            .build();

        match &q.kind {
            crate::tree::ast::query::QueryKind::Valid(valid) => {
                assert_eq!(valid.with_clauses.len(), 1);
            }
            _ => panic!("Expected valid query"),
        }
    }
}
