/*!
    Thread and Message Persistence Module

    This module provides all logic for managing threads and their messages, including creation, modification, deletion, and listing.
    Messages for each thread are persisted in a JSONL file (messages.jsonl) per thread directory.

    **Concurrency and Consistency Guarantee:**
    - All operations that write or modify messages for a thread are protected by a global, per-thread asynchronous lock.
    - This design ensures that only one operation can write to a thread's messages.jsonl file at a time, preventing race conditions.
    - As a result, the messages.jsonl file for each thread is always consistent and never corrupted, even under concurrent access.
*/

use serde::{Deserialize, Serialize};
use std::fs::{self, File};
use std::io::{BufRead, BufReader, Write};
use tauri::command;
use tauri::Runtime;
use uuid::Uuid;

// For async file write serialization
use once_cell::sync::Lazy;
use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::Mutex;

// Global per-thread locks for message file writes
static MESSAGE_LOCKS: Lazy<Mutex<HashMap<String, Arc<Mutex<()>>>>> =
    Lazy::new(|| Mutex::new(HashMap::new()));

use super::utils::{
    ensure_data_dirs, ensure_thread_dir_exists, get_data_dir, get_messages_path, get_thread_dir,
    get_thread_metadata_path, THREADS_FILE,
};

#[derive(Debug, Serialize, Deserialize, Clone)]
pub struct Thread {
    pub id: String,
    pub object: String,
    pub title: String,
    pub assistants: Vec<ThreadAssistantInfo>,
    pub created: i64,
    pub updated: i64,
    pub metadata: Option<serde_json::Value>,
}

#[derive(Debug, Serialize, Deserialize, Clone)]
pub struct ThreadMessage {
    pub id: String,
    pub object: String,
    pub thread_id: String,
    pub assistant_id: Option<String>,
    pub attachments: Option<Vec<Attachment>>,
    pub role: String,
    pub content: Vec<ThreadContent>,
    pub status: String,
    pub created_at: i64,
    pub completed_at: i64,
    pub metadata: Option<serde_json::Value>,
    pub type_: Option<String>,
    pub error_code: Option<String>,
    pub tool_call_id: Option<String>,
}

#[derive(Debug, Serialize, Deserialize, Clone)]
pub struct Attachment {
    pub file_id: Option<String>,
    pub tools: Option<Vec<Tool>>,
}

#[derive(Debug, Serialize, Deserialize, Clone)]
#[serde(tag = "type")]
pub enum Tool {
    #[serde(rename = "file_search")]
    FileSearch,
    #[serde(rename = "code_interpreter")]
    CodeInterpreter,
}

#[derive(Debug, Serialize, Deserialize, Clone)]
pub struct ThreadContent {
    pub type_: String,
    pub text: Option<ContentValue>,
    pub image_url: Option<ImageContentValue>,
}

#[derive(Debug, Serialize, Deserialize, Clone)]
pub struct ContentValue {
    pub value: String,
    pub annotations: Vec<String>,
}

#[derive(Debug, Serialize, Deserialize, Clone)]
pub struct ImageContentValue {
    pub detail: Option<String>,
    pub url: Option<String>,
}

#[derive(Debug, Serialize, Deserialize, Clone)]
pub struct ThreadAssistantInfo {
    pub assistant_id: String,
    pub assistant_name: String,
    pub model: ModelInfo,
    pub instructions: Option<String>,
    pub tools: Option<Vec<AssistantTool>>,
}

#[derive(Debug, Serialize, Deserialize, Clone)]
pub struct ModelInfo {
    pub id: String,
    pub name: String,
    pub settings: serde_json::Value,
}

#[derive(Debug, Serialize, Deserialize, Clone)]
#[serde(tag = "type")]
pub enum AssistantTool {
    #[serde(rename = "code_interpreter")]
    CodeInterpreter,
    #[serde(rename = "retrieval")]
    Retrieval,
    #[serde(rename = "function")]
    Function {
        name: String,
        description: Option<String>,
        parameters: Option<serde_json::Value>,
    },
}

