//! KHTST_RUST — refatoração Rust/PyO3 do KHTST com NCM (Next-Concept
//! Modeling), provas integradas (memoria-ncm-khtst n-02/n-03), atenção ADBA
//! (r-04) e ponte duplex (r-05). Documentação: docs/memorias/.
//!
//! Convenções do projeto (n-00 §4): nascimento neutro, leitura honesta,
//! Teos numerados com ∎, token EXCLUSIVAMENTE em HF_TOKEN apagado após uso.

use pyo3::prelude::*;

pub mod amostrador;
pub mod atencao;
pub mod camadas;
pub mod config;
pub mod controlador;
pub mod diagnostico;
pub mod gerador;
pub mod mascaras;
pub mod memoria;
pub mod perdas;
pub mod ponte;
pub mod quantizacao;
pub mod quantizador;
pub mod ssd;
pub mod stiefel;
pub mod telemetria;
pub mod tokens;
pub mod tokenizador;
pub mod treino;
pub mod tronco;

pub mod build_info {
    include!(concat!(env!("OUT_DIR"), "/built.rs"));
}

#[pyfunction]
fn versao() -> &'static str {
    "0.1.11 (KHTST_RUST — NCM + provas integradas + ADBA + ponte duplex)"
}

#[pyfunction]
fn info_build() -> String {
    format!(
        "rustc={}; perfil={}; features={}; alvo={}",
        build_info::RUSTC_VERSION,
        build_info::PROFILE,
        build_info::FEATURES_STR,
        build_info::TARGET,
    )
}

/// Lint de config com sugestões por Levenshtein (diagnóstico de digitação).
#[pyfunction]
fn lint_config(py: Python, caminho: &str) -> PyResult<PyObject> {
    let txt = std::fs::read_to_string(caminho)
        .map_err(|e| pyo3::exceptions::PyIOError::new_err(format!("{caminho}: {e}")))?;
    let v: serde_json::Value = serde_json::from_str(&txt)
        .map_err(|e| pyo3::exceptions::PyValueError::new_err(format!("JSON inválido: {e}")))?;
    let achatado = diagnostico::achatar(&v);
    let avisos = diagnostico::lint_chaves(&achatado);
    let saida: Vec<serde_json::Value> = avisos
        .iter()
        .map(|a| serde_json::json!({
            "chave": a.chave, "sugestao": a.sugestao, "causa": a.causa, "correcao": a.correcao,
        }))
        .collect();
    Ok(pythonize::pythonize(py, &serde_json::json!(saida))?.into())
}

#[pyclass(name = "Bpe")]
struct BpePy {
    bpe: tokens::Bpe,
}

#[pymethods]
impl BpePy {
    #[staticmethod]
    fn treinar(textos: Vec<String>, vocab_alvo: usize) -> PyResult<Self> {
        Ok(BpePy { bpe: tokens::Bpe::treinar(&textos, vocab_alvo) })
    }

    #[staticmethod]
    fn carregar(caminho: &str) -> PyResult<Self> {
        let txt = std::fs::read_to_string(caminho)
            .map_err(|e| pyo3::exceptions::PyIOError::new_err(format!("{e}")))?;
        let v: serde_json::Value = serde_json::from_str(&txt)
            .map_err(|e| pyo3::exceptions::PyValueError::new_err(format!("JSON inválido: {e}")))?;
        let bpe = tokens::Bpe::de_json(&v)
            .ok_or_else(|| pyo3::exceptions::PyValueError::new_err("formato BPE inválido"))?;
        Ok(BpePy { bpe })
    }

    fn codificar(&mut self, texto: &str) -> Vec<u32> {
        self.bpe.codificar(texto)
    }

    fn decodificar(&self, ids: Vec<u32>) -> String {
        self.bpe.decodificar(&ids)
    }

    fn salvar(&self, caminho: &str) -> PyResult<()> {
        std::fs::write(caminho, serde_json::to_string(&self.bpe.para_json()).unwrap())
            .map_err(|e| pyo3::exceptions::PyIOError::new_err(format!("{e}")))?;
        Ok(())
    }

    #[getter]
    fn n_vocab(&self) -> usize {
        self.bpe.n_vocab()
    }
}

