duckbridge 0.1.1

//! DDL/DML text emission for SQL Server 2019 (T-SQL) and PostgreSQL.
use super::model::*;
use crate::transpile::{self, Dialect as D};
use std::collections::HashMap;

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Flavor {
    Mssql,
    Postgres,
}

pub struct Emitter {
    pub flavor: Flavor,
    pub schema_map: HashMap<String, String>,
}

/// A value already rendered into a SQL literal by the reader.
pub type Literal = String;

impl Emitter {
    pub fn new(flavor: Flavor, schema_map: HashMap<String, String>) -> Self {
        Self { flavor, schema_map }
    }

    pub fn ident(&self, s: &str) -> String {
        match self.flavor {
            Flavor::Mssql => format!("[{}]", s.replace(']', "]]")),
            Flavor::Postgres => format!("\"{}\"", s.replace('"', "\"\"")),
        }
    }

    pub fn schema(&self, s: &str) -> String {
        self.schema_map.get(s).cloned().unwrap_or_else(|| s.to_string())
    }

    pub fn default_schema(&self) -> &'static str {
        match self.flavor {
            Flavor::Mssql => "dbo",
            Flavor::Postgres => "public",
        }
    }

    pub fn table_ref(&self, schema: &str, name: &str) -> String {
        format!("{}.{}", self.ident(&self.schema(schema)), self.ident(name))
    }

    pub fn str_lit(&self, s: &str) -> String {
        let body = s.replace('\'', "''");
        match self.flavor {
            Flavor::Mssql => format!("N'{body}'"),
            Flavor::Postgres => format!("'{body}'"),
        }
    }

    fn cname(&self, prefix: &str, t: &Table, cols: &[String]) -> String {
        let base = format!("{prefix}_{}_{}", t.name, cols.join("_"));
        let s: String = base.chars().map(|c| if c.is_alphanumeric() || c == '_' { c } else { '_' }).collect();
        s.chars().take(120).collect()
    }

    // ------------------------------------------------------------- defaults
    pub fn default_expr(&self, col: &mut Column) -> Option<String> {
        let d = col.default.clone()?.trim().to_string();
        if d.is_empty() || d.to_lowercase().contains("nextval(") {
            return None;
        }
        let low: String = d.to_lowercase().chars().filter(|c| !c.is_whitespace()).collect();
        let pg = self.flavor == Flavor::Postgres;
        match low.as_str() {
            "now()" | "current_timestamp" | "get_current_timestamp()" | "current_timestamp()" => {
                return Some(if pg { "now()" } else { "SYSDATETIME()" }.into())
            }
            "current_date" | "today()" | "current_date()" => return Some(if pg { "CURRENT_DATE" } else { "CAST(GETDATE() AS DATE)" }.into()),
            "uuid()" | "gen_random_uuid()" => return Some(if pg { "gen_random_uuid()" } else { "NEWID()" }.into()),
            "true" | "cast('t'asboolean)" | "cast('true'asboolean)" => return Some(if pg { "true" } else { "1" }.into()),
            "false" | "cast('f'asboolean)" | "cast('false'asboolean)" => return Some(if pg { "false" } else { "0" }.into()),
            _ => {}
        }
        // CAST(x AS type) -> x when x is a literal
        let mut inner = d.clone();
        if low.starts_with("cast(") && d.ends_with(')') {
            if let Some(i) = d.to_uppercase().rfind(" AS ") {
                inner = d[5..i].trim().to_string();
            }
        }
        if inner.parse::<f64>().is_ok() {
            return Some(inner);
        }
        if inner.starts_with('\'') && inner.ends_with('\'') && inner.len() >= 2 {
            return Some(if pg { inner } else { format!("N{inner}") });
        }
        match transpile::translate_expr(&d, D::DuckDb, self.dialect()) {
            Ok(e) if e.warnings.is_empty() => Some(e.sql),
            _ => {
                col.warnings.push(format!("default {d:?} not translated; dropped"));
                None
            }
        }
    }

    fn dialect(&self) -> D {
        match self.flavor {
            Flavor::Mssql => D::Tsql,
            Flavor::Postgres => D::Postgres,
        }
    }

    // ------------------------------------------------------------- DDL
    pub fn column_def(&self, col: &mut Column) -> String {
        let mut parts = vec![self.ident(&col.name), col.sql_type.clone()];
        if let Some((seed, inc)) = col.identity {
            parts.push(match self.flavor {
                Flavor::Mssql => format!("IDENTITY({seed},{inc})"),
                Flavor::Postgres => format!("GENERATED BY DEFAULT AS IDENTITY (START WITH {seed} INCREMENT BY {inc})"),
            });
        }
        parts.push(if col.nullable { "NULL" } else { "NOT NULL" }.into());
        if col.identity.is_none() {
            if let Some(d) = self.default_expr(col) {
                parts.push(format!("DEFAULT {d}"));
            }
        }
        parts.join(" ")
    }

    pub fn create_schema(&self, s: &str) -> String {
        let s = self.schema(s);
        match self.flavor {
            Flavor::Mssql => format!(
                "IF NOT EXISTS (SELECT 1 FROM sys.schemas WHERE name = N'{}')\n    EXEC(N'CREATE SCHEMA {}');",
                s.replace('\'', "''"), self.ident(&s).replace('\'', "''")
            ),
            Flavor::Postgres => format!("CREATE SCHEMA IF NOT EXISTS {};", self.ident(&s)),
        }
    }

    pub fn create_table(&self, t: &mut Table) -> String {
        self.create_table_opts(t, true)
    }

    pub fn pk_constraint(&self, t: &Table) -> Option<String> {
        let pk = t.pk()?;
        let cols: Vec<String> = pk.columns.iter().map(|c| self.ident(c)).collect();
        let clustered = if self.flavor == Flavor::Mssql { " CLUSTERED" } else { "" };
        Some(format!("CONSTRAINT {} PRIMARY KEY{clustered} ({})", self.ident(&self.cname("PK", t, &pk.columns)), cols.join(", ")))
    }

    /// `ALTER TABLE ... ADD CONSTRAINT PK ...` for heap-then-key loads.
    pub fn add_primary_key(&self, t: &Table) -> Option<String> {
        Some(format!("ALTER TABLE {} ADD {};", self.table_ref(&t.schema, &t.name), self.pk_constraint(t)?))
    }

    pub fn create_table_opts(&self, t: &mut Table, with_pk: bool) -> String {
        let tref = self.table_ref(&t.schema, &t.name);
        let pk = if with_pk { t.pk().cloned() } else { None };
        let mut lines: Vec<String> = t.columns.iter_mut().map(|c| self.column_def(c)).collect();
        if let Some(pk) = pk {
            let cols: Vec<String> = pk.columns.iter().map(|c| self.ident(c)).collect();
            let clustered = if self.flavor == Flavor::Mssql { " CLUSTERED" } else { "" };
            lines.push(format!("CONSTRAINT {} PRIMARY KEY{clustered} ({})", self.ident(&self.cname("PK", t, &pk.columns)), cols.join(", ")));
        }
        format!("CREATE TABLE {tref} (\n    {}\n);", lines.join(",\n    "))
    }

    pub fn drop_view_if_exists(&self, v: &View) -> String {
        format!("DROP VIEW IF EXISTS {};", self.table_ref(&v.schema, &v.name))
    }

    /// After loading explicit identity values, make the identity continue from the current maximum.
    pub fn reseed_identity(&self, t: &Table) -> Option<String> {
        if self.flavor != Flavor::Mssql || !t.has_identity() { return None; }
        let col = t.columns.iter().find(|c| c.identity.is_some())?;
        let tref = self.table_ref(&t.schema, &t.name);
        Some(format!("DECLARE @d2m_max BIGINT = (SELECT MAX({}) FROM {tref}); IF @d2m_max IS NOT NULL DBCC CHECKIDENT ('{}', RESEED, @d2m_max) WITH NO_INFOMSGS;", self.ident(&col.name), tref.replace('\'', "''")))
    }

    pub fn drop_table_if_exists(&self, t: &Table) -> String {
        let cascade = if self.flavor == Flavor::Postgres { " CASCADE" } else { "" };
        format!("DROP TABLE IF EXISTS {}{cascade};", self.table_ref(&t.schema, &t.name))
    }

    pub fn enum_checks(&self, t: &Table) -> Vec<String> {
        let mut out = vec![];
        for c in &t.columns {
            if let Some(vals) = &c.enum_values {
                let list: Vec<String> = vals.iter().map(|v| self.str_lit(v)).collect();
                out.push(format!(
                    "ALTER TABLE {} ADD CONSTRAINT {} CHECK ({} IN ({}));",
                    self.table_ref(&t.schema, &t.name), self.ident(&self.cname("CK_ENUM", t, &[c.name.clone()])), self.ident(&c.name), list.join(", ")
                ));
            }
        }
        out
    }

    /// Returns (statement, must_succeed). UNIQUE/FK must succeed; CHECK is best effort.
    pub fn post_constraints(&self, t: &Table) -> Vec<(String, bool)> {
        let tref = self.table_ref(&t.schema, &t.name);
        let mut out = vec![];
        for c in &t.constraints {
            let cols: Vec<String> = c.columns.iter().map(|x| self.ident(x)).collect();
            match c.kind {
                ConstraintKind::Unique => out.push((
                    format!("ALTER TABLE {tref} ADD CONSTRAINT {} UNIQUE ({});", self.ident(&self.cname("UQ", t, &c.columns)), cols.join(", ")), true)),
                ConstraintKind::ForeignKey => {
                    if let Some(rt) = &c.ref_table {
                        let rref = self.table_ref(c.ref_schema.as_deref().unwrap_or(&t.schema), rt);
                        let rcols: Vec<String> = (if c.ref_columns.is_empty() { &c.columns } else { &c.ref_columns }).iter().map(|x| self.ident(x)).collect();
                        out.push((format!("ALTER TABLE {tref} ADD CONSTRAINT {} FOREIGN KEY ({}) REFERENCES {rref} ({});",
                            self.ident(&self.cname("FK", t, &c.columns)), cols.join(", "), rcols.join(", ")), true));
                    }
                }
                ConstraintKind::Check => {
                    if let Some(expr) = &c.expression {
                        if let Ok(e) = transpile::translate_expr(expr, D::DuckDb, self.dialect()) {
                            if e.warnings.is_empty() {
                                let nm = if c.columns.is_empty() { vec!["expr".to_string()] } else { c.columns.clone() };
                                out.push((format!("ALTER TABLE {tref} ADD CONSTRAINT {} CHECK ({});", self.ident(&self.cname("CK", t, &nm)), e.sql), false));
                            }
                        }
                    }
                }
                ConstraintKind::PrimaryKey => {}
            }
        }
        out
    }

    pub fn create_indexes(&self, t: &Table) -> Vec<String> {
        let tref = self.table_ref(&t.schema, &t.name);
        let pk_cols = t.pk().map(|p| p.columns.clone());
        t.indexes.iter().filter(|i| Some(&i.columns) != pk_cols.as_ref()).map(|i| {
            let cols: Vec<String> = i.columns.iter().map(|c| self.ident(c)).collect();
            let u = if i.unique { "UNIQUE " } else { "" };
            match self.flavor {
                Flavor::Mssql => format!("CREATE {u}NONCLUSTERED INDEX {} ON {tref} ({});", self.ident(&i.name), cols.join(", ")),
                Flavor::Postgres => format!("CREATE {u}INDEX {} ON {tref} ({});", self.ident(&i.name), cols.join(", ")),
            }
        }).collect()
    }

    /// Translate a DuckDB view into the target dialect. Err carries the reason.
    pub fn create_view(&self, v: &View) -> Result<String, String> {
        let opts = transpile::Options { target_version: crate::core::typemap::mssql_version(), prelude: false, ..transpile::Options::default() };
        let out = transpile::translate_with(&v.sql, D::DuckDb, self.dialect(), &self.schema_map, &opts).map_err(|e| e.to_string())?;
        if !out.warnings.is_empty() {
            return Err(out.warnings.join("; "));
        }
        // DuckDB stores views as `CREATE VIEW name AS SELECT ...`; keep only the query body.
        let select = transpile::view_body(&out.sql, self.flavor == Flavor::Mssql);
        Ok(format!("CREATE VIEW {} AS\n{select};", self.table_ref(&v.schema, &v.name)))
    }

    pub fn comment(&self, t: &Table, col: Option<&Column>, text: &str) -> String {
        let txt = text.replace('\'', "''");
        match self.flavor {
            Flavor::Postgres => match col {
                Some(c) => format!("COMMENT ON COLUMN {}.{} IS '{txt}';", self.table_ref(&t.schema, &t.name), self.ident(&c.name)),
                None => format!("COMMENT ON TABLE {} IS '{txt}';", self.table_ref(&t.schema, &t.name)),
            },
            Flavor::Mssql => {
                let lvl2 = col.map(|c| format!(", @level2type=N'COLUMN', @level2name=N'{}'", c.name.replace('\'', "''"))).unwrap_or_default();
                format!("EXEC sys.sp_addextendedproperty @name=N'MS_Description', @value=N'{txt}', @level0type=N'SCHEMA', @level0name=N'{}', @level1type=N'TABLE', @level1name=N'{}'{lvl2};",
                    self.schema(&t.schema).replace('\'', "''"), t.name.replace('\'', "''"))
            }
        }
    }

    pub fn identity_insert(&self, t: &Table, on: bool) -> Option<String> {
        if self.flavor == Flavor::Mssql && t.has_identity() {
            Some(format!("SET IDENTITY_INSERT {} {};", self.table_ref(&t.schema, &t.name), if on { "ON" } else { "OFF" }))
        } else { None }
    }

    pub fn insert_batch(&self, t: &Table, rows: &[Vec<Literal>]) -> String {
        let cols: Vec<String> = t.columns.iter().map(|c| self.ident(&c.name)).collect();
        let vals: Vec<String> = rows.iter().map(|r| format!("({})", r.join(", "))).collect();
        format!("INSERT INTO {} ({}) VALUES\n{};", self.table_ref(&t.schema, &t.name), cols.join(", "), vals.join(",\n"))
    }

    pub fn max_rows_per_insert(&self) -> usize {
        match self.flavor { Flavor::Mssql => 1000, Flavor::Postgres => 5000 }
    }

    pub fn batch_terminator(&self) -> Option<&'static str> {
        match self.flavor { Flavor::Mssql => Some("GO"), Flavor::Postgres => None }
    }

    pub fn script_header(&self, cat: &Catalog, database: Option<&str>, create_db: bool) -> String {
        let mut h = vec![format!("-- Generated by duckbridge from {} (DuckDB {})", cat.source, cat.duckdb_version)];
        match self.flavor {
            Flavor::Mssql => {
                h.push("-- Target: Microsoft SQL Server 2019 or newer".into());
                h.push("SET NOCOUNT ON;".into());
                h.push("SET XACT_ABORT ON;".into());
                if let Some(db) = database {
                    let d = self.ident(db);
                    if create_db {
                        h.push(format!("IF DB_ID(N'{}') IS NULL CREATE DATABASE {d};", db.replace('\'', "''")));
                        h.push("GO".into());
                    }
                    h.push(format!("USE {d};"));
                    h.push("GO".into());
                }
            }
            Flavor::Postgres => {
                h.push("-- Target: PostgreSQL 12 or newer. Restore with: psql -d <db> -v ON_ERROR_STOP=1 -f <this file>".into());
                if let (Some(db), true) = (database, create_db) {
                    h.push(format!("-- Create the database first:  createdb {db}"));
                }
                h.push("SET client_encoding = 'UTF8';".into());
                h.push("SET standard_conforming_strings = on;".into());
            }
        }
        h.join("\n")
    }

    /// Wrap a best-effort statement so a failure prints a warning instead of aborting the script.
    pub fn guarded(&self, stmt: &str) -> String {
        let inner = stmt.replace('\'', "''");
        let short: String = inner.chars().take(3000).collect();
        match self.flavor {
            Flavor::Mssql => format!(
                "BEGIN TRY\n    EXEC(N'{inner}');\nEND TRY\nBEGIN CATCH\n    PRINT N'WARNING: statement failed: ' + ERROR_MESSAGE();\n    PRINT N'{short}';\nEND CATCH;"),
            Flavor::Postgres => format!(
                "DO $d2m$ BEGIN\n    EXECUTE '{inner}';\nEXCEPTION WHEN OTHERS THEN\n    RAISE WARNING 'statement failed: % -- %', SQLERRM, '{short}';\nEND $d2m$;"),
        }
    }
}