pub mod db;
pub mod entities;
pub mod llama;
pub mod mcp;
#[cfg(test)]
mod tests;

use crate::llama::LlamaClient;
use dotenvy::dotenv;
use sea_orm::DatabaseConnection;
use std::env;
use std::path::PathBuf;
use std::sync::Arc;
use tauri::menu::{Menu, MenuItem};
use tauri::tray::{TrayIconBuilder, TrayIconEvent};
use tauri::{Emitter, Manager};
use sea_orm::{ConnectionTrait, Statement, DatabaseBackend};

use std::sync::atomic::AtomicUsize;

pub struct AppState {
    pub db: DatabaseConnection,
    pub llama: Arc<LlamaClient>,
    pub mcp_tx: tokio::sync::broadcast::Sender<serde_json::Value>,
    pub connection_count: Arc<AtomicUsize>,
    pub app_handle: tauri::AppHandle,
}

// --- Path Resolution Helpers ---

// --- Path Resolution Helpers ---

pub fn find_build_assets_dir_logic(res_dir: Option<PathBuf>, exe_dir: Option<PathBuf>) -> Option<PathBuf> {
    // 1. Check Resource Dir (Release / Installed) - prioritizes nested structure
    if let Some(res_dir) = res_dir {
        let candidate = res_dir.join("build_assets");
        if candidate.exists() {
            log::info!("Found build_assets in resource_dir: {:?}", candidate);
            return Some(candidate);
        }
    }

    // 2. Check relative to Executable (Debug / Dev) - walks up directory tree
    if let Some(exe_dir) = exe_dir {
        let mut p = exe_dir;
        for _ in 0..10 { // Increased depth to check up to 9 levels up
            let candidate = p.join("build_assets");
            if candidate.exists() {
                log::info!("Found build_assets relative to exe: {:?}", candidate);
                return Some(candidate);
            }
            if !p.pop() { break; }
        }
    }
    None
}

pub fn find_build_assets_dir(app_handle: &tauri::AppHandle) -> Option<PathBuf> {
    let res_dir = app_handle.path().resource_dir().ok();
    let exe_dir = env::current_exe().ok().and_then(|p| p.parent().map(|p| p.to_path_buf()));
    find_build_assets_dir_logic(res_dir, exe_dir)
}

pub fn resolve_resource_path(app_handle: &tauri::AppHandle, relative_path: &str) -> Option<PathBuf> {
    find_build_assets_dir(app_handle).map(|dir| dir.join(relative_path))
}

fn diagnose_environment(app_handle: &tauri::AppHandle) {
    log::info!("--- [DIAGNOSIS] Environment Check Start ---");
    
    // 1. Build Assets Dir
    match find_build_assets_dir(app_handle) {
        Some(path) => log::info!("✅ build_assets found at: {:?}", path),
        None => log::error!("❌ build_assets NOT FOUND. Sidecar and resources will fail."),
    }

    // 2. Critical Files Check
    let critical_files = ["mcp.json", "config.json", "vec0.dll", "llama.dll"];
    for file in critical_files {
        match resolve_resource_path(app_handle, file) {
            Some(p) if p.exists() => log::info!("✅ found {}: {:?}", file, p),
            Some(p) => log::error!("❌ missing {}: path resolved to {:?} but file does not exist", file, p),
            None => log::error!("❌ could not resolve path for {}", file),
        }
    }
    
    // 3. Model file check (sample)
    if let Some(base) = find_build_assets_dir(app_handle) {
        let model_dir = base.join("models");
        if model_dir.exists() {
             log::info!("✅ models directory exists at: {:?}", model_dir);
        } else {
             log::warn!("⚠️ models directory NOT found at: {:?}", model_dir);
        }
    }
    
    log::info!("--- [DIAGNOSIS] Environment Check End ---");
}

#[tauri::command]
fn get_mcp_info(app_handle: tauri::AppHandle) -> Result<serde_json::Value, String> {
    let mcp_path = resolve_resource_path(&app_handle, "mcp.json")
        .ok_or_else(|| "Failed to resolve mcp.json path".to_string())?;
        
    if !mcp_path.exists() {
        return Err(format!("mcp.json not found at {:?}", mcp_path));
    }

    let content = std::fs::read_to_string(&mcp_path).map_err(|e| e.to_string())?;
    let mcp_data: serde_json::Value = serde_json::from_str(&content).map_err(|e| e.to_string())?;
    Ok(mcp_data)
}