#[pyclass(name = "Config")]
struct ConfigPy {
    cfg: config::ConfigKhtst,
    caminho: Option<String>,
}

#[pymethods]
impl ConfigPy {
    #[new]
    #[allow(never_type_fallback)]
    #[pyo3(signature = (caminho=None))]
    fn new(caminho: Option<String>) -> PyResult<Self> {
        let cfg = match &caminho {
            Some(c) => {
                let txt = std::fs::read_to_string(c)
                    .map_err(|e| pyo3::exceptions::PyIOError::new_err(format!("{c}: {e}")))?;
                let v: serde_json::Value = serde_json::from_str(&txt)
                    .map_err(|e| pyo3::exceptions::PyValueError::new_err(format!("JSON inválido: {e}")))?;
                serde_json::from_value(v)
                    .map_err(|e| pyo3::exceptions::PyValueError::new_err(format!("config: {e}")))?
            }
            None => config::ConfigKhtst::default(),
        };
        cfg.validar().map_err(pyo3::exceptions::PyValueError::new_err)?;
        Ok(ConfigPy { cfg, caminho })
    }

    fn validar(&self) -> PyResult<()> {
        self.cfg.validar().map_err(pyo3::exceptions::PyValueError::new_err)
    }

    fn checar_orcamento_cache(&self) -> PyResult<usize> {
        self.cfg.checar_orcamento_cache().map_err(pyo3::exceptions::PyValueError::new_err)
    }

    fn resumo(&self) -> String {
        format!(
            "modelo(vocab={}, d={}, camadas={}, ctx={}) ncm(L={}, niveis={:?}, K={}, d_lat={}, modo={}, ativo={})",
            self.cfg.modelo.vocab, self.cfg.modelo.d_modelo, self.cfg.modelo.n_camadas,
            self.cfg.modelo.comprimento_ctx, self.cfg.l_niveis(), self.cfg.ncm.niveis,
            self.cfg.ncm.k_por_nivel, self.cfg.ncm.d_lat, self.cfg.ncm.modo_quantizacao, self.cfg.ncm.ativo,
        )
    }

    #[getter]
    fn l_niveis(&self) -> usize {
        self.cfg.l_niveis()
    }
}

#[pyclass(name = "PonteDupla")]
struct PonteDuplaPy {
    ponte: ponte::PonteDupla,
}

#[pymethods]
impl PonteDuplaPy {
    #[new]
    #[pyo3(signature = (capacidade=64, timeout_ms=2000))]
    fn new(capacidade: usize, timeout_ms: u64) -> Self {
        PonteDuplaPy { ponte: ponte::PonteDupla::nova(capacidade, timeout_ms) }
    }

    fn rodada(&self, carga_a: Vec<f32>, carga_b: Vec<f32>, nivel: usize) -> PyResult<(Vec<f32>, Vec<f32>)> {
        let (ma, mb) = self.ponte.rodada(carga_a, carga_b, nivel)
            .map_err(pyo3::exceptions::PyRuntimeError::new_err)?;
        Ok((ma.carga, mb.carga))
    }

    fn demo_professor_aluno(&self, rodadas: usize, dim: usize) -> PyResult<PyObject> {
        let inicio: Vec<f32> = (0..dim).map(|i| (i as f32) * 0.01).collect();
        let stats = ponte::demo_professor_aluno(&self.ponte, rodadas, inicio, &|a, _| {
            a.iter().map(|x| x * 1.05).collect()
        })
        .map_err(pyo3::exceptions::PyRuntimeError::new_err)?;
        Python::with_gil(|py| {
            Ok(pythonize::pythonize(py, &serde_json::to_value(&stats).unwrap())?.into())
        })
    }

    fn estatisticas(&self) -> PyResult<PyObject> {
        let s = self.ponte.estatisticas();
        Python::with_gil(|py| {
            Ok(pythonize::pythonize(py, &serde_json::to_value(&s).unwrap())?.into())
        })
    }
}

#[cfg(feature = "ponte_async")]
#[pyclass(name = "PonteAsync")]
struct PonteAsyncPy {
    interna: std::sync::Arc<ponte::assincrona::PonteAsync>,
}