#[derive(Debug, Serialize, Deserialize, Clone)]
pub struct ThreadState {
    pub has_more: bool,
    pub waiting_for_response: bool,
    pub error: Option<String>,
    pub last_message: Option<String>,
}

/// Lists all threads by reading their metadata from the threads directory.
/// Returns a vector of thread metadata as JSON values.
#[command]
pub async fn list_threads<R: Runtime>(
    app_handle: tauri::AppHandle<R>,
) -> Result<Vec<serde_json::Value>, String> {
    ensure_data_dirs(app_handle.clone())?;
    let data_dir = get_data_dir(app_handle.clone());
    let mut threads = Vec::new();

    if !data_dir.exists() {
        return Ok(threads);
    }

    for entry in fs::read_dir(&data_dir).map_err(|e| e.to_string())? {
        let entry = entry.map_err(|e| e.to_string())?;
        let path = entry.path();
        if path.is_dir() {
            let thread_metadata_path = path.join(THREADS_FILE);
            if thread_metadata_path.exists() {
                let data = fs::read_to_string(&thread_metadata_path).map_err(|e| e.to_string())?;
                match serde_json::from_str(&data) {
                    Ok(thread) => threads.push(thread),
                    Err(e) => {
                        println!("Failed to parse thread file: {}", e);
                        continue; // skip invalid thread files
                    }
                }
            }
        }
    }

    Ok(threads)
}

/// Creates a new thread, assigns it a unique ID, and persists its metadata.
/// Ensures the thread directory exists and writes thread.json.
#[command]
pub async fn create_thread<R: Runtime>(
    app_handle: tauri::AppHandle<R>,
    mut thread: serde_json::Value,
) -> Result<serde_json::Value, String> {
    ensure_data_dirs(app_handle.clone())?;
    let uuid = Uuid::new_v4().to_string();
    thread["id"] = serde_json::Value::String(uuid.clone());
    let thread_dir = get_thread_dir(app_handle.clone(), &uuid);
    if !thread_dir.exists() {
        fs::create_dir_all(&thread_dir).map_err(|e| e.to_string())?;
    }
    let path = get_thread_metadata_path(app_handle.clone(), &uuid);
    let data = serde_json::to_string_pretty(&thread).map_err(|e| e.to_string())?;
    fs::write(path, data).map_err(|e| e.to_string())?;
    Ok(thread)
}

/// Modifies an existing thread's metadata by overwriting its thread.json file.
/// Returns an error if the thread directory does not exist.
#[command]
pub async fn modify_thread<R: Runtime>(
    app_handle: tauri::AppHandle<R>,
    thread: serde_json::Value,
) -> Result<(), String> {
    let thread_id = thread
        .get("id")
        .and_then(|id| id.as_str())
        .ok_or("Missing thread id")?;
    let thread_dir = get_thread_dir(app_handle.clone(), thread_id);
    if !thread_dir.exists() {
        return Err("Thread directory does not exist".to_string());
    }
    let path = get_thread_metadata_path(app_handle.clone(), thread_id);
    let data = serde_json::to_string_pretty(&thread).map_err(|e| e.to_string())?;
    fs::write(path, data).map_err(|e| e.to_string())?;
    Ok(())
}

/// Deletes a thread and all its associated files by removing its directory.
#[command]
pub async fn delete_thread<R: Runtime>(
    app_handle: tauri::AppHandle<R>,
    thread_id: String,
) -> Result<(), String> {
    let thread_dir = get_thread_dir(app_handle.clone(), &thread_id);
    if thread_dir.exists() {
        fs::remove_dir_all(thread_dir).map_err(|e| e.to_string())?;
    }
    Ok(())
}

