use crate::error::{TuskError, TuskResult};
use crate::models::ai::{
    AiProvider, AiSettings, FireworksChatRequest, FireworksChatResponse, FireworksModelsResponse,
    FireworksResponseFormat, OllamaChatMessage, OllamaChatRequest, OllamaChatResponse, OllamaModel,
    OllamaTagsResponse,
};
use crate::state::{AppState, DbFlavor};
use sqlx::Row;
use std::collections::{BTreeMap, HashMap};
use std::fs;
use std::sync::Arc;
use std::time::Duration;
use tauri::{AppHandle, Manager, State};

const MAX_RETRIES: u32 = 2;
const RETRY_DELAY_MS: u64 = 1000;
const FIREWORKS_BASE_URL: &str = "https://api.fireworks.ai/inference/v1";
const OPENROUTER_BASE_URL: &str = "https://openrouter.ai/api/v1";
/// Optional attribution headers OpenRouter uses for its public app leaderboard.
/// Harmless to send on every request; ignored by other OpenAI-compatible APIs.
const OPENROUTER_HEADERS: &[(&str, &str)] = &[
    ("HTTP-Referer", "https://github.com/codelab/tusk"),
    ("X-Title", "Tusk"),
];

fn http_client() -> &'static reqwest::Client {
    use std::sync::LazyLock;
    static CLIENT: LazyLock<reqwest::Client> = LazyLock::new(|| {
        reqwest::Client::builder()
            .connect_timeout(Duration::from_secs(5))
            .timeout(Duration::from_secs(300))
            .build()
            .unwrap_or_default()
    });
    &CLIENT
}

fn get_ai_settings_path(app: &AppHandle) -> TuskResult<std::path::PathBuf> {
    let dir = app
        .path()
        .app_data_dir()
        .map_err(|e| TuskError::Config(e.to_string()))?;
    fs::create_dir_all(&dir)?;
    Ok(dir.join("ai_settings.json"))
}

#[tauri::command]
pub async fn get_ai_settings(app: AppHandle) -> TuskResult<AiSettings> {
    let path = get_ai_settings_path(&app)?;
    if !path.exists() {
        return Ok(AiSettings::default());
    }
    let data = fs::read_to_string(&path)?;
    let settings: AiSettings = serde_json::from_str(&data)?;
    Ok(settings)
}

#[tauri::command]
pub async fn save_ai_settings(
    app: AppHandle,
    state: State<'_, Arc<AppState>>,
    settings: AiSettings,
) -> TuskResult<()> {
    let path = get_ai_settings_path(&app)?;
    let data = serde_json::to_string_pretty(&settings)?;
    fs::write(&path, data)?;
    // Update in-memory cache
    *state.ai_settings.write().await = Some(settings);
    Ok(())
}

#[tauri::command]
pub async fn list_ollama_models(ollama_url: String) -> TuskResult<Vec<OllamaModel>> {
    let url = format!("{}/api/tags", ollama_url.trim_end_matches('/'));
    let resp =
        http_client().get(&url).send().await.map_err(|e| {
            TuskError::Ai(format!("Cannot connect to Ollama at {}: {}", ollama_url, e))
        })?;

    if !resp.status().is_success() {
        let status = resp.status();
        let body = resp.text().await.unwrap_or_default();
        return Err(TuskError::Ai(format!(
            "Ollama error ({}): {}",
            status, body
        )));
    }

    let tags: OllamaTagsResponse = resp
        .json()
        .await
        .map_err(|e| TuskError::Ai(format!("Failed to parse Ollama response: {}", e)))?;

    Ok(tags.models)
}

#[tauri::command]
pub async fn list_fireworks_models(api_key: String) -> TuskResult<Vec<OllamaModel>> {
    list_openai_compatible_models(FIREWORKS_BASE_URL, &api_key, "Fireworks", &[]).await
}

#[tauri::command]
pub async fn list_openrouter_models(api_key: String) -> TuskResult<Vec<OllamaModel>> {
    list_openai_compatible_models(OPENROUTER_BASE_URL, &api_key, "OpenRouter", OPENROUTER_HEADERS)
        .await
}

/// List the available chat models for any OpenAI-compatible provider via its
/// `GET {base_url}/models` endpoint. `extra_headers` carries provider-specific
/// attribution headers (OpenRouter recommends `HTTP-Referer`/`X-Title`).
async fn list_openai_compatible_models(
    base_url: &str,
    api_key: &str,
    provider_label: &str,
    extra_headers: &[(&str, &str)],
) -> TuskResult<Vec<OllamaModel>> {
    let key = api_key.trim();
    if key.is_empty() {
        return Err(TuskError::Ai(format!("{} API key required", provider_label)));
    }

    let url = format!("{}/models", base_url);
    let mut req = http_client().get(&url).bearer_auth(key);
    for (name, value) in extra_headers {
        req = req.header(*name, *value);
    }
    let resp = req
        .send()
        .await
        .map_err(|e| TuskError::Ai(format!("Cannot reach {}: {}", provider_label, e)))?;

    if !resp.status().is_success() {
        let status = resp.status();
        let body = resp.text().await.unwrap_or_default();
        return Err(TuskError::Ai(format!(
            "{} error ({}): {}",
            provider_label, status, body
        )));
    }

    let parsed: FireworksModelsResponse = resp.json().await.map_err(|e| {
        TuskError::Ai(format!("Failed to parse {} models list: {}", provider_label, e))
    })?;

    Ok(parsed
        .data
        .into_iter()
        .map(|m| OllamaModel { name: m.id })
        .collect())
}