#[cfg(feature = "ponte_async")]
#[pymethods]
impl PonteAsyncPy {
    #[new]
    #[pyo3(signature = (capacidade=64, timeout_ms=2000))]
    fn new(capacidade: usize, timeout_ms: u64) -> Self {
        PonteAsyncPy { interna: std::sync::Arc::new(ponte::assincrona::PonteAsync::nova(capacidade, timeout_ms)) }
    }

    /// Corrotina: `await ponte.trocar(a, b, nivel)` — duplex simultânea.
    fn trocar<'p>(&self, py: Python<'p>, carga_a: Vec<f32>, carga_b: Vec<f32>, nivel: usize) -> PyResult<Bound<'p, PyAny>> {
        let p = self.interna.clone();
        pyo3_async_runtimes::tokio::future_into_py(py, async move {
            let (ma, mb) = p.rodada_async(carga_a, carga_b, nivel)
                .await
                .map_err(pyo3::exceptions::PyRuntimeError::new_err)?;
            Ok::<_, PyErr>((ma.carga, mb.carga))
        })
    }
}

/// Treina a fase A (tokenizador). Devolve dict com resumo + telemetria.
#[pyfunction]
#[pyo3(signature = (cfg, sequencias, caminho_telemetria=None, caminho_checkpoint=None, callback=None, teto_seg=None))]
fn treinar_fase_a(
    cfg: &ConfigPy,
    sequencias: Vec<Vec<u32>>,
    caminho_telemetria: Option<String>,
    caminho_checkpoint: Option<String>,
    callback: Option<Bound<'_, PyAny>>,
    teto_seg: Option<f64>,
) -> PyResult<PyObject> {
    let cb: treino::CallbackPy = callback.map(|c| c.unbind());
    let mut hub = telemetria::TelemetryHub::novo(caminho_telemetria);
    let dados = treino::DadosTreino::novo(sequencias, cfg.cfg.modelo.comprimento_ctx);
    let mut resumo = treino::treinar_fase_a(&cfg.cfg, &dados, &mut hub, &cb, caminho_checkpoint.as_deref(), teto_seg)
        .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("fase A: {e}")))?;
    let _ = &mut resumo;
    Python::with_gil(|py| {
        let dados = serde_json::json!({
            "resumo": {
                "passos": resumo.passos, "perda_inicial": resumo.perda_inicial,
                "perda_final": resumo.perda_final, "acc_recon": resumo.acc_recon,
                "delta_ort_max": resumo.delta_ort_max,
                "perplexidade_codebook": resumo.perplexidade_codebook,
                "monotonia_prefixo": resumo.monotonia_prefixo,
                "escore_aprendizado": resumo.escore_aprendizado,
            },
            "telemetria": hub.dump(),
        });
        Ok(pythonize::pythonize(py, &dados)?.into())
    })
}

/// Treina a fase B (gerador) com tokenizador congelado carregado de checkpoint.
#[pyfunction]
#[pyo3(signature = (cfg, sequencias, caminho_tokenizador, caminho_telemetria=None, caminho_checkpoint=None, callback=None, teto_seg=None))]
fn treinar_fase_b(
    cfg: &ConfigPy,
    sequencias: Vec<Vec<u32>>,
    caminho_tokenizador: &str,
    caminho_telemetria: Option<String>,
    caminho_checkpoint: Option<String>,
    callback: Option<Bound<'_, PyAny>>,
    teto_seg: Option<f64>,
) -> PyResult<PyObject> {
    let cb: treino::CallbackPy = callback.map(|c| c.unbind());
    let mut hub = telemetria::TelemetryHub::novo(caminho_telemetria);
    let dados = treino::DadosTreino::novo(sequencias, cfg.cfg.modelo.comprimento_ctx);
    // carrega o tokenizador congelado
    let dev = candle_core::Device::Cpu;
    let mut tok = tokenizador::TokenizadorConceitos::novo(&cfg.cfg, &dev)
        .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("construção: {e}")))?;
    carregar_tokenizador(&mut tok, caminho_tokenizador)
        .map_err(pyo3::exceptions::PyRuntimeError::new_err)?;
    let resumo = treino::treinar_fase_b(&cfg.cfg, &tok, &dados, &mut hub, &cb, caminho_checkpoint.as_deref(), teto_seg)
        .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("fase B: {e}")))?;
    Python::with_gil(|py| {
        let dados = serde_json::json!({
            "resumo": {
                "passos": resumo.passos, "perda_inicial": resumo.perda_inicial,
                "perda_final": resumo.perda_final, "acc_conceitos": resumo.acc_recon,
                "escore_aprendizado": resumo.escore_aprendizado,
            },
            "telemetria": hub.dump(),
        });
        Ok(pythonize::pythonize(py, &dados)?.into())
    })
}

