use std::any::Any;
use std::collections::HashMap;
use std::path::{Path, PathBuf};
use std::sync::Arc;

use async_trait::async_trait;
use tokio::sync::Mutex;

use crate::{Message, Result};

/// Metadata header written at the top of a persisted session file.
///
/// Stores the session ID, creation timestamp, agent-loop version, and
/// message count so that a [`SessionStore`] can verify compatibility
/// when loading a previously saved session.
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct SessionMeta {
    pub session_id: String,
    pub created_at_ms: u64,
    pub agent_loop_version: String,
    pub message_count: usize,
}

/// Pluggable persistence backend for [`AgentSession`](crate::AgentSession) history.
///
/// Implementations must be `Send + Sync` so they can be shared across
/// async tasks.  Two built-in implementations are provided:
///
/// * [`MemorySessionStore`] — in-process, for tests and short-lived sessions.
/// * [`FileSessionStore`] — JSONL-based, for durable persistence across restarts.
#[async_trait]
pub trait SessionStore: Send + Sync {
    async fn append(&self, session_id: &str, message: &Message) -> Result<()>;
    async fn flush(&self, session_id: &str) -> Result<()>;
    async fn load(&self, session_id: &str) -> Result<Vec<Message>>;
    async fn delete(&self, session_id: &str) -> Result<()>;
    async fn list(&self) -> Result<Vec<String>>;
    fn as_any(&self) -> &dyn Any;
}

#[derive(Default)]
pub struct MemorySessionStore {
    sessions: Arc<Mutex<HashMap<String, Vec<Message>>>>,
}

impl MemorySessionStore {
    pub fn new() -> Self {
        Self::default()
    }
}

#[async_trait]
impl SessionStore for MemorySessionStore {
    async fn append(&self, session_id: &str, message: &Message) -> Result<()> {
        let mut sessions = self.sessions.lock().await;
        sessions
            .entry(session_id.to_string())
            .or_default()
            .push(message.clone());
        Ok(())
    }

    async fn flush(&self, _session_id: &str) -> Result<()> {
        Ok(())
    }

    async fn load(&self, session_id: &str) -> Result<Vec<Message>> {
        let sessions = self.sessions.lock().await;
        Ok(sessions.get(session_id).cloned().unwrap_or_default())
    }

    async fn delete(&self, session_id: &str) -> Result<()> {
        let mut sessions = self.sessions.lock().await;
        sessions.remove(session_id);
        Ok(())
    }

    async fn list(&self) -> Result<Vec<String>> {
        let sessions = self.sessions.lock().await;
        let mut ids: Vec<String> = sessions.keys().cloned().collect();
        ids.sort();
        Ok(ids)
    }

    fn as_any(&self) -> &dyn Any {
        self
    }
}

#[derive(Default)]
struct FileStoreState {
    buffers: HashMap<String, Vec<Message>>,
}

pub struct FileSessionStore {
    base_dir: PathBuf,
    state: Arc<Mutex<FileStoreState>>,
    flush_task: Option<tokio::task::JoinHandle<()>>,
}

impl FileSessionStore {
    /// Create a file-backed session store.
    ///
    /// If a Tokio runtime is available, a background flush task is started
    /// (every 5s). Without a runtime, automatic periodic flushing is disabled;
    /// callers should invoke `flush()` explicitly for durability.
    pub fn new(base_dir: impl Into<PathBuf>) -> Self {
        let base_dir = base_dir.into();
        let state = Arc::new(Mutex::new(FileStoreState::default()));
        let state_clone = state.clone();
        let base_dir_clone = base_dir.clone();

        let flush_task = tokio::runtime::Handle::try_current().ok().map(|handle| {
            handle.spawn(async move {
                let mut interval = tokio::time::interval(std::time::Duration::from_secs(5));
                loop {
                    interval.tick().await;
                    let session_ids = {
                        let state = state_clone.lock().await;
                        state
                            .buffers
                            .iter()
                            .filter(|(_, messages)| !messages.is_empty())
                            .map(|(id, _)| id.clone())
                            .collect::<Vec<_>>()
                    };
                    for session_id in session_ids {
                        let _ = flush_session(&base_dir_clone, &state_clone, &session_id).await;
                    }
                }
            })
        });

        Self {
            base_dir,
            state,
            flush_task,
        }
    }

    fn session_path(&self, session_id: &str) -> PathBuf {
        self.base_dir.join(format!("{session_id}.jsonl"))
    }