#[tauri::command]
async fn get_db_stats(state: tauri::State<'_, Arc<AppState>>) -> Result<serde_json::Value, String> {
    use crate::entities::items;
    use sea_orm::{EntityTrait, PaginatorTrait};
    let count = items::Entity::find().count(&state.db).await.map_err(|e| e.to_string())?;
    Ok(serde_json::json!({ "itemCount": count }))
}


#[tauri::command]
async fn get_sidecar_status(state: tauri::State<'_, Arc<AppState>>) -> Result<bool, String> {
    Ok(state.llama.check_health().await)
}

// --- DB Browser Commands ---

fn validate_table_name(table_name: &str) -> Result<(), String> {
    if table_name.chars().all(|c| c.is_alphanumeric() || c == '_' || c == '-') {
        Ok(())
    } else {
        Err("Invalid table name".to_string())
    }
}

#[tauri::command]
async fn get_table_list(state: tauri::State<'_, Arc<AppState>>) -> Result<Vec<String>, String> {
    let results = state.db.query_all(Statement::from_string(
        DatabaseBackend::Sqlite,
        "SELECT name FROM sqlite_master WHERE type='table' AND name NOT LIKE 'sqlite_%'"
    )).await.map_err(|e| e.to_string())?;

    let mut tables = Vec::new();
    for res in results {
        let name: String = res.try_get("", "name").map_err(|e| e.to_string())?;
        tables.push(name);
    }
    Ok(tables)
}

#[tauri::command]
async fn get_table_schema(state: tauri::State<'_, Arc<AppState>>, table_name: String) -> Result<serde_json::Value, String> {
    validate_table_name(&table_name)?;
    let results = state.db.query_all(Statement::from_string(
        DatabaseBackend::Sqlite,
        format!("PRAGMA table_info({})", table_name)
    )).await.map_err(|e| e.to_string())?;

    let mut schema = Vec::new();
    for res in results {
        schema.push(serde_json::json!({
            "cid": res.try_get::<i32>("", "cid").map_err(|e| e.to_string())?,
            "name": res.try_get::<String>("", "name").map_err(|e| e.to_string())?,
            "type": res.try_get::<String>("", "type").map_err(|e| e.to_string())?,
            "notnull": res.try_get::<i32>("", "notnull").map_err(|e| e.to_string())?,
            "dflt_value": res.try_get::<Option<String>>("", "dflt_value").map_err(|e| e.to_string())?,
            "pk": res.try_get::<i32>("", "pk").map_err(|e| e.to_string())?,
        }));
    }
    Ok(serde_json::json!(schema))
}

#[tauri::command]
async fn get_table_data(
    state: tauri::State<'_, Arc<AppState>>,
    table_name: String,
    limit: i64,
    offset: i64
) -> Result<serde_json::Value, String> {
    validate_table_name(&table_name)?;
    
    // Get column names first to handle dynamic results
    let schema_res = get_table_schema(state.clone(), table_name.clone()).await?;
    let columns: Vec<String> = schema_res.as_array()
        .ok_or("Failed to parse schema")?
        .iter()
        .map(|c| c["name"].as_str().unwrap_or("").to_string())
        .collect();

    let results = state.db.query_all(Statement::from_string(
        DatabaseBackend::Sqlite,
        format!("SELECT * FROM {} LIMIT {} OFFSET {}", table_name, limit, offset)
    )).await.map_err(|e| e.to_string())?;

    let mut data = Vec::new();
    for res in results {
        let mut row = serde_json::Map::new();
        for col in &columns {
            // Try as different types since SeaORM Value is complex
            if let Ok(val) = res.try_get::<i32>("", col) {
                row.insert(col.clone(), serde_json::json!(val));
            } else if let Ok(val) = res.try_get::<String>("", col) {
                row.insert(col.clone(), serde_json::json!(val));
            } else if let Ok(val) = res.try_get::<Option<String>>("", col) {
                row.insert(col.clone(), serde_json::json!(val));
            } else if let Ok(val) = res.try_get::<f64>("", col) {
                row.insert(col.clone(), serde_json::json!(val));
            } else if let Ok(val) = res.try_get::<Vec<u8>>("", col) {
                row.insert(col.clone(), serde_json::json!(format!("<BLOB: {} bytes>", val.len())));
            } else {
                row.insert(col.clone(), serde_json::json!(null));
            }
        }
        data.push(serde_json::Value::Object(row));
    }
    
    let total_count: i64 = match table_name.as_str() {
        "items" => {
            use crate::entities::items;
            use sea_orm::{EntityTrait, PaginatorTrait};
            items::Entity::find().count(&state.db).await.map_err(|e| e.to_string())? as i64
        },
        _ => {
            let res = state.db.query_one(Statement::from_string(
                DatabaseBackend::Sqlite,
                format!("SELECT COUNT(*) as count FROM {}", table_name)
            )).await.map_err(|e| e.to_string())?;
            res.map(|r| r.try_get::<i64>("", "count").unwrap_or(0)).unwrap_or(0)
        }
    };

    Ok(serde_json::json!({
        "data": data,
        "total": total_count
    }))
}