async fn call_ai_with_retry<F, Fut, T>(
    _settings: &AiSettings,
    operation: &str,
    f: F,
) -> TuskResult<T>
where
    F: Fn() -> Fut,
    Fut: std::future::Future<Output = TuskResult<T>>,
{
    let mut last_error = None;

    for attempt in 0..MAX_RETRIES {
        match f().await {
            Ok(result) => return Ok(result),
            Err(e) => {
                last_error = Some(e);
                if attempt < MAX_RETRIES - 1 {
                    log::warn!(
                        "{} failed (attempt {}/{}), retrying in {}ms...",
                        operation,
                        attempt + 1,
                        MAX_RETRIES,
                        RETRY_DELAY_MS
                    );
                    tokio::time::sleep(Duration::from_millis(RETRY_DELAY_MS)).await;
                }
            }
        }
    }

    Err(last_error.unwrap_or_else(|| {
        TuskError::Ai(format!(
            "{} failed after {} attempts",
            operation, MAX_RETRIES
        ))
    }))
}

pub(crate) async fn load_ai_settings(app: &AppHandle, state: &AppState) -> TuskResult<AiSettings> {
    // Try in-memory cache first
    if let Some(cached) = state.ai_settings.read().await.clone() {
        return Ok(cached);
    }
    // Fallback to disk
    let path = get_ai_settings_path(app)?;
    if !path.exists() {
        return Err(TuskError::Ai(
            "No AI model selected. Open AI settings to choose a model.".to_string(),
        ));
    }
    let data = fs::read_to_string(&path)?;
    let settings: AiSettings = serde_json::from_str(&data)?;
    // Populate cache for future calls
    *state.ai_settings.write().await = Some(settings.clone());
    Ok(settings)
}

/// Provider-agnostic chat-completions dispatcher used by the chat agent.
/// Returns the model's raw text content.
pub(crate) async fn call_chat_messages(
    app: &AppHandle,
    state: &AppState,
    messages: Vec<OllamaChatMessage>,
    format: Option<String>,
) -> TuskResult<String> {
    let settings = load_ai_settings(app, state).await?;

    if settings.model.is_empty() {
        return Err(TuskError::Ai(
            "No AI model selected. Open AI settings to choose a model.".to_string(),
        ));
    }

    match settings.provider {
        AiProvider::Ollama => call_ollama(&settings, messages, format).await,
        AiProvider::Fireworks => {
            let api_key = require_api_key(
                settings.fireworks_api_key.as_deref(),
                "Fireworks API key not set. Open AI settings to add it.",
            )?;
            call_openai_compatible(
                &settings,
                FIREWORKS_BASE_URL,
                &api_key,
                "Fireworks",
                &[],
                messages,
                format,
            )
            .await
        }
        AiProvider::OpenRouter => {
            let api_key = require_api_key(
                settings.openrouter_api_key.as_deref(),
                "OpenRouter API key not set. Open AI settings to add it.",
            )?;
            call_openai_compatible(
                &settings,
                OPENROUTER_BASE_URL,
                &api_key,
                "OpenRouter",
                OPENROUTER_HEADERS,
                messages,
                format,
            )
            .await
        }
    }
}

/// Trim and validate an optional API key, returning a user-facing error when
/// it's missing or blank.
fn require_api_key(key: Option<&str>, missing_msg: &str) -> TuskResult<String> {
    key.map(|k| k.trim().to_string())
        .filter(|k| !k.is_empty())
        .ok_or_else(|| TuskError::Ai(missing_msg.to_string()))
}

async fn call_ollama(
    settings: &AiSettings,
    messages: Vec<OllamaChatMessage>,
    format: Option<String>,
) -> TuskResult<String> {
    let url = format!("{}/api/chat", settings.ollama_url.trim_end_matches('/'));
    let request = OllamaChatRequest {
        model: settings.model.clone(),
        messages,
        stream: false,
        format,
    };

    call_ai_with_retry(settings, "Ollama request", || {
        let url = url.clone();
        let request = request.clone();
        async move {
            let resp = http_client()
                .post(&url)
                .json(&request)
                .send()
                .await
                .map_err(|e| {
                    TuskError::Ai(format!("Cannot connect to Ollama at {}: {}", url, e))
                })?;

            if !resp.status().is_success() {
                let status = resp.status();
                let body = resp.text().await.unwrap_or_default();
                return Err(TuskError::Ai(format!(
                    "Ollama error ({}): {}",
                    status, body
                )));
            }

            let chat_resp: OllamaChatResponse = resp
                .json()
                .await
                .map_err(|e| TuskError::Ai(format!("Failed to parse Ollama response: {}", e)))?;

            Ok(chat_resp.message.content)
        }
    })
    .await
}

