//! Generate Python stub files for coglet.
//!
//! Run with: cargo run --bin stub_gen
//!
//! Custom generate logic: pyo3-stub-gen places classes from the native
//! `coglet._impl` module into the `coglet` parent package stub.  We redirect
//! that output to `coglet/_impl.pyi` so native module types are preserved in
//! the right place, then generate `coglet/__init__.pyi` ourselves to
//! re-export the public API — matching the hand-maintained `__init__.py`.

use pyo3_stub_gen::Result;
use std::fs;
use std::io::Write;

/// Public items re-exported from `coglet._impl` in `coglet/__init__.pyi`.
/// Uses the `X as X` pattern to mark explicit re-exports (PEP 484).
const PUBLIC_REEXPORTS: &[&str] = &["__build__", "__version__", "server", "CancelationException"];

/// Private submodules re-exported with `from . import X as X`.
///
/// These use a relative import (not `from coglet._impl`) because `_sdk` is a
/// subpackage that type checkers resolve via the filesystem, not an attribute
/// of the native extension module.  Not included in `__all__`.
const PRIVATE_REEXPORTS: &[&str] = &["_sdk"];

fn main() -> Result<()> {
    let stub = coglet::stub_info()?;

    for (name, module) in &stub.modules {
        let normalized = name.replace('-', "_");

        let dest = if normalized == "coglet" {
            // Native module classes land here — redirect to _impl.pyi
            stub.python_root.join("coglet").join("_impl.pyi")
        } else {
            // Submodules like "coglet._sdk" → coglet/_sdk/__init__.pyi
            let path = normalized.replace('.', "/");
            stub.python_root.join(&path).join("__init__.pyi")
        };

        let dir = dest.parent().expect("cannot get parent directory");
        if !dir.exists() {
            fs::create_dir_all(dir)?;
        }

        let mut f = fs::File::create(&dest)?;
        write!(f, "{module}")?;
        eprintln!("Generated stub: {}", dest.display());
    }

    // Generate coglet/__init__.pyi — re-exports from _impl
    let init_pyi = stub.python_root.join("coglet").join("__init__.pyi");
    let mut f = fs::File::create(&init_pyi)?;

    writeln!(f, "# This file is automatically generated by stub_gen")?;
    writeln!(f, "# ruff: noqa: E501, F401")?;
    writeln!(f)?;

    // `from coglet._impl import X as X, Y as Y, ...`
    let reexports: Vec<String> = PUBLIC_REEXPORTS
        .iter()
        .map(|name| format!("{name} as {name}"))
        .collect();
    writeln!(f, "from coglet._impl import {}", reexports.join(", "))?;

    // `from . import _sdk as _sdk` — relative import so ty resolves the
    // subpackage via coglet/_sdk/__init__.pyi, not through _impl.
    let private: Vec<String> = PRIVATE_REEXPORTS
        .iter()
        .map(|name| format!("{name} as {name}"))
        .collect();
    writeln!(f, "from . import {}", private.join(", "))?;
    writeln!(f)?;

    // __all__ only includes public items (no underscore-prefixed names)
    let all_items: Vec<String> = PUBLIC_REEXPORTS
        .iter()
        .map(|name| format!("'{name}'"))
        .collect();
    writeln!(f, "__all__ = [{}]", all_items.join(", "))?;

    eprintln!("Generated stub: {}", init_pyi.display());

    Ok(())
}