/// Lists all messages for a given thread by reading and parsing its messages.jsonl file.
/// Returns a vector of message JSON values.
#[command]
pub async fn list_messages<R: Runtime>(
    app_handle: tauri::AppHandle<R>,
    thread_id: String,
) -> Result<Vec<serde_json::Value>, String> {
    let path = get_messages_path(app_handle, &thread_id);
    if !path.exists() {
        return Ok(vec![]);
    }

    let file = File::open(&path).map_err(|e| {
        eprintln!("Error opening file {}: {}", path.display(), e);
        e.to_string()
    })?;
    let reader = BufReader::new(file);

    let mut messages = Vec::new();
    for line in reader.lines() {
        let line = line.map_err(|e| {
            eprintln!("Error reading line from file {}: {}", path.display(), e);
            e.to_string()
        })?;
        let message: serde_json::Value = serde_json::from_str(&line).map_err(|e| {
            eprintln!(
                "Error parsing JSON from line in file {}: {}",
                path.display(),
                e
            );
            e.to_string()
        })?;
        messages.push(message);
    }

    Ok(messages)
}

/// Appends a new message to a thread's messages.jsonl file.
/// Uses a per-thread async lock to prevent race conditions and ensure file consistency.
#[command]
pub async fn create_message<R: Runtime>(
    app_handle: tauri::AppHandle<R>,
    mut message: serde_json::Value,
) -> Result<serde_json::Value, String> {
    let thread_id = {
        let id = message
            .get("thread_id")
            .and_then(|v| v.as_str())
            .ok_or("Missing thread_id")?;
        id.to_string()
    };
    ensure_thread_dir_exists(app_handle.clone(), &thread_id)?;
    let path = get_messages_path(app_handle.clone(), &thread_id);

    let uuid = Uuid::new_v4().to_string();
    message["id"] = serde_json::Value::String(uuid);

    // Acquire per-thread lock before writing
    {
        let mut locks = MESSAGE_LOCKS.lock().await;
        let lock = locks
            .entry(thread_id.to_string())
            .or_insert_with(|| Arc::new(Mutex::new(())))
            .clone();
        drop(locks); // Release the map lock before awaiting the file lock

        let _guard = lock.lock().await;

        let mut file = fs::OpenOptions::new()
            .create(true)
            .append(true)
            .open(path)
            .map_err(|e| e.to_string())?;

        let data = serde_json::to_string(&message).map_err(|e| e.to_string())?;
        writeln!(file, "{}", data).map_err(|e| e.to_string())?;
    }

    Ok(message)
}

/// Modifies an existing message in a thread's messages.jsonl file.
/// Uses a per-thread async lock to prevent race conditions and ensure file consistency.
/// Rewrites the entire messages.jsonl file for the thread.
#[command]
pub async fn modify_message<R: Runtime>(
    app_handle: tauri::AppHandle<R>,
    message: serde_json::Value,
) -> Result<serde_json::Value, String> {
    let thread_id = message
        .get("thread_id")
        .and_then(|v| v.as_str())
        .ok_or("Missing thread_id")?;
    let message_id = message
        .get("id")
        .and_then(|v| v.as_str())
        .ok_or("Missing message id")?;

    // Acquire per-thread lock before modifying
    {
        let mut locks = MESSAGE_LOCKS.lock().await;
        let lock = locks
            .entry(thread_id.to_string())
            .or_insert_with(|| Arc::new(Mutex::new(())))
            .clone();
        drop(locks); // Release the map lock before awaiting the file lock

        let _guard = lock.lock().await;

        let mut messages = list_messages(app_handle.clone(), thread_id.to_string()).await?;
        if let Some(index) = messages
            .iter()
            .position(|m| m.get("id").and_then(|v| v.as_str()) == Some(message_id))
        {
            messages[index] = message.clone();

            // Rewrite all messages
            let path = get_messages_path(app_handle.clone(), thread_id);
            let mut file = File::create(path).map_err(|e| e.to_string())?;
            for msg in messages {
                let data = serde_json::to_string(&msg).map_err(|e| e.to_string())?;
                writeln!(file, "{}", data).map_err(|e| e.to_string())?;
            }
        }
    }
    Ok(message)
}