/// Chat-completions call for any OpenAI-compatible provider (Fireworks,
/// OpenRouter). `extra_headers` carries provider-specific attribution headers.
async fn call_openai_compatible(
    settings: &AiSettings,
    base_url: &str,
    api_key: &str,
    provider_label: &str,
    extra_headers: &[(&str, &str)],
    messages: Vec<OllamaChatMessage>,
    format: Option<String>,
) -> TuskResult<String> {
    let url = format!("{}/chat/completions", base_url);
    let response_format = format.as_deref().map(|f| FireworksResponseFormat {
        kind: if f == "json" {
            "json_object".to_string()
        } else {
            f.to_string()
        },
    });

    let request = FireworksChatRequest {
        model: settings.model.clone(),
        messages,
        temperature: 0.0,
        response_format,
    };

    let operation = format!("{} request", provider_label);
    call_ai_with_retry(settings, &operation, || {
        let url = url.clone();
        let request = request.clone();
        let api_key = api_key.to_string();
        async move {
            let mut req = http_client().post(&url).bearer_auth(&api_key);
            for (name, value) in extra_headers {
                req = req.header(*name, *value);
            }
            let resp = req
                .json(&request)
                .send()
                .await
                .map_err(|e| {
                    TuskError::Ai(format!("Cannot reach {} at {}: {}", provider_label, url, e))
                })?;

            if !resp.status().is_success() {
                let status = resp.status();
                let body = resp.text().await.unwrap_or_default();
                // Detect the prompt-overflow case and rewrite it as actionable
                // user-facing guidance instead of leaking raw API JSON.
                let lc = body.to_lowercase();
                if status.as_u16() == 400
                    && (lc.contains("prompt is too long")
                        || lc.contains("maximum context length"))
                {
                    return Err(TuskError::Ai(
                        "Context exceeded the model's window. Click `Compact` (or send `/compact`) \
                         to summarize the conversation, then retry. If this keeps happening, \
                         start a new chat for unrelated questions.".to_string(),
                    ));
                }
                return Err(TuskError::Ai(format!(
                    "{} error ({}): {}",
                    provider_label, status, body
                )));
            }

            let parsed: FireworksChatResponse = resp.json().await.map_err(|e| {
                TuskError::Ai(format!("Failed to parse {} response: {}", provider_label, e))
            })?;

            parsed
                .choices
                .into_iter()
                .next()
                .map(|c| c.message.content)
                .ok_or_else(|| {
                    TuskError::Ai(format!("{} returned no choices", provider_label))
                })
        }
    })
    .await
}

// ---------------------------------------------------------------------------
// Lite overview builder (chat v2)
// ---------------------------------------------------------------------------

const OVERVIEW_CACHE_TTL: std::time::Duration = std::time::Duration::from_secs(300);

/// Build a compact overview of a connection's databases and active-DB tables.
/// Designed to be small enough to inject in every system prompt.
/// Cached per connection_id; invalidated by switch_database / disconnect.
pub(crate) async fn build_overview_context(
    state: &AppState,
    connection_id: &str,
) -> TuskResult<String> {
    if let Some(cached) = state.overview_cache.read().await.get(connection_id).cloned() {
        if cached.cached_at.elapsed() < OVERVIEW_CACHE_TTL {
            return Ok(cached.value);
        }
    }

    let flavor = state.get_flavor(connection_id).await;
    let text = match flavor {
        crate::state::DbFlavor::ClickHouse => build_overview_clickhouse(state, connection_id).await?,
        _ => build_overview_postgres(state, connection_id).await?,
    };

    let mut cache = state.overview_cache.write().await;
    cache.insert(
        connection_id.to_string(),
        crate::state::CachedString {
            value: text.clone(),
            cached_at: std::time::Instant::now(),
        },
    );
    Ok(text)
}

async fn build_overview_postgres(state: &AppState, connection_id: &str) -> TuskResult<String> {
    let pool = state.get_pool(connection_id).await?;

    let (version_res, current_db_res, dbs_res, tables_res) = tokio::join!(
        sqlx::query_scalar::<_, String>("SELECT version()").fetch_one(&pool),
        sqlx::query_scalar::<_, String>("SELECT current_database()").fetch_one(&pool),
        sqlx::query_scalar::<_, String>(
            "SELECT datname FROM pg_database \
             WHERE datistemplate = false AND datallowconn = true \
             ORDER BY datname"
        )
        .fetch_all(&pool),
        sqlx::query_as::<_, (String, String)>(
            "SELECT table_schema, table_name FROM information_schema.tables \
             WHERE table_schema NOT IN ('pg_catalog','information_schema','pg_toast','gp_toolkit') \
               AND table_type = 'BASE TABLE' \
             ORDER BY table_schema, table_name"
        )
        .fetch_all(&pool),
    );

    let version = version_res.unwrap_or_default();
    let current_db = current_db_res.unwrap_or_default();
    let all_dbs = dbs_res.unwrap_or_default();
    let tables = tables_res.unwrap_or_default();

    let short_version = version
        .split_whitespace()
        .take(2)
        .collect::<Vec<_>>()
        .join(" ");

    let mut out = Vec::<String>::new();
    out.push(format!("DATABASE: {}", short_version));
    if !current_db.is_empty() {
        out.push(format!("ACTIVE DATABASE: {}", current_db));
    }
    out.push(String::new());

    if !all_dbs.is_empty() {
        out.push("DATABASES ON THIS SERVER:".to_string());
        for db in &all_dbs {
            if db == &current_db {
                out.push(format!("  * {} (active)", db));
            } else {
                out.push(format!("    {}", db));
            }
        }
        out.push(String::new());
    }

    // Cap the inline table list so a multi-thousand-table catalog (typical on
    // Greenplum data-lake clusters) doesn't blow past the model's context
    // window when the overview is re-injected into every turn's system prompt.
    // Beyond the cap, we summarize per-schema counts so the agent still knows
    // which schemas exist and can call list_tables for the missing ones.
    const OVERVIEW_TABLE_CAP: usize = 400;
    out.push(format!(
        "TABLES IN ACTIVE DATABASE ({}):",
        tables.len()
    ));
    if tables.len() <= OVERVIEW_TABLE_CAP {
        for (schema, name) in &tables {
            out.push(format!("  {}.{}", schema, name));
        }
    } else {
        for (schema, name) in tables.iter().take(OVERVIEW_TABLE_CAP) {
            out.push(format!("  {}.{}", schema, name));
        }
        out.push(format!(
            "  ... + {} more tables (truncated). Per-schema counts of remaining tables:",
            tables.len() - OVERVIEW_TABLE_CAP
        ));
        let mut remaining: BTreeMap<&str, usize> = BTreeMap::new();
        for (schema, _) in tables.iter().skip(OVERVIEW_TABLE_CAP) {
            *remaining.entry(schema.as_str()).or_default() += 1;
        }
        for (schema, count) in &remaining {
            out.push(format!("    {} ({} tables)", schema, count));
        }
        out.push(
            "  → call list_tables({\"schema\":\"<schema>\"}) to enumerate a specific schema."
                .to_string(),
        );
    }
    out.push(String::new());
    out.push(
        "NOTE: Tables of other databases are not enumerated here. \
         Call list_tables({\"database\":\"<name>\"}) to see them. For PostgreSQL this also \
         requires switch_database before you can actually query data there."
            .to_string(),
    );

    if matches!(state.get_flavor(connection_id).await, DbFlavor::Greenplum) {
        out.push(String::new());
        out.push(
            "GREENPLUM NOTES (this is an MPP cluster, not vanilla PostgreSQL):\n\
             - When you call get_columns on a table you'll see a `-- GP: DISTRIBUTED BY (...) | STORAGE: ...` line. \
               Use it for join planning.\n\
             - JOINs on the distribution key avoid Redistribute Motion (often 10-100x faster); \
               JOINs on non-distribution columns broadcast or redistribute the smaller side.\n\
             - AO/AOCO storage is optimized for large analytical scans; heap is for point lookups. \
               Don't run wide row-by-row predicates against AOCO without a covering filter.\n\
             - For data-skew diagnostics on a table T:\n\
               SELECT gp_segment_id, count(*) FROM T GROUP BY 1 ORDER BY 2 DESC LIMIT 10;\n\
             - Greenplum extends PostgreSQL syntax — most things work, but some features differ \
               by major version (GP6 ≈ PG9.4 catalog, GP7 ≈ PG14)."
                .to_string(),
        );
    }

    Ok(out.join("\n"))
}

