pub mod chat;
pub mod file_system;
pub mod web;

use std::{
    collections::{HashMap, HashSet}, fmt::Display, io::BufRead, path::{Path, PathBuf}, sync::Arc
};

use anyhow::Result;
use chat::ChatMessage;
use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
use tokio::sync::Mutex;

use uuid::Uuid;


use crate::llms::coding::Language;

/**
 * Gather and manage context.
 *
 * Main features:
 * - Start a context gathering session
 * - Trait for context gathering modules to add context to the session
 * - Format the context for display
 * - Save context to a file in binary or json format
 * - Format context for loading into an LLM
 * - Format context for loading into a vector database
 *
 */

#[derive(Serialize, Deserialize, Debug, Clone)]
pub struct Context {
    items: Vec<ContextItem>,
    index: HashMap<Uuid, usize>,   // id -> vector index for fast lookup
    content_hash: HashSet<String>, // set lookup for fast duplicate detection
}

impl Default for Context {
    fn default() -> Self {
        Self::new()
    }
}

impl Context {
    pub fn new() -> Self {
        Self {
            items: vec![],
            index: HashMap::new(),
            content_hash: HashSet::new(),
        }
    }

    pub fn add_item(&mut self, item: ContextItem) -> Uuid {
        let id = item.id;
        // Check for duplicates based on content
        let item_string = item.format_context();
        if self.content_hash.contains(&item_string) {
            return id;
        }

        let index = self.items.len();
        self.index.insert(item.id, index);
        self.content_hash.insert(item_string);
        self.items.push(item);
        id
    }

    pub fn add_tag_to_item(&mut self, id: &Uuid, tag: Tags) {
        if let Some(index) = self.index.get(id) {
            self.items[*index].add_tag(tag);
        }
    }

    pub fn get_item(&self, id: &Uuid) -> Option<&ContextItem> {
        self.index.get(id).map(|index| &self.items[*index])
    }

    pub fn get_items(&self) -> &Vec<ContextItem> {
        &self.items
    }

    pub fn get_index(&self) -> &HashMap<Uuid, usize> {
        &self.index
    }

    pub fn get_item_ids(&self) -> Vec<Uuid> {
        self.index.keys().cloned().collect()
    }

    pub fn get_item_by_index(&self, index: usize) -> Option<&ContextItem> {
        self.items.get(index)
    }

    pub fn get_item_index(&self, id: &Uuid) -> Option<usize> {
        self.index.get(id).copied()
    }

    pub fn delete_item(&mut self, id: &Uuid) -> Result<()> {
        if let Some(index) = self.index.remove(id) {
            self.items.remove(index);
            Ok(())
        } else {
            Err(anyhow::anyhow!("Item not found"))
        }
    }

    pub fn format_context(&self) -> String {
        self.items
            .iter()
            .map(|item| item.format_context())
            .collect::<Vec<String>>()
            .join("\n")
    }

    pub fn get_response(&self, question: &str) -> Option<&String> {
        for item in self.items.iter() {
            if let ContextContent::Text { identifier, content } = &item.get_content().unwrap() {
                if identifier == question {
                    return Some(content);
                }
            }
        }
        None
    }

    pub fn add_response(&mut self, question: &str, answer: String) {
        self.add_item(ContextItem::new(ContextContent::Text {
            identifier: question.to_string(),
            content: answer,
        }));
    }

    pub fn from_toml(path: &Path) -> Result<Self> {
        let content = std::fs::read_to_string(path)?;
        let responses: HashMap<String, String> = toml::from_str(&content)?;
        let mut context = Context::new();
        
        for (question, answer) in responses {
            context.add_item(ContextItem::new(ContextContent::Text {
                identifier: question,
                content: answer,
            }));
        }
        
        Ok(context)
    }

