use std::collections::BTreeMap;
use std::fs;
use std::path::PathBuf;

use anyhow::{Context, Result, bail};
use serde::Serialize;
use serde_json::Value;

use super::common::{CATALOG_CHAT_SCRIPT, CATALOG_MODELS_SUFFIX, run_catalog_npm_script, ts_catalog_to_json};

#[derive(Serialize)]
pub struct CatalogIndexEntry {
    #[serde(rename = "providerId")]
    pub provider_id: String,
    #[serde(rename = "rustMod")]
    pub rust_mod: String,
    pub count: usize,
}

pub struct ChatOptions {
    pub catalog_dir: PathBuf,
    pub skip_scripts: bool,
    pub models_dir: PathBuf,
    pub catalog_rs: PathBuf,
    pub no_regenerate_catalog: bool,
}

pub fn generate_chat(options: ChatOptions) -> Result<()> {
    if !options.catalog_dir.join(CATALOG_CHAT_SCRIPT).is_file() {
        bail!(
            "catalog source package not found at {} (expected npm catalog scripts)",
            options.catalog_dir.display()
        );
    }

    if !options.skip_scripts {
        run_catalog_npm_script(&options.catalog_dir, "generate-models")?;
    }

    let providers_dir = options.catalog_dir.join("src/providers");
    if !providers_dir.is_dir() {
        bail!(
            "missing catalog source providers directory at {}",
            providers_dir.display()
        );
    }

    fs::create_dir_all(&options.models_dir).context("create models output directory")?;

    let mut catalogs: BTreeMap<String, (String, Value)> = BTreeMap::new();
    for entry in fs::read_dir(&providers_dir).context("read catalog providers directory")? {
        let entry = entry?;
        let path = entry.path();
        let Some(file_name) = path.file_name().and_then(|n| n.to_str()) else {
            continue;
        };
        let Some(provider_id) = file_name.strip_suffix(CATALOG_MODELS_SUFFIX) else {
            continue;
        };
        let rust_mod = provider_id.replace('-', "_");
        let ts = fs::read_to_string(&path).with_context(|| format!("read {}", path.display()))?;
        let json = ts_catalog_to_json(&ts).with_context(|| format!("convert {}", path.display()))?;
        let count = json.as_object().map(|m| m.len()).unwrap_or(0);
        if count == 0 {
            continue;
        }
        catalogs.insert(provider_id.to_string(), (rust_mod, json));
        println!("Converted {provider_id}: {count} models");
    }

    if catalogs.is_empty() {
        bail!("no model catalog files found under {}", providers_dir.display());
    }

    let mut index = Vec::new();
    for (provider_id, (rust_mod, json)) in &catalogs {
        let out_path = options.models_dir.join(format!("{rust_mod}.json"));
        let pretty = serde_json::to_string_pretty(json).context("serialize catalog json")?;
        fs::write(&out_path, format!("{pretty}\n")).with_context(|| format!("write {}", out_path.display()))?;
        index.push(CatalogIndexEntry {
            provider_id: provider_id.clone(),
            rust_mod: rust_mod.clone(),
            count: json.as_object().map(|m| m.len()).unwrap_or(0),
        });
    }

    let index_path = options.models_dir.join("index.json");
    let index_json = serde_json::to_string_pretty(&index).context("serialize index.json")?;
    fs::write(&index_path, format!("{index_json}\n")).context("write index.json")?;

    if options.no_regenerate_catalog {
        println!(
            "\nWrote {} chat catalogs to {} (skipped catalog.rs regeneration)",
            catalogs.len(),
            options.models_dir.display()
        );
    } else {
        let catalog_source = render_chat_catalog_rs(&index);
        fs::write(&options.catalog_rs, catalog_source).context("write src/models/catalog.rs")?;
        println!(
            "\nWrote {} chat catalogs to {} and regenerated {}",
            catalogs.len(),
            options.models_dir.display(),
            options.catalog_rs.display()
        );
    }

    Ok(())
}