async fn build_overview_clickhouse(state: &AppState, connection_id: &str) -> TuskResult<String> {
    let client = state.get_ch_client(connection_id).await?;
    let active_db = client.database.clone();

    let dbs_rows = client
        .fetch_objects(
            "SELECT name FROM system.databases \
             WHERE name NOT IN ('system','INFORMATION_SCHEMA','information_schema') \
             ORDER BY name",
        )
        .await
        .unwrap_or_default();
    let table_rows = client
        .fetch_objects(
            "SELECT database, name FROM system.tables \
             WHERE database NOT IN ('system','INFORMATION_SCHEMA','information_schema') \
             ORDER BY database, name",
        )
        .await
        .unwrap_or_default();
    let version = client.ping().await.unwrap_or_default();

    let mut out = Vec::<String>::new();
    out.push("DATABASE: ClickHouse".to_string());
    if !version.is_empty() {
        out.push(format!("VERSION: {}", version.trim()));
    }
    out.push(format!("ACTIVE DATABASE: {}", active_db));
    out.push(String::new());

    if !dbs_rows.is_empty() {
        out.push("DATABASES ON THIS SERVER:".to_string());
        for row in &dbs_rows {
            let name = row.get("name").and_then(|v| v.as_str()).unwrap_or("");
            if name == active_db {
                out.push(format!("  * {} (active)", name));
            } else {
                out.push(format!("    {}", name));
            }
        }
        out.push(String::new());
    }

    out.push(format!("TABLES ACROSS ALL DATABASES ({}):", table_rows.len()));
    for row in &table_rows {
        let dbn = row.get("database").and_then(|v| v.as_str()).unwrap_or("");
        let tbl = row.get("name").and_then(|v| v.as_str()).unwrap_or("");
        if !dbn.is_empty() && !tbl.is_empty() {
            out.push(format!("  {}.{}", dbn, tbl));
        }
    }
    out.push(String::new());
    out.push(
        "NOTE: ClickHouse allows fully-qualified `db.table` queries — \
         you can reference any table in this list directly without switching databases."
            .to_string(),
    );

    Ok(out.join("\n"))
}

// ---------------------------------------------------------------------------
// Schema query helpers
// ---------------------------------------------------------------------------

#[derive(Clone)]
pub(crate) struct ColumnInfo {
    pub(crate) schema: String,
    pub(crate) table: String,
    pub(crate) column: String,
    pub(crate) data_type: String,
    pub(crate) not_null: bool,
    pub(crate) is_pk: bool,
    pub(crate) column_default: Option<String>,
}

