// Copyright 2022 RisingWave Labs
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use std::collections::HashSet;

use either::Either;
use pgwire::pg_response::StatementType;
use risingwave_common::catalog::{ColumnCatalog, ColumnDesc};
use risingwave_pb::ddl_service::TableJobType;
use risingwave_sqlparser::ast::{ColumnDef, ObjectName, OnConflict, Query, Statement};

use super::{HandlerArgs, RwPgResponse};
use crate::binder::BoundStatement;
use crate::error::{ErrorCode, Result};
use crate::handler::create_table::{
    ColumnIdGenerator, CreateTableProps, gen_create_table_plan_without_source,
};
use crate::handler::query::handle_query;
use crate::handler::util::{
    LongRunningNotificationAction, execute_with_long_running_notification,
    reject_internal_table_dependencies,
};
use crate::stream_fragmenter::GraphJobType;
use crate::{Binder, OptimizerContext, build_graph};

pub async fn handle_create_as(
    handler_args: HandlerArgs,
    table_name: ObjectName,
    if_not_exists: bool,
    query: Box<Query>,
    column_defs: Vec<ColumnDef>,
    append_only: bool,
    on_conflict: Option<OnConflict>,
    with_version_columns: Vec<String>,
    ast_engine: risingwave_sqlparser::ast::Engine,
) -> Result<RwPgResponse> {
    if column_defs.iter().any(|column| column.data_type.is_some()) {
        return Err(ErrorCode::InvalidInputSyntax(
            "Should not specify data type in CREATE TABLE AS".into(),
        )
        .into());
    }
    let engine = match ast_engine {
        risingwave_sqlparser::ast::Engine::Hummock => risingwave_common::catalog::Engine::Hummock,
        risingwave_sqlparser::ast::Engine::Iceberg => risingwave_common::catalog::Engine::Iceberg,
    };

    let session = handler_args.session.clone();

    if let Either::Right(resp) = session.check_relation_name_duplicated(
        table_name.clone(),
        StatementType::CREATE_TABLE,
        if_not_exists,
    )? {
        return Ok(resp);
    }

    let mut col_id_gen = ColumnIdGenerator::new_initial();

    // Generate catalog descs from query
    let (mut columns, dependent_relations): (Vec<_>, HashSet<_>) = {
        let mut binder = Binder::new_for_batch(&session);
        let bound = binder.bind(Statement::Query(query.clone()))?;
        if let BoundStatement::Query(query) = bound {
            // Create ColumnCatelog by Field
            let columns = query
                .schema()
                .fields()
                .iter()
                .map(|field| ColumnCatalog {
                    column_desc: ColumnDesc::from_field_without_column_id(field),
                    is_hidden: false,
                })
                .collect();
            (columns, binder.included_relations().clone())
        } else {
            unreachable!()
        }
    };

    reject_internal_table_dependencies(&session, &dependent_relations, "CREATE TABLE AS")?;

    // Generate column id.
    for c in &mut columns {
        col_id_gen.generate(c)?;
    }

    if column_defs.len() > columns.len() {
        return Err(ErrorCode::InvalidInputSyntax(
            "too many column names were specified".to_owned(),
        )
        .into());
    }

    // Override column name if it specified in create statement.
    column_defs.iter().enumerate().for_each(|(idx, column)| {
        columns[idx].column_desc.name = column.name.real_value();
    });

    let (graph, source, table) = {
        let context = OptimizerContext::from_handler_args(handler_args.clone());
        let (_, secret_refs, connection_refs) = context.with_options().clone().into_parts();
        if !secret_refs.is_empty() || !connection_refs.is_empty() {
            return Err(crate::error::ErrorCode::InvalidParameterValue(
                "Secret reference and Connection reference are not allowed in options for CREATE TABLE AS".to_owned(),
            )
            .into());
        }
        let (plan, table) = gen_create_table_plan_without_source(
            context,
            table_name.clone(),
            columns,
            vec![],
            vec![],
            vec![], // No watermark should be defined in for `CREATE TABLE AS`
            col_id_gen.into_version(),
            CreateTableProps {
                // Note: by providing and persisting an empty definition, querying the definition of the table
                // will hit the purification logic, which will construct it based on the catalog.
                definition: "".to_owned(),
                append_only,
                on_conflict: on_conflict.into(),
                with_version_columns,
                webhook_info: None,
                engine,
            },
        )?;
        let graph = build_graph(plan, Some(GraphJobType::Table))?;

        (graph, None, table)
    };

    tracing::trace!(
        "name={}, graph=\n{}",
        table_name,
        serde_json::to_string_pretty(&graph).unwrap()
    );

    let catalog_writer = session.catalog_writer()?;
    execute_with_long_running_notification(
        catalog_writer.create_table(
            source,
            table.to_prost(),
            graph,
            TableJobType::Unspecified,
            if_not_exists,
            HashSet::default(),
        ),
        &session,
        "CREATE TABLE AS",
        LongRunningNotificationAction::DiagnoseBarrierLatency,
    )
    .await?;

    // Generate insert
    let insert = Statement::Insert {
        table_name,
        columns: vec![],
        source: query,
        returning: vec![],
    };

    handle_query(handler_args, insert, vec![]).await
}