    pub fn from_template(template: &str) -> Result<Self> {
        let template_path = Path::new(template);
        if !template_path.exists() {
            println!("Template file does not exist: {}", template_path.display());
            return Err(anyhow::anyhow!("Template file does not exist"));
        }

        Self::from_toml(template_path)
    }
}

#[derive(Debug, Clone)]
pub struct ContextSession {
    start_time: DateTime<Utc>,
    context: Context,
    shared_context: Arc<std::sync::Mutex<Context>>,
    session_dir: PathBuf,
}

pub fn get_base_context_dir() -> Result<PathBuf> {
    let current_dir = std::env::current_dir()?;
    let session_dir = current_dir.join(".context").join("sessions");
    Ok(session_dir)
}

impl ContextSession {
    pub fn new() -> Result<Self> {
        let timestamp = Utc::now();
        let session_dir = get_base_context_dir()?;
        let session_path =
            session_dir.join(format!("session_{}", timestamp.format("%Y-%m-%d_%H-%M-%S")));

        let session_dir = get_base_context_dir()?;
        if !session_dir.exists() {
            std::fs::create_dir_all(&session_dir)?;
        }

        std::fs::create_dir_all(&session_path)?;

        Ok(Self {
            start_time: timestamp,
            context: Context::new(),
            shared_context: Arc::new(std::sync::Mutex::new(Context::new())),
            session_dir: session_path,
        })
    }

    /// Resume from the last session in the session directory
    pub fn resume() -> Result<Self> {
        let sessions = Self::list_sessions()?;
        let last_session = sessions
            .first()
            .ok_or(anyhow::anyhow!("No sessions found"))?;

        // Create the path to the context.bin file
        let context_file = last_session.0.join("context.bin");

        let context = match Self::load_context(&context_file) {
            Ok(context) => context,
            Err(_) => {
                println!("No context found, starting new session");
                Context::new()
            }
        };

        Ok(Self {
            start_time: last_session.1,
            context: context.clone(),
            shared_context: Arc::new(std::sync::Mutex::new(context)),
            session_dir: last_session.0.clone(),
        })
    }

    pub fn try_resume() -> Result<Self> {
        Self::resume().or_else(|_| Self::new())
    }
    pub fn get_context(&self) -> Context {
        self.merge_shared_context().unwrap()
    }

    /// Add an item to the context
    ///
    /// **Arguments**
    ///
    /// * `item` - The item to add to the context
    ///
    /// **Returns**
    ///
    /// The id of the added item
    fn add_item_to_context(&mut self, item: ContextItem) -> Uuid {
        self.context.add_item(item)
    }

    /// Add a text item to the context
    ///
    /// **Arguments**
    ///
    /// * `identifier` - The identifier of the text
    /// * `content` - The content of the text
    ///
    /// **Returns**
    ///
    /// The id of the added item
    pub fn add_text_to_context(&mut self, identifier: &str, content: &str) -> Uuid {
        self.add_item_to_context(ContextItem::new(ContextContent::Text {
            identifier: identifier.to_string(),
            content: content.to_string(),
        }))
    }

    /// Add a file item to the context
    ///
    /// **Arguments**
    ///
    /// * `path` - The path of the file
    /// * `content` - The content of the file
    ///
    /// **Returns**
    ///
    /// The id of the added item
    pub fn add_file_to_context(&mut self, path: &str, content: &str) -> Uuid {
        self.add_item_to_context(ContextItem::new(ContextContent::File {
            path: path.to_string(),
            content: content.to_string(),
        }))
    }

    /// Add a webpage item to the context
    ///
    /// **Arguments**
    ///
    /// * `url` - The url of the webpage
    /// * `content` - The content of the webpage
    ///
    /// **Returns**
    ///
    /// The id of the added item
    pub fn add_webpage_to_context(&mut self, url: &str, content: &str) -> Uuid {
        self.add_item_to_context(ContextItem::new(ContextContent::Webpage {
            url: url.to_string(),
            content: content.to_string(),
        }))
    }