/// Render a single table's column block in the human/LLM-readable schema format.
/// Used by the chat agent's `get_columns` tool.
#[allow(clippy::too_many_arguments)]
pub(crate) fn format_table_block(
    full_name: &str,
    columns: &[ColumnInfo],
    tbl_comments: &HashMap<String, String>,
    col_comments: &HashMap<(String, String, String), String>,
    fk_inline: &HashMap<(String, String, String), String>,
    enum_map: &BTreeMap<String, Vec<String>>,
    unique_map: &HashMap<(String, String), Vec<String>>,
    varchar_values: &HashMap<(String, String, String), Vec<String>>,
    jsonb_keys: &HashMap<(String, String, String), Vec<String>>,
    gp_extras: Option<&HashMap<(String, String), GpTableExtras>>,
    output: &mut Vec<String>,
) {
    let tbl_comment = tbl_comments.get(full_name).map(|c| c.as_str());
    match tbl_comment {
        Some(comment) => output.push(format!("\nTABLE {} -- {}", full_name, comment)),
        None => output.push(format!("\nTABLE {}", full_name)),
    }

    // Greenplum-specific table attributes (distribution + storage), printed
    // immediately under the TABLE header so the agent sees them before any
    // column. Skipped on non-GP connections.
    if let Some(extras_map) = gp_extras {
        let parts: Vec<&str> = full_name.splitn(2, '.').collect();
        if parts.len() == 2 {
            if let Some(extras) =
                extras_map.get(&(parts[0].to_string(), parts[1].to_string()))
            {
                let mut bits: Vec<String> = Vec::new();
                if !extras.distribution.is_empty() {
                    bits.push(extras.distribution.clone());
                }
                if let Some(storage) = &extras.storage {
                    bits.push(format!("STORAGE: {}", storage));
                }
                if !bits.is_empty() {
                    output.push(format!("  -- GP: {}", bits.join(" | ")));
                }
            }
        }
    }

    for ci in columns {
        let mut parts: Vec<String> = vec![ci.column.clone(), ci.data_type.clone()];

        if ci.is_pk {
            parts.push("PK".to_string());
        }
        if ci.not_null && !ci.is_pk {
            parts.push("NOT NULL".to_string());
        }

        if let Some(ref_target) =
            fk_inline.get(&(ci.schema.clone(), ci.table.clone(), ci.column.clone()))
        {
            parts.push(format!("FK->{}", ref_target));
        }

        if let Some(ref def) = ci.column_default {
            let simplified = simplify_default(def);
            if !simplified.is_empty() {
                parts.push(format!("DEFAULT {}", simplified));
            }
        }

        let col_key = (ci.schema.clone(), ci.table.clone(), ci.column.clone());
        let col_comment = col_comments.get(&col_key);
        let enum_annotation = enum_map.get(&ci.data_type);

        let mut suffix_parts: Vec<String> = Vec::new();
        if let Some(vals) = enum_annotation {
            let vals_str = vals
                .iter()
                .map(|v| format!("'{}'", v))
                .collect::<Vec<_>>()
                .join(", ");
            suffix_parts.push(format!("enum: {}", vals_str));
        }

        if enum_annotation.is_none() {
            if let Some(vals) = varchar_values.get(&col_key) {
                let vals_str = vals
                    .iter()
                    .take(15)
                    .map(|v| format!("'{}'", v))
                    .collect::<Vec<_>>()
                    .join(", ");
                suffix_parts.push(format!("values: {}", vals_str));
            }
        }

        if let Some(keys) = jsonb_keys.get(&col_key) {
            suffix_parts.push(format!("json keys: {}", keys.join(", ")));
        }

        if let Some(comment) = col_comment {
            suffix_parts.push(comment.clone());
        }

        if suffix_parts.is_empty() {
            output.push(format!("  {}", parts.join(" ")));
        } else {
            output.push(format!(
                "  {} -- {}",
                parts.join(" "),
                suffix_parts.join("; ")
            ));
        }
    }

    // Unique constraints
    let schema_table: Vec<&str> = full_name.splitn(2, '.').collect();
    if schema_table.len() == 2 {
        if let Some(uqs) =
            unique_map.get(&(schema_table[0].to_string(), schema_table[1].to_string()))
        {
            for uq_cols in uqs {
                output.push(format!("  UNIQUE({})", uq_cols));
            }
        }
    }
}

pub(crate) async fn fetch_columns(pool: &sqlx::PgPool) -> TuskResult<Vec<ColumnInfo>> {
    let rows = sqlx::query(
        "SELECT \
            c.table_schema, c.table_name, c.column_name, \
            CASE \
                WHEN c.data_type = 'USER-DEFINED' THEN c.udt_name \
                WHEN c.data_type = 'ARRAY' THEN c.udt_name || '[]' \
                ELSE c.data_type \
            END AS data_type, \
            c.is_nullable = 'NO' AS not_null, \
            c.column_default, \
            EXISTS( \
                SELECT 1 FROM information_schema.table_constraints tc \
                JOIN information_schema.key_column_usage kcu \
                    ON tc.constraint_name = kcu.constraint_name AND tc.table_schema = kcu.table_schema \
                WHERE tc.constraint_type = 'PRIMARY KEY' \
                    AND tc.table_schema = c.table_schema \
                    AND tc.table_name = c.table_name \
                    AND kcu.column_name = c.column_name \
            ) AS is_pk \
         FROM information_schema.columns c \
         WHERE c.table_schema NOT IN ('pg_catalog', 'information_schema', 'pg_toast', 'gp_toolkit') \
         ORDER BY c.table_schema, c.table_name, c.ordinal_position",
    )
    .fetch_all(pool)
    .await
    .map_err(TuskError::Database)?;

    Ok(rows
        .iter()
        .map(|r| ColumnInfo {
            schema: r.get(0),
            table: r.get(1),
            column: r.get(2),
            data_type: r.get(3),
            not_null: r.get(4),
            column_default: r.get(5),
            is_pk: r.get(6),
        })
        .collect())
}

pub(crate) struct ForeignKeyInfo {
    pub(crate) schema: String,
    pub(crate) table: String,
    pub(crate) columns: Vec<String>,
    pub(crate) ref_schema: String,
    pub(crate) ref_table: String,
    pub(crate) ref_columns: Vec<String>,
}