fn carregar_tokenizador(tok: &mut tokenizador::TokenizadorConceitos, caminho: &str) -> Result<(), String> {
    let txt = std::fs::read_to_string(caminho).map_err(|e| format!("{e}"))?;
    let v: serde_json::Value = serde_json::from_str(&txt).map_err(|e| format!("{e}"))?;
    let obj = v.as_object().ok_or("formato de checkpoint inválido")?;
    let vars = tok.vars();
    for (i, var) in vars.iter().enumerate() {
        if let Some(entrada) = obj.get(&format!("v{i}")) {
            let forma: Vec<usize> = entrada["forma"].as_array().map(|a| a.iter().filter_map(|x| x.as_u64().map(|y| y as usize)).collect()).unwrap_or_default();
            let dados: Vec<f32> = entrada["dados"].as_array().map(|a| a.iter().filter_map(|x| x.as_f64().map(|y| y as f32)).collect()).unwrap_or_default();
            if forma.is_empty() || dados.is_empty() {
                continue;
            }
            let dev = candle_core::Device::Cpu;
            let t2 = candle_core::Tensor::from_vec(dados, forma.clone(), &dev)
                .and_then(|t| t.reshape(forma))
                .map_err(|e| format!("{e}"))?;
            var.set(&t2).map_err(|e| format!("{e}"))?;
        }
    }
    Ok(())
}

/// Gera amostras com NFE = L e devolve textos decodificados (modo conceitos).
#[pyfunction]
#[pyo3(signature = (cfg, caminho_tokenizador, caminho_gerador, n_amostras=2, parar_no_nivel=None, greedy=true))]
fn gerar_amostras(
    cfg: &ConfigPy,
    caminho_tokenizador: &str,
    caminho_gerador: &str,
    n_amostras: usize,
    parar_no_nivel: Option<usize>,
    greedy: bool,
) -> PyResult<PyObject> {
    let dev = candle_core::Device::Cpu;
    let mut tok = tokenizador::TokenizadorConceitos::novo(&cfg.cfg, &dev)
        .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("{e}")))?;
    carregar_tokenizador(&mut tok, caminho_tokenizador)
        .map_err(pyo3::exceptions::PyRuntimeError::new_err)?;
    let mut ger = gerador::GeradorConceitos::novo(&cfg.cfg, &dev)
        .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("{e}")))?;
    carregar_gerador(&mut ger, caminho_gerador)
        .map_err(pyo3::exceptions::PyRuntimeError::new_err)?;
    let mut amostrador = amostrador::AmostradorConceitos::novo(
        &mut ger, &tok, cfg.cfg.ncm.cfg_scale, cfg.cfg.ncm.temperaturas.limites, cfg.cfg.ncm.temperaturas.entropia_alvo,
    );
    let mut rng = <rand::prelude::StdRng as rand::SeedableRng>::seed_from_u64(4242);
    let inicio = std::time::Instant::now();
    let (codigos, nfe) = amostrador.gerar(n_amostras.max(1), &dev, &mut rng, parar_no_nivel, greedy)
        .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("amostragem: {e}")))?;
    let tempo_ms = inicio.elapsed().as_secs_f64() * 1000.0;
    // decodifica: ẑ final → decoder AR greedy → tokens
    let nivel_final = codigos.len();
    let zhats = tok.reconstruir_de_codigos(&codigos, nivel_final, n_amostras.max(1), &dev)
        .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("{e}")))?;
    let textos_ids = decodificar_textos(&tok, zhats.last(), n_amostras.max(1))
        .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("{e}")))?;
    Python::with_gil(|py| {
        let d = serde_json::json!({
            "nfe": nfe, "l": cfg.cfg.l_niveis(), "tempo_ms": tempo_ms,
            "niveis_usados": nivel_final, "ids_tokens": textos_ids,
            "temperaturas": amostrador.temperaturas,
        });
        Ok(pythonize::pythonize(py, &d)?.into())
    })
}