fn render_chat_catalog_rs(index: &[CatalogIndexEntry]) -> String {
    let mut out = String::from(
        "//! Embedded builtin model catalogs (auto-generated by `generate-models chat` — do not edit).\n\n\
         use std::collections::HashMap;\n\n\
         use once_cell::sync::Lazy;\n\
         use serde::Deserialize;\n\n\
         use crate::types::{\n\
             AnthropicMessagesCompat, Model, ModelCost, OpenAICompletionsCompat, OpenAIResponsesCompat, ThinkingLevelMap,\n\
         };\n\n\
         #[derive(Debug, Deserialize)]\n\
         struct RawModel {\n\
             id: String,\n\
             name: String,\n\
             api: String,\n\
             provider: String,\n\
             #[serde(rename = \"baseUrl\")]\n\
             base_url: String,\n\
             reasoning: bool,\n\
             #[serde(default)]\n\
             #[serde(rename = \"thinkingLevelMap\")]\n\
             thinking_level_map: Option<ThinkingLevelMap>,\n\
             input: Vec<String>,\n\
             cost: RawCost,\n\
             #[serde(rename = \"contextWindow\")]\n\
             context_window: u32,\n\
             #[serde(rename = \"maxTokens\")]\n\
             max_tokens: u32,\n\
             #[serde(default)]\n\
             headers: Option<HashMap<String, String>>,\n\
             #[serde(default)]\n\
             compat: Option<serde_json::Value>,\n\
         }\n\n\
         #[derive(Debug, Deserialize)]\n\
         struct RawCost {\n\
             input: f64,\n\
             output: f64,\n\
             #[serde(rename = \"cacheRead\")]\n\
             cache_read: f64,\n\
             #[serde(rename = \"cacheWrite\")]\n\
             cache_write: f64,\n\
         }\n\n\
         fn parse_models(json: &str) -> Vec<Model> {\n\
             let raw: HashMap<String, RawModel> = serde_json::from_str(json).expect(\"invalid embedded model catalog\");\n\
             raw.into_values().map(convert_model).collect()\n\
         }\n\n\
         fn convert_model(raw: RawModel) -> Model {\n\
             let (openai_completions_compat, openai_responses_compat, anthropic_compat) =\n\
                 parse_compat(&raw.api, raw.compat.as_ref());\n\n\
             Model {\n\
                 id: raw.id,\n\
                 name: raw.name,\n\
                 api: raw.api,\n\
                 provider: raw.provider,\n\
                 base_url: raw.base_url,\n\
                 reasoning: raw.reasoning,\n\
                 thinking_level_map: raw.thinking_level_map,\n\
                 input: raw.input,\n\
                 cost: ModelCost {\n\
                     input: raw.cost.input,\n\
                     output: raw.cost.output,\n\
                     cache_read: raw.cost.cache_read,\n\
                     cache_write: raw.cost.cache_write,\n\
                 },\n\
                 context_window: raw.context_window,\n\
                 max_tokens: raw.max_tokens,\n\
                 headers: raw.headers,\n\
                 openai_completions_compat,\n\
                 openai_responses_compat,\n\
                 anthropic_compat,\n\
             }\n\
         }\n\n\
         fn parse_compat(\n\
             api: &str,\n\
             compat: Option<&serde_json::Value>,\n\
         ) -> (\n\
             Option<OpenAICompletionsCompat>,\n\
             Option<OpenAIResponsesCompat>,\n\
             Option<AnthropicMessagesCompat>,\n\
         ) {\n\
             let Some(compat) = compat else {\n\
                 return (None, None, None);\n\
             };\n\
             match api {\n\
                 \"openai-completions\" => (serde_json::from_value(compat.clone()).ok(), None, None),\n\
                 \"openai-responses\" | \"azure-openai-responses\" | \"openai-codex-responses\" => {\n\
                     (None, serde_json::from_value(compat.clone()).ok(), None)\n\
                 }\n\
                 \"anthropic-messages\" => (None, None, serde_json::from_value(compat.clone()).ok()),\n\
                 _ => (None, None, None),\n\
             }\n\
         }\n\n\
         macro_rules! define_catalog {\n\
             ($name:ident, $file:literal) => {\n\
                 pub static $name: Lazy<Vec<Model>> = Lazy::new(|| {\n\
                     parse_models(include_str!(concat!(\n\
                         env!(\"CARGO_MANIFEST_DIR\"),\n\
                         \"/models/\",\n\
                         $file\n\
                     )))\n\
                 });\n\
             };\n\
         }\n\n",
    );

    for entry in index {
        let const_name = chat_catalog_const_name(&entry.provider_id);
        out.push_str(&format!(
            "define_catalog!({const_name}, \"{}.json\");\n",
            entry.rust_mod
        ));
    }

    out.push_str("\npub fn all_builtin_models() -> HashMap<&'static str, &'static [Model]> {\n    HashMap::from([\n");
    for entry in index {
        let const_name = chat_catalog_const_name(&entry.provider_id);
        out.push_str(&format!(
            "        (\"{}\", {}.as_slice()),\n",
            entry.provider_id, const_name
        ));
    }
    out.push_str(
        "    ])\n}\n\n\
         pub fn get_builtin_model(provider: &str, id: &str) -> Option<Model> {\n\
             all_builtin_models().get(provider)?.iter().find(|m| m.id == id).cloned()\n\
         }\n\n\
         pub fn get_builtin_models(provider: &str) -> Vec<Model> {\n\
             all_builtin_models()\n\
                 .get(provider)\n\
                 .map(|models| models.to_vec())\n\
                 .unwrap_or_default()\n\
         }\n\n\
         pub fn get_builtin_providers() -> Vec<&'static str> {\n\
             let mut providers: Vec<_> = all_builtin_models().keys().copied().collect();\n\
             providers.sort_unstable();\n\
             providers\n\
         }\n",
    );
    out
}

fn chat_catalog_const_name(provider_id: &str) -> String {
    let mut out = String::new();
    for ch in provider_id.chars() {
        if ch.is_ascii_alphanumeric() {
            out.push(ch.to_ascii_uppercase());
        } else {
            out.push('_');
        }
    }
    format!("{out}_MODELS")
}