pub(crate) async fn fetch_foreign_keys_raw(pool: &sqlx::PgPool) -> TuskResult<Vec<ForeignKeyInfo>> {
    let rows = sqlx::query(
        "SELECT \
           cn.nspname AS schema_name, cl.relname AS table_name, \
           array_agg(DISTINCT a.attname ORDER BY a.attname) AS columns, \
           cnf.nspname AS ref_schema, clf.relname AS ref_table, \
           array_agg(DISTINCT af.attname ORDER BY af.attname) AS ref_columns \
         FROM pg_constraint con \
         JOIN pg_class cl ON con.conrelid = cl.oid \
         JOIN pg_namespace cn ON cl.relnamespace = cn.oid \
         JOIN pg_class clf ON con.confrelid = clf.oid \
         JOIN pg_namespace cnf ON clf.relnamespace = cnf.oid \
         JOIN pg_attribute a ON a.attrelid = con.conrelid AND a.attnum = ANY(con.conkey) \
         JOIN pg_attribute af ON af.attrelid = con.confrelid AND af.attnum = ANY(con.confkey) \
         WHERE con.contype = 'f' \
           AND cn.nspname NOT IN ('pg_catalog','information_schema','pg_toast','gp_toolkit') \
         GROUP BY cn.nspname, cl.relname, cnf.nspname, clf.relname, con.oid",
    )
    .fetch_all(pool)
    .await
    .map_err(TuskError::Database)?;

    Ok(rows
        .iter()
        .map(|r| ForeignKeyInfo {
            schema: r.get(0),
            table: r.get(1),
            columns: r.get(2),
            ref_schema: r.get(3),
            ref_table: r.get(4),
            ref_columns: r.get(5),
        })
        .collect())
}

/// Returns BTreeMap<enum_type_name, Vec<enum_values>> ordered by type name
pub(crate) async fn fetch_enum_types(pool: &sqlx::PgPool) -> TuskResult<BTreeMap<String, Vec<String>>> {
    let rows = sqlx::query(
        "SELECT t.typname, \
                array_agg(e.enumlabel ORDER BY e.enumsortorder) AS vals \
         FROM pg_enum e \
         JOIN pg_type t ON e.enumtypid = t.oid \
         JOIN pg_namespace n ON t.typnamespace = n.oid \
         WHERE n.nspname NOT IN ('pg_catalog', 'information_schema') \
         GROUP BY t.typname \
         ORDER BY t.typname",
    )
    .fetch_all(pool)
    .await
    .map_err(TuskError::Database)?;

    let mut map = BTreeMap::new();
    for r in &rows {
        let name: String = r.get(0);
        let vals: Vec<String> = r.get(1);
        map.insert(name, vals);
    }
    Ok(map)
}

/// Returns HashMap<"schema.table", comment>
pub(crate) async fn fetch_table_comments(pool: &sqlx::PgPool) -> TuskResult<HashMap<String, String>> {
    let rows = sqlx::query(
        "SELECT n.nspname, c.relname, obj_description(c.oid, 'pg_class') \
         FROM pg_class c \
         JOIN pg_namespace n ON c.relnamespace = n.oid \
         WHERE c.relkind IN ('r', 'v', 'p', 'm') \
           AND n.nspname NOT IN ('pg_catalog', 'information_schema', 'pg_toast', 'gp_toolkit') \
           AND obj_description(c.oid, 'pg_class') IS NOT NULL",
    )
    .fetch_all(pool)
    .await
    .map_err(TuskError::Database)?;

    let mut map = HashMap::new();
    for r in &rows {
        let schema: String = r.get(0);
        let table: String = r.get(1);
        let comment: String = r.get(2);
        map.insert(format!("{}.{}", schema, table), comment);
    }
    Ok(map)
}

/// Returns HashMap<(schema, table, column), comment>
pub(crate) async fn fetch_column_comments(
    pool: &sqlx::PgPool,
) -> TuskResult<HashMap<(String, String, String), String>> {
    let rows = sqlx::query(
        "SELECT n.nspname, c.relname, a.attname, d.description \
         FROM pg_description d \
         JOIN pg_class c ON d.objoid = c.oid \
         JOIN pg_namespace n ON c.relnamespace = n.oid \
         JOIN pg_attribute a ON a.attrelid = c.oid AND a.attnum = d.objsubid \
         WHERE d.objsubid > 0 \
           AND n.nspname NOT IN ('pg_catalog', 'information_schema', 'pg_toast', 'gp_toolkit')",
    )
    .fetch_all(pool)
    .await
    .map_err(TuskError::Database)?;

    let mut map = HashMap::new();
    for r in &rows {
        let schema: String = r.get(0);
        let table: String = r.get(1);
        let column: String = r.get(2);
        let comment: String = r.get(3);
        map.insert((schema, table, column), comment);
    }
    Ok(map)
}

/// Returns Vec<(schema, table, Vec<column_names>)> for UNIQUE constraints
pub(crate) async fn fetch_unique_constraints(
    pool: &sqlx::PgPool,
) -> TuskResult<Vec<(String, String, Vec<String>)>> {
    let rows = sqlx::query(
        "SELECT n.nspname, cl.relname, \
                array_agg(a.attname ORDER BY a.attnum) AS cols \
         FROM pg_constraint con \
         JOIN pg_class cl ON con.conrelid = cl.oid \
         JOIN pg_namespace n ON cl.relnamespace = n.oid \
         JOIN pg_attribute a ON a.attrelid = con.conrelid AND a.attnum = ANY(con.conkey) \
         WHERE con.contype = 'u' \
           AND n.nspname NOT IN ('pg_catalog', 'information_schema', 'pg_toast', 'gp_toolkit') \
         GROUP BY n.nspname, cl.relname, con.oid \
         ORDER BY n.nspname, cl.relname",
    )
    .fetch_all(pool)
    .await
    .map_err(TuskError::Database)?;

    Ok(rows
        .iter()
        .map(|r| {
            let schema: String = r.get(0);
            let table: String = r.get(1);
            let cols: Vec<String> = r.get(2);
            (schema, table, cols)
        })
        .collect())
}