fn decodificar_textos(tok: &tokenizador::TokenizadorConceitos, zhat: Option<&candle_core::Tensor>, b: usize) -> Result<Vec<Vec<u32>>, String> {
    let Some(z) = zhat else { return Ok(vec![vec![]; b]) };
    tok.decodificar_livre(z, 48).map_err(|e| format!("{e}"))
}

fn carregar_gerador(ger: &mut gerador::GeradorConceitos, caminho: &str) -> Result<(), String> {
    let txt = std::fs::read_to_string(caminho).map_err(|e| format!("{e}"))?;
    let v: serde_json::Value = serde_json::from_str(&txt).map_err(|e| format!("{e}"))?;
    let obj = v.as_object().ok_or("formato inválido")?;
    let vars = ger.vars();
    for (i, var) in vars.iter().enumerate() {
        if let Some(entrada) = obj.get(&format!("v{i}")) {
            let forma: Vec<usize> = entrada["forma"].as_array().map(|a| a.iter().filter_map(|x| x.as_u64().map(|y| y as usize)).collect()).unwrap_or_default();
            let dados: Vec<f32> = entrada["dados"].as_array().map(|a| a.iter().filter_map(|x| x.as_f64().map(|y| y as f32)).collect()).unwrap_or_default();
            if forma.is_empty() || dados.is_empty() {
                continue;
            }
            let dev = candle_core::Device::Cpu;
            let t = candle_core::Tensor::from_vec(dados, forma.clone(), &dev)
                .and_then(|t| t.reshape(forma))
                .map_err(|e| format!("{e}"))?;
            var.set(&t).map_err(|e| format!("{e}"))?;
        }
    }
    Ok(())
}

/// Utilitários de prova expostos ao Python (V-série da suíte).
#[pyfunction]
fn verificar_mascara(cfg: &ConfigPy) -> PyResult<usize> {
    let dev = candle_core::Device::Cpu;
    let m = mascaras::mascara_block_causal(&cfg.cfg.ncm.niveis, gerador::N_PREFIXO, &dev)
        .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("{e}")))?;
    mascaras::contar_vazamentos(&cfg.cfg.ncm.niveis, gerador::N_PREFIXO, &m)
        .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("{e}")))
}

#[pyfunction]
fn verificar_stiefel(k: usize, d: usize) -> PyResult<f64> {
    let dev = candle_core::Device::Cpu;
    let c = stiefel::init_stiefel(k, d, &dev)
        .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("{e}")))?;
    stiefel::erro_ortogonalidade(&c.transpose(0, 1).unwrap()).map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("{e}")))
}

#[pymodule]
fn nucleo_rust(m: &Bound<'_, PyModule>) -> PyResult<()> {
    pyo3_log::init();
    m.add_function(wrap_pyfunction!(versao, m)?)?;
    m.add_function(wrap_pyfunction!(info_build, m)?)?;
    m.add_function(wrap_pyfunction!(lint_config, m)?)?;
    m.add_function(wrap_pyfunction!(treinar_fase_a, m)?)?;
    m.add_function(wrap_pyfunction!(treinar_fase_b, m)?)?;
    m.add_function(wrap_pyfunction!(gerar_amostras, m)?)?;
    m.add_function(wrap_pyfunction!(verificar_mascara, m)?)?;
    m.add_function(wrap_pyfunction!(verificar_stiefel, m)?)?;
    m.add_class::<BpePy>()?;
    m.add_class::<ConfigPy>()?;
    m.add_class::<PonteDuplaPy>()?;
    #[cfg(feature = "ponte_async")]
    m.add_class::<PonteAsyncPy>()?;
    Ok(())
}