    pub async fn load_meta(&self, session_id: &str) -> Result<Option<SessionMeta>> {
        let path = self.session_path(session_id);
        load_meta_from_path(&path).await
    }
}

impl Drop for FileSessionStore {
    fn drop(&mut self) {
        if let Some(task) = &self.flush_task {
            task.abort();
        }
    }
}

#[async_trait]
impl SessionStore for FileSessionStore {
    async fn append(&self, session_id: &str, message: &Message) -> Result<()> {
        let should_flush = {
            let mut state = self.state.lock().await;
            let buffer = state.buffers.entry(session_id.to_string()).or_default();
            buffer.push(message.clone());
            buffer.len() >= 20
        };

        if should_flush {
            self.flush(session_id).await?;
        }

        Ok(())
    }

    async fn flush(&self, session_id: &str) -> Result<()> {
        flush_session(&self.base_dir, &self.state, session_id).await
    }

    async fn load(&self, session_id: &str) -> Result<Vec<Message>> {
        let path = self.session_path(session_id);
        load_messages_from_path(&path).await
    }

    async fn delete(&self, session_id: &str) -> Result<()> {
        {
            let mut state = self.state.lock().await;
            state.buffers.remove(session_id);
        }

        let path = self.session_path(session_id);
        match tokio::fs::remove_file(path).await {
            Ok(()) => Ok(()),
            Err(err) if err.kind() == std::io::ErrorKind::NotFound => Ok(()),
            Err(err) => Err(err.into()),
        }
    }

    async fn list(&self) -> Result<Vec<String>> {
        if !self.base_dir.exists() {
            return Ok(Vec::new());
        }

        let mut out = Vec::new();
        let mut entries = tokio::fs::read_dir(&self.base_dir).await?;
        while let Some(entry) = entries.next_entry().await? {
            let path = entry.path();
            if path.extension().and_then(|ext| ext.to_str()) != Some("jsonl") {
                continue;
            }
            if let Some(stem) = path.file_stem().and_then(|stem| stem.to_str()) {
                out.push(stem.to_string());
            }
        }
        out.sort();
        Ok(out)
    }

    fn as_any(&self) -> &dyn Any {
        self
    }
}

async fn flush_session(
    base_dir: &Path,
    state: &Arc<Mutex<FileStoreState>>,
    session_id: &str,
) -> Result<()> {
    let pending = {
        let mut state = state.lock().await;
        state.buffers.remove(session_id).unwrap_or_default()
    };

    if pending.is_empty() {
        return Ok(());
    }

    tokio::fs::create_dir_all(base_dir).await?;
    let path = base_dir.join(format!("{session_id}.jsonl"));
    let mut existing = load_messages_from_path(&path).await?;
    existing.extend(pending);

    let meta = SessionMeta {
        session_id: session_id.to_string(),
        created_at_ms: now_ms(),
        agent_loop_version: env!("CARGO_PKG_VERSION").to_string(),
        message_count: existing.len(),
    };

    let mut lines = Vec::with_capacity(existing.len() + 1);
    lines.push(serde_json::to_string(&meta)?);
    for message in existing {
        lines.push(serde_json::to_string(&message)?);
    }

    let payload = format!("{}\n", lines.join("\n"));
    tokio::fs::write(path, payload).await?;
    Ok(())
}

async fn load_meta_from_path(path: &Path) -> Result<Option<SessionMeta>> {
    if !path.exists() {
        return Ok(None);
    }
    let content = tokio::fs::read_to_string(path).await?;
    let mut lines = content.lines();
    let Some(first) = lines.next() else {
        return Ok(None);
    };
    let meta = serde_json::from_str(first)?;
    Ok(Some(meta))
}

async fn load_messages_from_path(path: &Path) -> Result<Vec<Message>> {
    if !path.exists() {
        return Ok(Vec::new());
    }

    let content = tokio::fs::read_to_string(path).await?;
    let mut messages = Vec::new();

    for (idx, line) in content.lines().enumerate() {
        if line.trim().is_empty() {
            continue;
        }
        if idx == 0 && serde_json::from_str::<SessionMeta>(line).is_ok() {
            continue;
        }
        messages.push(serde_json::from_str::<Message>(line)?);
    }

    Ok(messages)
}

fn now_ms() -> u64 {
    std::time::SystemTime::now()
        .duration_since(std::time::UNIX_EPOCH)
        .map(|d| d.as_millis() as u64)
        .unwrap_or_default()
}