/// Returns HashMap<(schema, table, column), Vec<distinct_values>> for varchar columns
/// with few distinct values (pseudo-enums), using pg_stats for zero-cost discovery.
/// Returns None if pg_stats is not accessible (graceful degradation).
#[allow(dead_code)] // re-exposed by profile_table tool (PR2)
async fn fetch_varchar_values(
    pool: &sqlx::PgPool,
) -> Option<HashMap<(String, String, String), Vec<String>>> {
    let rows = match sqlx::query(
        "SELECT s.schemaname, s.tablename, s.attname, \
                s.most_common_vals::text AS vals \
         FROM pg_stats s \
         JOIN information_schema.columns c \
             ON c.table_schema = s.schemaname \
             AND c.table_name = s.tablename \
             AND c.column_name = s.attname \
         WHERE s.schemaname NOT IN ('pg_catalog', 'information_schema', 'pg_toast', 'gp_toolkit') \
           AND c.data_type = 'character varying' \
           AND s.n_distinct > 0 AND s.n_distinct <= 20 \
           AND s.most_common_vals IS NOT NULL \
         ORDER BY s.schemaname, s.tablename, s.attname",
    )
    .fetch_all(pool)
    .await
    {
        Ok(r) => r,
        Err(e) => {
            log::warn!("Failed to fetch varchar values from pg_stats: {}", e);
            return None;
        }
    };

    let mut map = HashMap::new();
    for r in &rows {
        let schema: String = r.get(0);
        let table: String = r.get(1);
        let column: String = r.get(2);
        let vals_text: String = r.get(3);
        let vals = parse_pg_array_text(&vals_text);
        if !vals.is_empty() {
            map.insert((schema, table, column), vals);
        }
    }
    Some(map)
}

// ---------------------------------------------------------------------------
// Greenplum-specific table attributes
// ---------------------------------------------------------------------------

/// Per-table Greenplum metadata that aggressively shapes query performance:
/// the data-distribution policy and the storage type. Without these visible
/// to the chat agent it tends to write PG-optimal SQL that triggers
/// Redistribute Motion in GP and runs orders of magnitude slower than needed.
#[derive(Debug, Clone, Default)]
pub(crate) struct GpTableExtras {
    /// Pre-formatted policy string: e.g. "DISTRIBUTED BY (id, created_at)" or
    /// "DISTRIBUTED RANDOMLY". Empty string means "policy unknown".
    pub distribution: String,
    /// Storage kind: "heap" / "AO row" / "AO column" / "external" / etc.
    /// `None` when storage couldn't be determined for this catalog version.
    pub storage: Option<String>,
}

/// Fetch DISTRIBUTED BY policy + storage kind for every user table in the
/// current GP database. Two parallel queries; storage uses the GP-major-aware
/// catalog (GP6 has `pg_class.relstorage`, GP7 dropped it and uses access
/// methods via `pg_am`). Returns an empty map on any catalog miss so callers
/// degrade gracefully on unfamiliar GP releases.
pub(crate) async fn fetch_gp_table_extras(
    pool: &sqlx::PgPool,
    gp_major: u8,
) -> HashMap<(String, String), GpTableExtras> {
    let dist_query = "
        SELECT n.nspname, c.relname,
               COALESCE(
                   (SELECT array_agg(a.attname ORDER BY ord.idx)
                    FROM regexp_split_to_table(NULLIF(trim(p.distkey::text), ''), ' ')
                         WITH ORDINALITY AS ord(attnum_str, idx)
                    JOIN pg_attribute a
                      ON a.attrelid = c.oid
                     AND a.attnum::int = ord.attnum_str::int),
                   ARRAY[]::text[]
               ) AS dist_columns
        FROM gp_distribution_policy p
        JOIN pg_class c ON p.localoid = c.oid
        JOIN pg_namespace n ON c.relnamespace = n.oid
        WHERE n.nspname NOT IN ('pg_catalog','information_schema','pg_toast','gp_toolkit')
    ";

    let storage_query = if gp_major <= 6 {
        // GP6 (PG9.4 base) — reltype is encoded in pg_class.relstorage.
        "SELECT n.nspname, c.relname,
                CASE c.relstorage
                    WHEN 'h' THEN 'heap'
                    WHEN 'a' THEN 'AO row'
                    WHEN 'c' THEN 'AO column'
                    WHEN 'x' THEN 'external'
                    ELSE c.relstorage::text
                END AS storage
         FROM pg_class c
         JOIN pg_namespace n ON c.relnamespace = n.oid
         WHERE c.relkind IN ('r','m')
           AND n.nspname NOT IN ('pg_catalog','information_schema','pg_toast','gp_toolkit')"
    } else {
        // GP7 (PG14 base) — relstorage is gone; access method on pg_class.relam
        // identifies AO row / AO column tables (extension AM names are
        // 'ao_row' and 'ao_column').
        "SELECT n.nspname, c.relname,
                CASE
                    WHEN am.amname = 'heap' THEN 'heap'
                    WHEN am.amname = 'ao_row' THEN 'AO row'
                    WHEN am.amname = 'ao_column' THEN 'AO column'
                    ELSE COALESCE(am.amname, 'unknown')
                END AS storage
         FROM pg_class c
         JOIN pg_namespace n ON c.relnamespace = n.oid
         LEFT JOIN pg_am am ON c.relam = am.oid
         WHERE c.relkind IN ('r','m')
           AND n.nspname NOT IN ('pg_catalog','information_schema','pg_toast','gp_toolkit')"
    };

    let (dist_res, storage_res) = tokio::join!(
        sqlx::query(dist_query).fetch_all(pool),
        sqlx::query(storage_query).fetch_all(pool),
    );

    let mut map: HashMap<(String, String), GpTableExtras> = HashMap::new();

    if let Ok(rows) = dist_res {
        for r in &rows {
            let schema: String = r.get(0);
            let table: String = r.get(1);
            let cols: Vec<String> = r.get(2);
            let distribution = if cols.is_empty() {
                "DISTRIBUTED RANDOMLY".to_string()
            } else {
                format!("DISTRIBUTED BY ({})", cols.join(", "))
            };
            map.entry((schema, table)).or_default().distribution = distribution;
        }
    } else if let Err(e) = dist_res {
        log::warn!("gp_distribution_policy fetch failed: {}", e);
    }

    if let Ok(rows) = storage_res {
        for r in &rows {
            let schema: String = r.get(0);
            let table: String = r.get(1);
            let storage: String = r.get(2);
            map.entry((schema, table)).or_default().storage = Some(storage);
        }
    } else if let Err(e) = storage_res {
        log::warn!("Greenplum storage fetch failed: {}", e);
    }

    map
}