    /// Add a chat item to the context
    ///
    /// **Arguments**
    ///
    /// * `messages` - The messages of the chat
    ///
    /// **Returns**
    ///
    /// The id of the added item
    pub fn add_chat_to_context(&mut self, messages: Vec<ChatMessage>) -> Uuid {
        self.add_item_to_context(ContextItem::new(ContextContent::Chat { messages }))
    }

    /// Add an api item to the context
    ///
    /// **Arguments**
    ///
    /// * `url` - The url of the api
    /// * `content` - The content of the api
    ///
    /// **Returns**
    ///
    /// The id of the added item
    pub fn add_api_to_context(&mut self, url: &str, content: &str) -> Uuid {
        self.add_item_to_context(ContextItem::new(ContextContent::Api {
            url: url.to_string(),
            content: content.to_string(),
        }))
    }

    /// Add a project path item to the context
    ///
    /// **Arguments**
    ///
    /// * `path` - The path of the project
    ///
    /// **Returns**
    ///
    /// The id of the added item
    pub fn add_project_path_to_context(&mut self, path: &str) -> Uuid {
        self.add_item_to_context(ContextItem::new(ContextContent::ProjectPath {
            path: path.to_string(),
        }))
    }

    /// Add a code item to the context
    ///
    /// **Arguments**
    ///
    /// * `language` - The language of the code
    /// * `descriptor` - The descriptor of the code
    /// * `content` - The content of the code
    ///
    /// **Returns**
    ///
    /// The id of the added item
    pub fn add_code_to_context(
        &mut self,
        language: Language,
        descriptor: CodeDescriptor,
        content: &str,
    ) -> Uuid {
        self.add_item_to_context(ContextItem::new(ContextContent::Code {
            language,
            descriptor,
            content: content.to_string(),
        }))
    }

    /// Add a web context to the session
    ///
    /// **Arguments**
    ///
    /// * `web_context` - The web context to add
    ///
    /// **Returns**
    ///
    /// The id of the added item
    pub fn add_web_context_to_session(&mut self, web_context: web::WebContext) -> Result<()> {
        let static_site_items = web_context.collect_static_site_into_context_items()?;
        let api_items = web_context.collect_api_into_context_items()?;

        for item in static_site_items {
            self.add_item_to_context(item);
        }
        for item in api_items {
            self.add_item_to_context(item);
        }
        Ok(())
    }

    // Add method to get ordered list of sessions
    pub fn list_sessions() -> Result<Vec<(PathBuf, DateTime<Utc>)>> {
        let base_dir = get_base_context_dir()?;

        // Create directory if it doesn't exist
        if !base_dir.exists() {
            return Ok(Vec::new());
        }

        let mut sessions = Vec::new();

        // Read directory entries
        for entry in std::fs::read_dir(base_dir)? {
            let entry = entry?;
            let path = entry.path();

            // Parse the timestamp from directory name
            if let Some(dir_name) = path.file_name().and_then(|n| n.to_str()) {
                if dir_name.starts_with("session_") {
                    // Extract timestamp string (after "session_")
                    let timestamp_str = &dir_name[8..];
                    if let Ok(timestamp) = DateTime::parse_from_str(
                        &format!("{} +0000", timestamp_str.replace('_', " ")),
                        "%Y-%m-%d %H-%M-%S %z",
                    ) {
                        sessions.push((path, timestamp.with_timezone(&Utc)));
                    }
                }
            }
        }

        // Sort by timestamp, newest first
        sessions.sort_by(|a, b| b.1.cmp(&a.1));

        Ok(sessions)
    }

    pub fn get_session_dir(&self) -> &PathBuf {
        &self.session_dir
    }

    pub fn save_context(&self) -> Result<()> {
        let merged_context = self.merge_shared_context()?;
        let context_file = self.session_dir.join("context.bin");
        let encoded: Vec<u8> = bincode::serialize(&merged_context)?;
        std::fs::write(context_file, encoded)?;
        Ok(())
    }

