// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you 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.

//! Benchmark derived from IMDB dataset.

use datafusion::{
    arrow::datatypes::{DataType, Field, Schema},
    common::plan_err,
    error::Result,
};
mod convert;
pub use convert::ConvertOpt;

use std::fs;
mod run;
pub use run::RunOpt;

// we have 21 tables in the IMDB dataset
pub const IMDB_TABLES: &[&str] = &[
    "aka_name",
    "aka_title",
    "cast_info",
    "char_name",
    "comp_cast_type",
    "company_name",
    "company_type",
    "complete_cast",
    "info_type",
    "keyword",
    "kind_type",
    "link_type",
    "movie_companies",
    "movie_info_idx",
    "movie_keyword",
    "movie_link",
    "name",
    "role_type",
    "title",
    "movie_info",
    "person_info",
];

pub const IMDB_QUERY_START_ID: usize = 1;
pub const IMDB_QUERY_END_ID: usize = 113;

/// Get the schema for the IMDB dataset tables
/// see benchmarks/data/imdb/schematext.sql
pub fn get_imdb_table_schema(table: &str) -> Schema {
    match table {
        "aka_name" => Schema::new(vec![
            Field::new("id", DataType::UInt32, false),
            Field::new("person_id", DataType::Int32, false),
            Field::new("name", DataType::Utf8, true),
            Field::new("imdb_index", DataType::Utf8, true),
            Field::new("name_pcode_cf", DataType::Utf8, true),
            Field::new("name_pcode_nf", DataType::Utf8, true),
            Field::new("surname_pcode", DataType::Utf8, true),
            Field::new("md5sum", DataType::Utf8, true),
        ]),
        "aka_title" => Schema::new(vec![
            Field::new("id", DataType::UInt32, false),
            Field::new("movie_id", DataType::Int32, false),
            Field::new("title", DataType::Utf8, true),
            Field::new("imdb_index", DataType::Utf8, true),
            Field::new("kind_id", DataType::Int32, false),
            Field::new("production_year", DataType::Int32, true),
            Field::new("phonetic_code", DataType::Utf8, true),
            Field::new("episode_of_id", DataType::Int32, true),
            Field::new("season_nr", DataType::Int32, true),
            Field::new("episode_nr", DataType::Int32, true),
            Field::new("note", DataType::Utf8, true),
            Field::new("md5sum", DataType::Utf8, true),
        ]),
        "cast_info" => Schema::new(vec![
            Field::new("id", DataType::UInt32, false),
            Field::new("person_id", DataType::Int32, false),
            Field::new("movie_id", DataType::Int32, false),
            Field::new("person_role_id", DataType::Int32, true),
            Field::new("note", DataType::Utf8, true),
            Field::new("nr_order", DataType::Int32, true),
            Field::new("role_id", DataType::Int32, false),
        ]),
        "char_name" => Schema::new(vec![
            Field::new("id", DataType::UInt32, false),
            Field::new("name", DataType::Utf8, false),
            Field::new("imdb_index", DataType::Utf8, true),
            Field::new("imdb_id", DataType::Int32, true),
            Field::new("name_pcode_nf", DataType::Utf8, true),
            Field::new("surname_pcode", DataType::Utf8, true),
            Field::new("md5sum", DataType::Utf8, true),
        ]),
        "comp_cast_type" => Schema::new(vec![
            Field::new("id", DataType::UInt32, false),
            Field::new("kind", DataType::Utf8, false),
        ]),
        "company_name" => Schema::new(vec![
            Field::new("id", DataType::UInt32, false),
            Field::new("name", DataType::Utf8, false),
            Field::new("country_code", DataType::Utf8, true),
            Field::new("imdb_id", DataType::Int32, true),
            Field::new("name_pcode_nf", DataType::Utf8, true),
            Field::new("name_pcode_sf", DataType::Utf8, true),
            Field::new("md5sum", DataType::Utf8, true),
        ]),
        "company_type" => Schema::new(vec![
            Field::new("id", DataType::UInt32, false),
            Field::new("kind", DataType::Utf8, true),
        ]),
        "complete_cast" => Schema::new(vec![
            Field::new("id", DataType::UInt32, false),
            Field::new("movie_id", DataType::Int32, true),
            Field::new("subject_id", DataType::Int32, false),
            Field::new("status_id", DataType::Int32, false),
        ]),
        "info_type" => Schema::new(vec![
            Field::new("id", DataType::UInt32, false),
            Field::new("info", DataType::Utf8, false),
        ]),
        "keyword" => Schema::new(vec![
            Field::new("id", DataType::UInt32, false),
            Field::new("keyword", DataType::Utf8, false),
            Field::new("phonetic_code", DataType::Utf8, true),
        ]),
        "kind_type" => Schema::new(vec![
            Field::new("id", DataType::UInt32, false),
            Field::new("kind", DataType::Utf8, true),
        ]),
        "link_type" => Schema::new(vec![
            Field::new("id", DataType::UInt32, false),
            Field::new("link", DataType::Utf8, false),
        ]),
        "movie_companies" => Schema::new(vec![
            Field::new("id", DataType::UInt32, false),
            Field::new("movie_id", DataType::Int32, false),
            Field::new("company_id", DataType::Int32, false),
            Field::new("company_type_id", DataType::Int32, false),
            Field::new("note", DataType::Utf8, true),
        ]),
        "movie_info_idx" => Schema::new(vec![
            Field::new("id", DataType::UInt32, false),
            Field::new("movie_id", DataType::Int32, false),
            Field::new("info_type_id", DataType::Int32, false),
            Field::new("info", DataType::Utf8, false),
            Field::new("note", DataType::Utf8, true),
        ]),
        "movie_keyword" => Schema::new(vec![
            Field::new("id", DataType::UInt32, false),
            Field::new("movie_id", DataType::Int32, false),
            Field::new("keyword_id", DataType::Int32, false),
        ]),
        "movie_link" => Schema::new(vec![
            Field::new("id", DataType::UInt32, false),
            Field::new("movie_id", DataType::Int32, false),
            Field::new("linked_movie_id", DataType::Int32, false),
            Field::new("link_type_id", DataType::Int32, false),
        ]),
        "name" => Schema::new(vec![
            Field::new("id", DataType::UInt32, false),
            Field::new("name", DataType::Utf8, false),
            Field::new("imdb_index", DataType::Utf8, true),
            Field::new("imdb_id", DataType::Int32, true),
            Field::new("gender", DataType::Utf8, true),
            Field::new("name_pcode_cf", DataType::Utf8, true),
            Field::new("name_pcode_nf", DataType::Utf8, true),
            Field::new("surname_pcode", DataType::Utf8, true),
            Field::new("md5sum", DataType::Utf8, true),
        ]),
        "role_type" => Schema::new(vec![
            Field::new("id", DataType::UInt32, false),
            Field::new("role", DataType::Utf8, false),
        ]),
        "title" => Schema::new(vec![
            Field::new("id", DataType::UInt32, false),
            Field::new("title", DataType::Utf8, false),
            Field::new("imdb_index", DataType::Utf8, true),
            Field::new("kind_id", DataType::Int32, false),
            Field::new("production_year", DataType::Int32, true),
            Field::new("imdb_id", DataType::Int32, true),
            Field::new("phonetic_code", DataType::Utf8, true),
            Field::new("episode_of_id", DataType::Int32, true),
            Field::new("season_nr", DataType::Int32, true),
            Field::new("episode_nr", DataType::Int32, true),
            Field::new("series_years", DataType::Utf8, true),
            Field::new("md5sum", DataType::Utf8, true),
        ]),
        "movie_info" => Schema::new(vec![
            Field::new("id", DataType::UInt32, false),
            Field::new("movie_id", DataType::Int32, false),
            Field::new("info_type_id", DataType::Int32, false),
            Field::new("info", DataType::Utf8, false),
            Field::new("note", DataType::Utf8, true),
        ]),
        "person_info" => Schema::new(vec![
            Field::new("id", DataType::UInt32, false),
            Field::new("person_id", DataType::Int32, false),
            Field::new("info_type_id", DataType::Int32, false),
            Field::new("info", DataType::Utf8, false),
            Field::new("note", DataType::Utf8, true),
        ]),
        _ => unimplemented!("Schema for table {} is not implemented", table),
    }
}

/// Get the SQL statements from the specified query file
pub fn get_query_sql(query: &str) -> Result<Vec<String>> {
    let possibilities = vec![
        format!("queries/imdb/{query}.sql"),
        format!("benchmarks/queries/imdb/{query}.sql"),
    ];
    let mut errors = vec![];
    for filename in possibilities {
        match fs::read_to_string(&filename) {
            Ok(contents) => {
                return Ok(contents
                    .split(';')
                    .map(|s| s.trim())
                    .filter(|s| !s.is_empty())
                    .map(|s| s.to_string())
                    .collect());
            }
            Err(e) => errors.push(format!("{filename}: {e}")),
        };
    }
    plan_err!("invalid query. Could not find query: {:?}", errors)
}