fn get_config(app_handle: &tauri::AppHandle) -> serde_json::Value {
    // 1. Try unified resource path first
    if let Some(path) = resolve_resource_path(app_handle, "config.json") {
        log::info!("Checking config.json at resolved path: {:?}", path);
        if path.exists() {
             if let Ok(content) = std::fs::read_to_string(&path) {
                if let Ok(parsed) = serde_json::from_str::<serde_json::Value>(&content) { 
                    return parsed; 
                }
            }
        }
    }

    // 2. Fallbacks for flexibility (AppConfig direct, CWD)
    let mut config_paths = vec![];
    if let Ok(app_data) = app_handle.path().app_config_dir() { config_paths.push(app_data.join("config.json")); }
    if let Ok(cwd) = env::current_dir() { config_paths.push(cwd.join("config.json")); }

    for path in config_paths {
        if path.exists() {
            if let Ok(content) = std::fs::read_to_string(path) {
                if let Ok(parsed) = serde_json::from_str::<serde_json::Value>(&content) { return parsed; }
            }
        }
    }
    
    // Default fallback
    serde_json::json!({
        "database": { "path": "data/vector.db" },
        "model": { "path": "models/embeddinggemma-300m-q4_0.gguf" },
        "llama_server": { "port": 8080 }
    })
}

fn resolve_db_path(app_handle: &tauri::AppHandle, config: &serde_json::Value) -> String {
    if let Ok(p) = env::var("DB_PATH") { return p; }
    if let Some(p) = config.get("database").and_then(|d| d.get("path")).and_then(|p| p.as_str()) {
        let mut candidate = PathBuf::from(p);
        if candidate.is_relative() {
            if cfg!(debug_assertions) {
                if let Ok(exe_path) = env::current_exe() {
                    if let Some(exe_dir) = exe_path.parent() {
                        let mut pr = exe_dir.to_path_buf();
                        for _ in 0..4 {
                            if pr.join("config.json").exists() { candidate = pr.join(p); break; }
                            if !pr.pop() { break; }
                        }
                    }
                }
            } else {
                let mut p_base = app_handle.path().app_data_dir().expect("App data dir not found");
                p_base.push("data");
                let _ = std::fs::create_dir_all(&p_base);
                p_base.push("vector.db");
                candidate = p_base;
            }
        }
        return candidate.to_string_lossy().to_string();
    }
    
    // Default fallback
    if cfg!(debug_assertions) { "data/vector.db".to_string() } else {
        let mut p = app_handle.path().app_data_dir().expect("App data dir not found");
        p.push("data");
        let _ = std::fs::create_dir_all(&p);
        p.push("vector.db");
        p.to_string_lossy().to_string()
    }
}

fn resolve_extension_path(app_handle: &tauri::AppHandle) -> String {
    // 1. Check verified resource path (build_assets/vec0.dll)
    if let Some(path) = resolve_resource_path(app_handle, "vec0.dll") {
        if path.exists() {
             return path.to_string_lossy().to_string();
        }
    }

    // 2. Search relative to EXE for local dev
    if let Ok(exe_path) = env::current_exe() {
        if let Some(mut p) = exe_path.parent() {
            let mut pr = p.to_path_buf();
            for _ in 0..10 {
                // Check in bin/
                let bin_cand = pr.join("bin").join("vec0.dll");
                if bin_cand.exists() { return bin_cand.to_string_lossy().to_string(); }
                
                // Check in node_modules (legacy/other)
                let nm_cand = pr.join("node_modules/sqlite-vec-windows-x64/vec0.dll");
                if nm_cand.exists() { return nm_cand.to_string_lossy().to_string(); }

                if !pr.pop() { break; }
            }
        }
    }

    "vec0.dll".to_string()
}