    // Optional: Add a method to load context from file
    pub fn load_context(path: &PathBuf) -> Result<Context> {
        let bytes = std::fs::read(path)?;
        let context: Context = bincode::deserialize(&bytes)?;
        Ok(context)
    }

    // Remove the context folder for a session
    pub fn destroy_context_session_storage(&mut self) -> Result<()> {
        std::fs::remove_dir_all(&self.session_dir)?;
        Ok(())
    }

    /// Merges this context with another context and returns a new combined Context
    /// without modifying either original context
    pub fn merge_context(&self, other: &Context) -> Context {
        let mut combined = Context::new();
        
        // Create vectors of references to items from both contexts
        let mut all_items: Vec<&ContextItem> = Vec::new();
        all_items.extend(self.context.get_items().iter());
        all_items.extend(other.get_items().iter());
        
        // Sort items by timestamp
        all_items.sort_by(|a, b| a.timestamp.cmp(&b.timestamp));
        
        // Add sorted items to combined context
        for item in all_items {
            combined.add_item(item.clone());
        }
        
        combined
    }

    /// Merges the current shared context with the session's context and returns 
    /// the combined context without modifying either
    pub fn merge_shared_context(&self) -> Result<Context> {
        let guard = self.shared_context.lock().map_err(|e| anyhow::anyhow!("{}", e))?;
        let shared = guard.clone();
        drop(guard);
        Ok(self.merge_context(&shared))
    }

    /// Gets a clone of the shared context
    pub fn get_shared_context(&self) -> Result<Context> {
        let guard = self.shared_context.lock().map_err(|e| anyhow::anyhow!("{}", e))?;
        Ok(guard.clone())
    }

    /// Gets a mutable reference to the session's context
    pub fn get_context_mut(&mut self) -> &mut Context {
        &mut self.context
    }

    /// Creates a thread-safe clone of this session that can be used across threads
    pub fn into_thread_safe(self) -> Arc<std::sync::Mutex<Self>> {
        Arc::new(std::sync::Mutex::new(self))
    }

    /// Add text to context in a thread-safe way, using the shared context
    pub fn add_text_to_shared_context(&self, identifier: &str, content: &str) -> Result<()> {
        if let Ok(mut shared) = self.shared_context.lock() {
            shared.add_item(ContextItem::new(ContextContent::Text {
                identifier: identifier.to_string(),
                content: content.to_string(),
            }));
        }
        Ok(())
    }

    /// Spawn a background task that processes lines from a reader and adds them to shared context
    pub fn spawn_line_processor<R>(&self, reader: R, identifier: String) -> Result<()> 
    where
        R: BufRead + Send + 'static,
    {
        let shared_context = self.shared_context.clone();
        std::thread::spawn(move || {
            for line in reader.lines() {
                if let Ok(line) = line {
                    if let Ok(mut context) = shared_context.lock() {
                        context.add_item(ContextItem::new(ContextContent::Text {
                            identifier: identifier.clone(),
                            content: line,
                        }));
                    }
                }
            }
        });
        Ok(())
    }
}

#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, Eq)]
pub enum Tags {
    General,
    Project,
    Code,
    Web,
    File,
}

#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, Eq)]
pub struct ContextItem {
    id: Uuid,
    timestamp: DateTime<Utc>,
    content: ContextContent,
    tags: Vec<Tags>,
}

impl ContextItem {
    pub fn new(content: ContextContent) -> Self {
        Self {
            id: Uuid::new_v4(),
            timestamp: Utc::now(),
            content,
            tags: vec![],
        }
    }

    pub fn add_tag(&mut self, tag: Tags) {
        self.tags.push(tag);
    }