/// Deletes a message from a thread's messages.jsonl file by message ID.
/// Rewrites the entire messages.jsonl file for the thread.
#[command]
pub async fn delete_message<R: Runtime>(
    app_handle: tauri::AppHandle<R>,
    thread_id: String,
    message_id: String,
) -> Result<(), String> {
    let mut messages = list_messages(app_handle.clone(), thread_id.clone()).await?;
    messages.retain(|m| m.get("id").and_then(|v| v.as_str()) != Some(message_id.as_str()));

    // Rewrite remaining messages
    let path = get_messages_path(app_handle.clone(), &thread_id);
    let mut file = File::create(path).map_err(|e| e.to_string())?;
    for msg in messages {
        let data = serde_json::to_string(&msg).map_err(|e| e.to_string())?;
        writeln!(file, "{}", data).map_err(|e| e.to_string())?;
    }

    Ok(())
}

/// Retrieves the first assistant associated with a thread.
/// Returns an error if the thread or assistant is not found.
#[command]
pub async fn get_thread_assistant<R: Runtime>(
    app_handle: tauri::AppHandle<R>,
    thread_id: String,
) -> Result<serde_json::Value, String> {
    let path = get_thread_metadata_path(app_handle, &thread_id);
    if !path.exists() {
        return Err("Thread not found".to_string());
    }
    let data = fs::read_to_string(&path).map_err(|e| e.to_string())?;
    let thread: serde_json::Value = serde_json::from_str(&data).map_err(|e| e.to_string())?;
    if let Some(assistants) = thread.get("assistants").and_then(|a| a.as_array()) {
        if let Some(first) = assistants.get(0) {
            Ok(first.clone())
        } else {
            Err("Assistant not found".to_string())
        }
    } else {
        Err("Assistant not found".to_string())
    }
}

/// Adds a new assistant to a thread's metadata.
/// Updates thread.json with the new assistant information.
#[command]
pub async fn create_thread_assistant<R: Runtime>(
    app_handle: tauri::AppHandle<R>,
    thread_id: String,
    assistant: serde_json::Value,
) -> Result<serde_json::Value, String> {
    let path = get_thread_metadata_path(app_handle.clone(), &thread_id);
    if !path.exists() {
        return Err("Thread not found".to_string());
    }
    let mut thread: serde_json::Value = {
        let data = fs::read_to_string(&path).map_err(|e| e.to_string())?;
        serde_json::from_str(&data).map_err(|e| e.to_string())?
    };
    if let Some(assistants) = thread.get_mut("assistants").and_then(|a| a.as_array_mut()) {
        assistants.push(assistant.clone());
    } else {
        thread["assistants"] = serde_json::Value::Array(vec![assistant.clone()]);
    }
    let data = serde_json::to_string_pretty(&thread).map_err(|e| e.to_string())?;
    fs::write(&path, data).map_err(|e| e.to_string())?;
    Ok(assistant)
}

/// Modifies an existing assistant's information in a thread's metadata.
/// Updates thread.json with the modified assistant data.
#[command]
pub async fn modify_thread_assistant<R: Runtime>(
    app_handle: tauri::AppHandle<R>,
    thread_id: String,
    assistant: serde_json::Value,
) -> Result<serde_json::Value, String> {
    let path = get_thread_metadata_path(app_handle.clone(), &thread_id);
    if !path.exists() {
        return Err("Thread not found".to_string());
    }
    let mut thread: serde_json::Value = {
        let data = fs::read_to_string(&path).map_err(|e| e.to_string())?;
        serde_json::from_str(&data).map_err(|e| e.to_string())?
    };
    let assistant_id = assistant
        .get("id")
        .and_then(|v| v.as_str())
        .ok_or("Missing assistant_id")?;
    if let Some(assistants) = thread
        .get_mut("assistants")
        .and_then(|a: &mut serde_json::Value| a.as_array_mut())
    {
        if let Some(index) = assistants
            .iter()
            .position(|a| a.get("id").and_then(|v| v.as_str()) == Some(assistant_id))
        {
            assistants[index] = assistant.clone();
            let data = serde_json::to_string_pretty(&thread).map_err(|e| e.to_string())?;
            fs::write(&path, data).map_err(|e| e.to_string())?;
        }
    }
    Ok(assistant)
}