// ---------------------------------------------------------------------------
// Helpers
// ---------------------------------------------------------------------------

/// Parses PostgreSQL text representation of arrays: {val1,val2,"val with comma"}
#[allow(dead_code)] // helper for fetch_varchar_values; re-exposed by profile_table tool (PR2)
fn parse_pg_array_text(s: &str) -> Vec<String> {
    let s = s.trim();
    let s = s.strip_prefix('{').unwrap_or(s);
    let s = s.strip_suffix('}').unwrap_or(s);
    if s.is_empty() {
        return Vec::new();
    }

    let mut result = Vec::new();
    let mut current = String::new();
    let mut in_quotes = false;
    let mut chars = s.chars().peekable();

    while let Some(ch) = chars.next() {
        match ch {
            '"' if !in_quotes => {
                in_quotes = true;
            }
            '"' if in_quotes => {
                if chars.peek() == Some(&'"') {
                    current.push('"');
                    chars.next();
                } else {
                    in_quotes = false;
                }
            }
            ',' if !in_quotes => {
                result.push(current.trim().to_string());
                current = String::new();
            }
            _ => {
                current.push(ch);
            }
        }
    }
    if !current.is_empty() || !result.is_empty() {
        result.push(current.trim().to_string());
    }
    result
}

fn simplify_default(raw: &str) -> String {
    let s = raw.trim();
    if s.contains("nextval(") {
        return "auto-increment".to_string();
    }
    // Shorten common defaults
    if s == "now()" || s == "CURRENT_TIMESTAMP" || s == "current_timestamp" {
        return "now()".to_string();
    }
    if s == "true" || s == "false" {
        return s.to_string();
    }
    // Numeric/string literals — keep short ones, skip very long generated defaults
    if s.len() > 50 {
        return String::new();
    }
    s.to_string()
}

#[cfg(test)]
mod tests {
    use super::*;

    // ── Fireworks provider ───────────────────────────────────

    #[test]
    fn serializes_fireworks_provider() {
        let json = serde_json::to_string(&AiProvider::Fireworks).unwrap();
        assert_eq!(json, "\"fireworks\"");
    }

    #[test]
    fn deserializes_legacy_settings_with_dropped_provider_keys() {
        // Old config files may include `openai_api_key`/`anthropic_api_key` and a
        // legacy `"provider": "openai"` value — both must be tolerated, with the
        // unknown provider coerced to Ollama.
        let legacy = r#"{
            "provider": "openai",
            "ollama_url": "http://localhost:11434",
            "openai_api_key": "sk-deprecated",
            "anthropic_api_key": null,
            "model": "qwen2.5-coder:7b"
        }"#;
        let parsed: AiSettings = serde_json::from_str(legacy).unwrap();
        assert_eq!(parsed.provider, AiProvider::Ollama);
        assert_eq!(parsed.ollama_url, "http://localhost:11434");
        assert!(parsed.fireworks_api_key.is_none());
        assert_eq!(parsed.model, "qwen2.5-coder:7b");
    }

    #[test]
    fn parses_fireworks_chat_response() {
        let body = r#"{
            "choices": [
                {"message": {"role": "assistant", "content": "hi"}}
            ]
        }"#;
        let parsed: FireworksChatResponse = serde_json::from_str(body).unwrap();
        assert_eq!(parsed.choices.len(), 1);
        assert_eq!(parsed.choices[0].message.role, "assistant");
        assert_eq!(parsed.choices[0].message.content, "hi");
    }

    // ── OpenRouter provider ──────────────────────────────────

    #[test]
    fn serializes_openrouter_provider() {
        let json = serde_json::to_string(&AiProvider::OpenRouter).unwrap();
        assert_eq!(json, "\"openrouter\"");
    }

    #[test]
    fn deserializes_openrouter_settings() {
        let cfg = r#"{
            "provider": "openrouter",
            "ollama_url": "http://localhost:11434",
            "openrouter_api_key": "sk-or-v1-abc",
            "model": "anthropic/claude-3.5-sonnet"
        }"#;
        let parsed: AiSettings = serde_json::from_str(cfg).unwrap();
        assert_eq!(parsed.provider, AiProvider::OpenRouter);
        assert_eq!(parsed.openrouter_api_key.as_deref(), Some("sk-or-v1-abc"));
        assert_eq!(parsed.model, "anthropic/claude-3.5-sonnet");
    }

    #[test]
    fn parses_fireworks_models_list() {
        let body = r#"{
            "data": [
                {"id": "accounts/fireworks/models/qwen2p5-coder-32b-instruct"},
                {"id": "accounts/fireworks/models/deepseek-v3"}
            ]
        }"#;
        let parsed: FireworksModelsResponse = serde_json::from_str(body).unwrap();
        let names: Vec<String> = parsed.data.into_iter().map(|m| m.id).collect();
        assert_eq!(names.len(), 2);
        assert_eq!(
            names[0],
            "accounts/fireworks/models/qwen2p5-coder-32b-instruct"
        );
        assert_eq!(names[1], "accounts/fireworks/models/deepseek-v3");
    }
}
