ebook-rs 0.15.3

use ebook_rs::{Book, ReaderServer, generate_sample_epub};
use std::env;
use std::fs;
use std::path::Path;

#[cfg(feature = "mimalloc")]
#[global_allocator]
static GLOBAL: mimalloc::MiMalloc = mimalloc::MiMalloc;

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let args: Vec<String> = env::args().collect();

    let mode = args.get(1).map(|s| s.as_str()).unwrap_or("serve");

    match mode {
        "mcp" => {
            eprintln!("🚀 Starting ebook-rs Model Context Protocol (MCP) server on stdio...");
            ebook_rs::run_mcp_server()?;
        }
        "parse" => {
            let path = args.get(2).ok_or("Usage: ebook-rs parse <path.epub>")?;
            let book = Book::from_file(path)?;
            println!(
                "📖 Metadata:\n{}",
                serde_json::to_string_pretty(book.metadata())?
            );
            println!(
                "\n📌 Table of Contents:\n{}",
                serde_json::to_string_pretty(book.toc())?
            );
            println!(
                "\n📚 Spine Items ({} total):\n{}",
                book.spine().len(),
                serde_json::to_string_pretty(book.spine())?
            );
        }
        "search" => {
            let path = args
                .get(2)
                .ok_or("Usage: ebook-rs search <path.epub> <query>")?;
            let query = args
                .get(3)
                .ok_or("Usage: ebook-rs search <path.epub> <query>")?;
            let book = Book::from_file(path)?;
            let results = book.search(query);
            println!(
                "🔍 Search results for '{}' ({} found):",
                query,
                results.len()
            );
            for r in results {
                println!(
                    "[Spine {}] CFI: {}\n    Snippet: {}\n",
                    r.spine_index, r.cfi, r.snippet
                );
            }
        }
        "locations" => {
            let path = args.get(2).ok_or("Usage: ebook-rs locations <path.epub>")?;
            let book = Book::from_file(path)?;
            println!(
                "📍 Generated Locations ({} total):",
                book.locations.total_locations
            );
            println!("{}", serde_json::to_string_pretty(&book.locations.entries)?);
        }
        "sample" => {
            let out_path = args.get(2).map(|s| s.as_str()).unwrap_or("sample.epub");
            let bytes = generate_sample_epub()?;
            fs::write(out_path, bytes)?;
            println!("✅ Generated sample EPUB file at: {}", out_path);
        }
        "convert" => {
            let in_path = args
                .get(2)
                .ok_or("Usage: ebook-rs convert <input.epub/mobi/pdf/fb2/txt/kfx> <output.epub|output.kfx|output.json>")?;
            let out_path = args
                .get(3)
                .ok_or("Usage: ebook-rs convert <input.epub/mobi/pdf/fb2/txt/kfx> <output.epub|output.kfx|output.json>")?;

            println!("🔄 Loading eBook from: {}", in_path);
            let book = Book::from_file(in_path)?;

            if out_path.ends_with(".epub") {
                let epub_bytes = ebook_rs::UniversalEpub3Exporter::export(&book)?;
                fs::write(out_path, epub_bytes)?;
                println!("✅ Converted and exported EPUB 3 file to: {}", out_path);
            } else if out_path.ends_with(".kfx") {
                let kfx_bytes = ebook_rs::UniversalKfxExporter::export(&book)?;
                fs::write(out_path, kfx_bytes)?;
                println!("✅ Converted and exported Amazon KFX file to: {}", out_path);
            } else if out_path.ends_with(".json") {
                let chunks = book.to_rag_chunks(&ebook_rs::RagChunkConfig::default());
                let json = serde_json::to_string_pretty(&chunks)?;
                fs::write(out_path, json)?;
                println!("✅ Exported RAG chunks JSON file to: {}", out_path);
            } else {
                return Err("Unsupported output extension. Supported: .epub, .kfx, .json".into());
            }
        }
        "rag" => {
            let path = args
                .get(2)
                .ok_or("Usage: ebook-rs rag <path.epub> [max_tokens]")?;
            let max_tokens: usize = args.get(3).and_then(|s| s.parse().ok()).unwrap_or(512);

            let book = Book::from_file(path)?;
            let config = ebook_rs::RagChunkConfig {
                max_tokens,
                ..Default::default()
            };
            let chunks = book.to_rag_chunks(&config);
            println!("{}", serde_json::to_string_pretty(&chunks)?);
        }
        _ => {
            let file_arg = if args.len() > 1 && !args[1].starts_with('-') && mode != "serve" {
                Some(&args[1])
            } else if args.len() > 2 {
                Some(&args[2])
            } else {
                None
            };

            let port: u16 = env::var("PORT")
                .unwrap_or_else(|_| "8080".to_string())
                .parse()
                .unwrap_or(8080);

            let book = if let Some(p) = file_arg {
                if Path::new(p).exists() {
                    println!("📖 Loading EPUB from file: {}", p);
                    Book::from_file(p)?
                } else {
                    println!(
                        "⚠️ Specified file '{}' not found. Generating sample EPUB in memory...",
                        p
                    );
                    let bytes = generate_sample_epub()?;
                    Book::from_bytes(&bytes)?
                }
            } else {
                println!("💡 No EPUB file path provided. Generating built-in sample EPUB 3...");
                let bytes = generate_sample_epub()?;
                Book::from_bytes(&bytes)?
            };

            println!("✨ Title: {}", book.metadata().title);
            println!("✨ Author(s): {}", book.metadata().creators.join(", "));
            println!("✨ Spine Sections: {}", book.spine().len());
            println!("✨ Locations Generated: {}", book.locations.total_locations);

            let server = ReaderServer::new(book, port);
            server.listen()?;
        }
    }

    Ok(())
}