#[cfg(test)]
mod tests {
    use crate::core::cmd::get_jan_data_folder_path;

    use super::*;
    use serde_json::json;
    use std::fs;
    use std::path::PathBuf;
    use tauri::test::{mock_app, MockRuntime};

    // Helper to create a mock app handle with a temp data dir
    fn mock_app_with_temp_data_dir() -> (tauri::App<MockRuntime>, PathBuf) {
        let app = mock_app();
        let data_dir = get_jan_data_folder_path(app.handle().clone());
        println!("Mock app data dir: {}", data_dir.display());
        // Patch get_data_dir to use temp dir (requires get_data_dir to be overridable or injectable)
        // For now, we assume get_data_dir uses tauri::api::path::app_data_dir(&app_handle)
        // and that we can set the environment variable to redirect it.
        (app, data_dir)
    }

    #[tokio::test]
    async fn test_create_and_list_threads() {
        let (app, data_dir) = mock_app_with_temp_data_dir();
        // Create a thread
        let thread = json!({
            "object": "thread",
            "title": "Test Thread",
            "assistants": [],
            "created": 1234567890,
            "updated": 1234567890,
            "metadata": null
        });
        let created = create_thread(app.handle().clone(), thread.clone())
            .await
            .unwrap();
        assert_eq!(created["title"], "Test Thread");

        // List threads
        let threads = list_threads(app.handle().clone()).await.unwrap();
        assert!(threads.len() > 0);

        // Clean up
        fs::remove_dir_all(data_dir).unwrap();
    }

    #[tokio::test]
    async fn test_create_and_list_messages() {
        let (app, data_dir) = mock_app_with_temp_data_dir();
        // Create a thread first
        let thread = json!({
            "object": "thread",
            "title": "Msg Thread",
            "assistants": [],
            "created": 123,
            "updated": 123,
            "metadata": null
        });
        let created = create_thread(app.handle().clone(), thread.clone())
            .await
            .unwrap();
        let thread_id = created["id"].as_str().unwrap().to_string();

        // Create a message
        let message = json!({
            "object": "message",
            "thread_id": thread_id,
            "assistant_id": null,
            "attachments": null,
            "role": "user",
            "content": [],
            "status": "sent",
            "created_at": 123,
            "completed_at": 123,
            "metadata": null,
            "type_": null,
            "error_code": null,
            "tool_call_id": null
        });
        let created_msg = create_message(app.handle().clone(), message).await.unwrap();
        assert_eq!(created_msg["role"], "user");

        // List messages
        let messages = list_messages(app.handle().clone(), thread_id.clone())
            .await
            .unwrap();
        assert!(messages.len() > 0);
        assert_eq!(messages[0]["role"], "user");

        // Clean up
        fs::remove_dir_all(data_dir).unwrap();
    }

    #[tokio::test]
    async fn test_create_and_get_thread_assistant() {
        let (app, data_dir) = mock_app_with_temp_data_dir();
        // Create a thread
        let thread = json!({
            "object": "thread",
            "title": "Assistant Thread",
            "assistants": [],
            "created": 1,
            "updated": 1,
            "metadata": null
        });
        let created = create_thread(app.handle().clone(), thread.clone())
            .await
            .unwrap();
        let thread_id = created["id"].as_str().unwrap().to_string();

        // Add assistant
        let assistant = json!({
            "id": "assistant-1",
            "assistant_name": "Test Assistant",
            "model": {
                "id": "model-1",
                "name": "Test Model",
                "settings": json!({})
            },
            "instructions": null,
            "tools": null
        });
        let _ = create_thread_assistant(app.handle().clone(), thread_id.clone(), assistant.clone())
            .await
            .unwrap();

        // Get assistant
        let got = get_thread_assistant(app.handle().clone(), thread_id.clone())
            .await
            .unwrap();
        assert_eq!(got["assistant_name"], "Test Assistant");

        // Clean up
        fs::remove_dir_all(data_dir).unwrap();
    }
}