fn spawn_llama_server(app_handle: &tauri::AppHandle, _config: &serde_json::Value) {
    use tauri_plugin_shell::ShellExt;

    let sidecar = match app_handle.shell().sidecar("llama-server") {
        Ok(s) => s,
        Err(e) => {
            eprintln!("CRITICAL: Failed to create sidecar handle: {}", e);
            return;
        }
    };

    // 1. Resolve base_dir using unified logic
    let base_dir = match find_build_assets_dir(app_handle) {
        Some(p) => p,
        None => {
            log::error!("CRITICAL: Could not find build_assets. Sidecar will likely fail due to missing DLLs/Models.");
            // Fallback to something sane? Or just exe_dir?
            env::current_exe()
                .map(|p| p.parent().unwrap_or(&p).to_path_buf())
                .unwrap_or_else(|_| PathBuf::from("."))
        }
    };

    // 2. Pre-launch Integrity Check (DLLs)
    let required_dlls = ["llama.dll", "ggml.dll"]; // Add others if known
    for dll in required_dlls {
        let dll_path = base_dir.join(dll);
        if !dll_path.exists() {
             log::error!("CRITICAL: Required DLL not found at: {:?}. Llama server may fail to start.", dll_path);
        } else {
             log::info!("Verified DLL exists: {:?}", dll_path);
        }
    }
    
    // Model path should be in base_dir/models/...
    let mut model_path = base_dir.join("models").join("embeddinggemma-300m-q4_0.gguf");
    if !model_path.exists() {
        log::warn!("Model not found at default location: {:?}. Checking fallback...", model_path);
         // Double check if config has an absolute path? 
         // For now, assuming standard layout.
    }
    log::info!("Using base_dir for sidecar: {:?}", base_dir);
    log::info!("Using model path for sidecar: {:?}", model_path);

    // Prepare arguments
    let args = [
        "--model", model_path.to_str().unwrap_or(""),
        "--port", "8080",
        "--embedding",
        "--host", "127.0.0.1",
        "-c", "8192", "-b", "8192", "-ub", "8192",
        "--parallel", "1"
    ];

    // 3. Update PATH environment variable to include base_dir and bin dir
    let old_path = std::env::var("PATH").unwrap_or_default();
    let bin_dir = base_dir.parent().map(|p| p.join("bin")).unwrap_or_else(|| PathBuf::from("bin"));
    let new_path = format!("{};{};{}", base_dir.display(), bin_dir.display(), old_path).replace("\\\\?\\", "");
    log::info!("Setting Sidecar PATH: {}", new_path);
    
    let mut sidecar_cmd = sidecar.args(args).env("PATH", new_path);
    
    // Set current directory to base_dir as well
    if base_dir.exists() {
        sidecar_cmd = sidecar_cmd.current_dir(&base_dir);
    }

    let (mut rx, _child) = match sidecar_cmd.spawn() {
        Ok(res) => res,
        Err(e) => {
            log::error!("CRITICAL: Failed to spawn sidecar: {}", e);
            return;
        }
    };

    println!("llama-server started (Sidecar)");

    // Handle output
    tauri::async_runtime::spawn(async move {
        while let Some(event) = rx.recv().await {
             match event {
                tauri_plugin_shell::process::CommandEvent::Stdout(line) => {
                    let s = String::from_utf8_lossy(&line);
                    if s.contains("HTTP server listening") {
                        println!("llama-server: Ready");
                    }
                }
                tauri_plugin_shell::process::CommandEvent::Stderr(line) => {
                    eprintln!("llama-server error: {}", String::from_utf8_lossy(&line));
                }
                tauri_plugin_shell::process::CommandEvent::Terminated(status) => {
                    println!("llama-server terminated with status: {:?}", status.code);
                    break;
                }
                _ => {}
            }
        }
    });
}

