use std::sync::atomic::{AtomicUsize, Ordering};

use crate::error::ExportError;
use config::database::{DatabaseConfig, DatabaseRole as ConfigRole};
use sea_orm::prelude::async_trait::async_trait;
use sea_orm::{Database, DatabaseConnection, DbBackend, DbErr, ExecResult, QueryResult, Statement};

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum DatabaseRole {
    Primary,
    Replica,
}

impl From<ConfigRole> for DatabaseRole {
    fn from(role: ConfigRole) -> Self {
        match role {
            ConfigRole::Primary => DatabaseRole::Primary,
            ConfigRole::Replica => DatabaseRole::Replica,
        }
    }
}

impl From<DatabaseRole> for ConfigRole {
    fn from(role: DatabaseRole) -> Self {
        match role {
            DatabaseRole::Primary => ConfigRole::Primary,
            DatabaseRole::Replica => ConfigRole::Replica,
        }
    }
}

pub struct AppDatabase {
    config: DatabaseConfig,
    primary: DatabaseConnection,
    replicas: Vec<DatabaseConnection>,
    round_robin: AtomicUsize,
}

impl AppDatabase {
    pub async fn connect(config: DatabaseConfig) -> Result<Self, ExportError> {
        config.validate().map_err(ExportError::InvalidConfig)?;

        let primary_opts = config
            .primary_connect_options()
            .map_err(ExportError::InvalidConfig)?;
        let primary = Database::connect(primary_opts).await?;

        let mut replicas = Vec::with_capacity(config.replica_count());
        for i in 0..config.replica_count() {
            let opts = config
                .replica_connect_options(i)
                .map_err(ExportError::InvalidConfig)?;
            let conn = Database::connect(opts).await?;
            replicas.push(conn);
        }

        Ok(Self {
            config,
            primary,
            replicas,
            round_robin: AtomicUsize::new(0),
        })
    }

    pub fn config(&self) -> &DatabaseConfig {
        &self.config
    }

    pub fn primary(&self) -> &DatabaseConnection {
        &self.primary
    }

    pub fn replicas(&self) -> &[DatabaseConnection] {
        &self.replicas
    }

    pub fn replica(&self, index: usize) -> Option<&DatabaseConnection> {
        self.replicas.get(index)
    }

    pub fn replica_count(&self) -> usize {
        self.replicas.len()
    }

    pub fn has_replicas(&self) -> bool {
        !self.replicas.is_empty()
    }

    pub fn write(&self) -> &DatabaseConnection {
        &self.primary
    }

    pub fn read(&self) -> &DatabaseConnection {
        self.pick_replica().unwrap_or(&self.primary)
    }

    pub fn for_role(&self, role: DatabaseRole) -> &DatabaseConnection {
        match role {
            DatabaseRole::Primary => self.write(),
            DatabaseRole::Replica => self.read(),
        }
    }

    pub fn route_for(&self, sql: &str) -> &DatabaseConnection {
        if is_write_sql(sql) {
            self.write()
        } else {
            self.read()
        }
    }

    fn pick_replica(&self) -> Option<&DatabaseConnection> {
        if self.replicas.is_empty() {
            return None;
        }
        let idx = self.round_robin.fetch_add(1, Ordering::Relaxed) % self.replicas.len();
        Some(&self.replicas[idx])
    }
}

fn strip_sql_lead(sql: &str) -> &str {
    let mut s = sql.trim_start();
    loop {
        if let Some(rest) = s.strip_prefix("--") {
            let end = rest.find('\n').unwrap_or(rest.len());
            s = rest[end..].trim_start();
            continue;
        }
        if let Some(rest) = s.strip_prefix("/*") {
            let end = rest.find("*/").map(|i| i + 2).unwrap_or(rest.len());
            s = rest[end..].trim_start();
            continue;
        }
        break;
    }
    s
}

fn first_keyword(sql: &str) -> Option<String> {
    let lead = strip_sql_lead(sql);
    let head: String = lead
        .chars()
        .take_while(|c| c.is_ascii_alphanumeric() || *c == '_')
        .collect();
    if head.is_empty() {
        None
    } else {
        Some(head.to_uppercase())
    }
}

fn is_write_sql(sql: &str) -> bool {
    let Some(kw) = first_keyword(sql) else {
        return true;
    };
    if kw == "WITH" {
        return with_clause_is_write(sql);
    }
    matches!(
        kw.as_str(),
        "INSERT"
            | "UPDATE"
            | "DELETE"
            | "MERGE"
            | "UPSERT"
            | "REPLACE"
            | "TRUNCATE"
            | "CREATE"
            | "DROP"
            | "ALTER"
            | "RENAME"
            | "GRANT"
            | "REVOKE"
            | "COMMIT"
            | "ROLLBACK"
            | "SAVEPOINT"
            | "LOCK"
            | "UNLOCK"
            | "COPY"
            | "CALL"
            | "SET"
            | "RESET"
            | "DO"
    )
}

fn with_clause_is_write(sql: &str) -> bool {
    let upper = sql.to_uppercase();
    contains_top_level(&upper, "INSERT")
        || contains_top_level(&upper, "UPDATE")
        || contains_top_level(&upper, "DELETE")
        || contains_top_level(&upper, "MERGE")
}

fn contains_top_level(haystack: &str, needle: &str) -> bool {
    let bytes = haystack.as_bytes();
    let n = needle.as_bytes();
    let mut depth: i32 = 0;
    let mut in_string = false;
    let mut i = 0;
    while i < bytes.len() {
        let c = bytes[i] as char;
        if in_string {
            if c == '\'' {
                if i + 1 < bytes.len() && bytes[i + 1] == b'\'' {
                    i += 2;
                    continue;
                }
                in_string = false;
            }
            i += 1;
            continue;
        }
        match c {
            '\'' => in_string = true,
            '(' => depth += 1,
            ')' => depth -= 1,
            _ if depth == 0 && i + n.len() <= bytes.len() && &bytes[i..i + n.len()] == n => {
                let before_ok = i == 0
                    || !(bytes[i - 1] as char).is_ascii_alphanumeric() && bytes[i - 1] != b'_';
                let after_ok = i + n.len() == bytes.len()
                    || !(bytes[i + n.len()] as char).is_ascii_alphanumeric()
                        && bytes[i + n.len()] != b'_';
                if before_ok && after_ok {
                    return true;
                }
            }
            _ => {}
        }
        i += 1;
    }
    false
}

#[async_trait]
impl sea_orm::ConnectionTrait for AppDatabase {
    fn get_database_backend(&self) -> DbBackend {
        self.primary.get_database_backend()
    }

    async fn execute_raw(&self, stmt: Statement) -> Result<ExecResult, DbErr> {
        self.route_for(&stmt.sql).execute_raw(stmt).await
    }

    async fn execute_unprepared(&self, sql: &str) -> Result<ExecResult, DbErr> {
        self.route_for(sql).execute_unprepared(sql).await
    }

    async fn query_one_raw(&self, stmt: Statement) -> Result<Option<QueryResult>, DbErr> {
        self.route_for(&stmt.sql).query_one_raw(stmt).await
    }

    async fn query_all_raw(&self, stmt: Statement) -> Result<Vec<QueryResult>, DbErr> {
        self.route_for(&stmt.sql).query_all_raw(stmt).await
    }
}