    pub fn format_context(&self) -> String {
        let human_readable_timestamp = self.timestamp.format("%Y-%m-%d %H:%M:%S").to_string();
        match &self.content {
            ContextContent::Text {
                identifier,
                content,
            } => {
                format!(
                    "------\nContext gathered at: {}\n---\n{}\n---\n{}\n------ \n",
                    human_readable_timestamp, identifier, content
                )
            }
            ContextContent::File { path, content } => {
                format!(
                    "------\nContext gathered at: {}\n---\nFrom File: {}\n---\n{}\n------ \n",
                    human_readable_timestamp, path, content
                )
            }
            ContextContent::Webpage { url, content } => {
                format!(
                    "------\nContext gathered at: {}\n---\nFrom Webpage: {}\n---\n{}\n------ \n",
                    human_readable_timestamp, url, content
                )
            }
            ContextContent::Chat { messages } => {
                format!(
                    "------\nContext gathered at: {}\n---\nFrom Chat: {:#?}\n------ \n",
                    human_readable_timestamp, messages
                )
            }
            ContextContent::Api { url, content } => {
                format!(
                    "------\nContext gathered at: {}\n---\nFrom API: {}\n---\n{}\n------ \n",
                    human_readable_timestamp, url, content
                )
            }
            ContextContent::ProjectPath { path } => {
                format!(
                    "------\nContext gathered at: {}\n---\nFrom Project Path: {}\n------ \n",
                    human_readable_timestamp, path
                )
            }
            ContextContent::Code {
                language,
                descriptor,
                content,
            } => {
                format!(
                    "------\nContext gathered at: {}\n---\nFrom Code ({}): {}\n---\n{}\n------ \n",
                    human_readable_timestamp, language, descriptor, content
                )
            }
        }
    }

    pub fn get_content(&self) -> Result<&ContextContent> {
        Ok(&self.content)
    }
}

#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, Eq)]
pub enum CodeDescriptor {
    MooseDataModel,
    MooseIngestionScript,
}

impl Display for CodeDescriptor {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            CodeDescriptor::MooseDataModel => write!(f, "moose-data-model"),
            CodeDescriptor::MooseIngestionScript => write!(f, "moose-ingestion-script"),
        }
    }
}

impl CodeDescriptor {
    fn from_str(s: &str) -> Result<Self> {
        match s {
            "moose-data-model" => Ok(CodeDescriptor::MooseDataModel),
            "moose-ingestion-script" => Ok(CodeDescriptor::MooseIngestionScript),
            _ => Err(anyhow::anyhow!("Unsupported code descriptor: {}", s)),
        }
    }
}

#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, Eq)]
pub enum ContextContent {
    Text {
        identifier: String,
        content: String,
    },
    File {
        path: String,
        content: String,
    },
    Webpage {
        url: String,
        content: String,
    },
    Chat {
        messages: Vec<ChatMessage>,
    },
    Api {
        url: String,
        content: String,
    },
    ProjectPath {
        path: String,
    },
    // Add more variants as needed
    Code {
        language: Language,
        descriptor: CodeDescriptor,
        content: String,
    },
    // Url { url: String },
    // etc.
}

#[cfg(test)]
mod tests {
    use chat::ChatMessageModel;

    use crate::llms::open_ai::messages::OpenAIMessage;

    use super::*;

    #[test]
    fn test_text_format_context() {
        let item = ContextItem::new(ContextContent::Text {
            identifier: "identifier".to_string(),
            content: "content".to_string(),
        });
        let formatted = item.format_context();
        // Just verify the structure and content, not the exact timestamp
        assert!(formatted.contains("Context gathered at: "));
        assert!(formatted.contains("content"));
        assert!(formatted.starts_with("------\n"));
        assert!(formatted.ends_with("------ \n"));
    }