#[cfg_attr(mobile, tauri::mobile_entry_point)]
pub fn run() {
    tauri::Builder::default()
        .plugin(tauri_plugin_shell::init())
        .on_window_event(|window, event| {
            if let tauri::WindowEvent::CloseRequested { api, .. } = event {
                api.prevent_close();
                let _ = window.hide();
            }
        })
        .setup(|app| {
            // 1. Initialize Logging FIRST for crash diagnostics
            let mut log_builder = tauri_plugin_log::Builder::default()
                .targets([
                    tauri_plugin_log::Target::new(tauri_plugin_log::TargetKind::Stdout),
                    tauri_plugin_log::Target::new(tauri_plugin_log::TargetKind::LogDir { file_name: None }),
                ])
                .max_file_size(10 * 1024 * 1024)
                .level(log::LevelFilter::Info);
            
            if cfg!(debug_assertions) {
                log_builder = log_builder.target(tauri_plugin_log::Target::new(tauri_plugin_log::TargetKind::Folder { path: std::path::PathBuf::from("logs"), file_name: None }));
            }
            let _ = app.handle().plugin(log_builder.build());

            log::info!("Application starting (Tauri 2)...");
            
            // 2. DIAGNOSIS (New)
            diagnose_environment(app.handle());

            // 3. Load env

            dotenv().ok();
            let app_handle = app.handle().clone();

            // 3. Start Sidecar
            let config = get_config(&app_handle);
            spawn_llama_server(&app_handle, &config);

            // 4. Database and MCP initialization
            tauri::async_runtime::block_on(async move {
                let db_path = resolve_db_path(&app_handle, &config);
                let ext_path = resolve_extension_path(&app_handle);
                log::info!("DB Path: {}, Ext Path: {}", db_path, ext_path);
                
                let vec_dim = env::var("VEC_DIM").unwrap_or_else(|_| "768".to_string()).parse::<usize>().unwrap_or(768);
                let conn = db::init_db(&db_path, &ext_path, vec_dim).await.expect("Failed to init db");

                let state = Arc::new(AppState {
                    db: conn,
                    llama: Arc::new(LlamaClient::new(
                        env::var("LLAMA_CPP_BASE_URL").unwrap_or_else(|_| "http://localhost:8080".to_string()),
                        env::var("LLAMA_CPP_EMBEDDING_MODEL").unwrap_or_else(|_| "nomic-embed-text".to_string()),
                        env::var("LLAMA_CPP_MODEL").unwrap_or_else(|_| "mistral".to_string()),
                    )),
                    mcp_tx: tokio::sync::broadcast::channel(100).0,
                    connection_count: Arc::new(AtomicUsize::new(0)),
                    app_handle: app_handle.clone(),
                });
                app_handle.manage(state.clone());
                let port = env::var("MCP_PORT").unwrap_or_else(|_| "3000".to_string()).parse::<u16>().unwrap_or(3000);
                tokio::spawn(async move { mcp::start_mcp_server(state, port).await; });
            });

            // 5. Tray setup
            let quit_i = MenuItem::with_id(app, "quit", "Quit", true, None::<&str>)?;
            let show_i = MenuItem::with_id(app, "show", "Show", true, None::<&str>)?;
            let menu = Menu::with_items(app, &[&show_i, &quit_i])?;
            let _tray = TrayIconBuilder::new()
                .icon(app.default_window_icon().unwrap().clone())
                .menu(&menu)
                .on_menu_event(|app, event| match event.id.as_ref() {
                    "quit" => app.exit(0),
                    "show" => if let Some(w) = app.get_webview_window("main") { let _ = w.show(); let _ = w.set_focus(); }
                    _ => {}
                })
                .on_tray_icon_event(|tray, event| if let TrayIconEvent::DoubleClick { .. } = event {
                    let app = tray.app_handle();
                    if let Some(w) = app.get_webview_window("main") { let _ = w.show(); let _ = w.set_focus(); }
                })
                .build(app)?;
            
            log::info!("Tauri setup completed");
            Ok(())
        })
        .invoke_handler(tauri::generate_handler![
            get_mcp_info,
            get_db_stats,
            get_sidecar_status,
            get_table_list,
            get_table_schema,
            get_table_data
        ])
        .run(tauri::generate_context!())
        .expect("error while running tauri application");
}