    #[test]
    fn test_file_format_context() {
        let item = ContextItem::new(ContextContent::File {
            path: "path".to_string(),
            content: "content".to_string(),
        });
        let formatted = item.format_context();

        // Just verify the structure and content, not the exact timestamp
        assert!(formatted.contains("Context gathered at: "));
        assert!(formatted.contains("From File: path"));
        assert!(formatted.contains("content"));
        assert!(formatted.starts_with("------\n"));
        assert!(formatted.ends_with("------ \n"));
    }

    #[test]
    fn test_webpage_format_context() {
        let item = ContextItem::new(ContextContent::Webpage {
            url: "url".to_string(),
            content: "content".to_string(),
        });
        let formatted = item.format_context();

        // Just verify the structure and content, not the exact timestamp
        assert!(formatted.contains("Context gathered at: "));
        assert!(formatted.contains("From Webpage: url"));
        assert!(formatted.contains("content"));
        assert!(formatted.starts_with("------\n"));
        assert!(formatted.ends_with("------ \n"));
    }

    #[test]
    fn test_api_format_context() {
        let item = ContextItem::new(ContextContent::Api {
            url: "url".to_string(),
            content: "content".to_string(),
        });
        let formatted = item.format_context();

        // Just verify the structure and content, not the exact timestamp
        assert!(formatted.contains("Context gathered at: "));
        assert!(formatted.contains("From API: url"));
        assert!(formatted.contains("content"));
        assert!(formatted.starts_with("------\n"));
        assert!(formatted.ends_with("------ \n"));
    }

    #[test]
    fn test_chat_format_context() {
        let open_ai_message = OpenAIMessage::User {
            content: "content".to_string(),
        };

        let item = ContextItem::new(ContextContent::Chat {
            messages: vec![ChatMessage::from_open_ai_message(
                open_ai_message,
                ChatMessageModel::GPT4o,
            )],
        });
        let formatted = item.format_context();

        // Just verify the structure and content, not the exact timestamp
        assert!(formatted.contains("Context gathered at: "));
        assert!(formatted.contains("From Chat: "));
        assert!(formatted.contains("content"));
        assert!(formatted.starts_with("------\n"));
        assert!(formatted.ends_with("------ \n"));
    }

    #[test]
    fn test_context_session_initialization() {
        let session = ContextSession::new().unwrap();
        assert_eq!(session.context.items.len(), 0);
        assert!(session.get_session_dir().exists());
    }

    #[test]
    fn test_add_item() {
        let mut session = ContextSession::new().unwrap();
        let item = ContextItem::new(ContextContent::Text {
            identifier: "identifier".to_string(),
            content: "content".to_string(),
        });
        session.context.add_item(item.clone());
        assert_eq!(session.context.items.len(), 1);
        assert_eq!(session.context.get_item_index(&item.id), Some(0));
    }

    #[test]
    fn test_get_item_with_wrong_id() {
        let mut session = ContextSession::new().unwrap();
        let item = ContextItem::new(ContextContent::Text {
            identifier: "identifier".to_string(),
            content: "content".to_string(),
        });
        session.context.add_item(item.clone());
        assert_eq!(session.context.get_item_index(&Uuid::new_v4()), None);
    }

    #[test]
    fn test_add_and_remove_item() {
        let mut session = ContextSession::new().unwrap();
        let item = ContextItem::new(ContextContent::Text {
            identifier: "identifier".to_string(),
            content: "content".to_string(),
        });
        session.context.add_item(item.clone());
        assert_eq!(session.context.items.len(), 1);
        session.context.delete_item(&item.id).unwrap();
        assert_eq!(session.context.items.len(), 0);
        assert_eq!(session.context.get_item_index(&item.id), None);
    }

    #[test]
    fn test_delete_on_empty_context() {
        let mut session = ContextSession::new().unwrap();
        let item = ContextItem::new(ContextContent::Text {
            identifier: "identifier".to_string(),
            content: "content".to_string(),
        });
        session.context.add_item(item.clone());
        assert!(session.context.delete_item(&item.id).is_ok());
    }
}
