use anyhow::{Context, Result};
use chrono::Utc;
use serde::{Deserialize, Serialize};
use std::collections::BTreeMap;
use std::path::{Path, PathBuf};

use crate::canonical::{LocalSessionState, SessionLocator, WorkspaceSessionState};

use super::{atomic_write, session_overrides};

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SessionStateStore {
    #[serde(default = "current_version")]
    pub version: u32,
    #[serde(default)]
    pub sessions: BTreeMap<String, BTreeMap<String, LocalSessionState>>,
}

#[derive(Debug, Clone, Serialize, Deserialize, Default, PartialEq, Eq)]
pub struct ResolvedLocalSessionState {
    #[serde(skip_serializing_if = "Option::is_none")]
    pub display_title: Option<String>,
    #[serde(default)]
    pub archived: bool,
    #[serde(default)]
    pub hidden: bool,
    #[serde(default)]
    pub pinned: bool,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub notes: Option<String>,
    #[serde(default, skip_serializing_if = "Vec::is_empty")]
    pub tags: Vec<String>,
    #[serde(default, skip_serializing_if = "Vec::is_empty")]
    pub preferred_targets: Vec<String>,
}

#[derive(Debug, Clone, Serialize, Deserialize, Default, PartialEq, Eq)]
pub struct SessionLocalStateUpdate {
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub display_title: Option<Option<String>>,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub hidden: Option<bool>,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub pinned: Option<bool>,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub notes: Option<Option<String>>,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub tags: Option<Vec<String>>,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub preferred_targets: Option<Vec<String>>,
}

fn current_version() -> u32 {
    2
}

impl Default for SessionStateStore {
    fn default() -> Self {
        Self {
            version: current_version(),
            sessions: BTreeMap::new(),
        }
    }
}

pub fn session_state_path() -> Result<PathBuf> {
    let home = dirs::home_dir().context("Unable to locate user home directory")?;
    Ok(home.join(".memorph").join("session_state.json"))
}

pub fn load_state_store() -> Result<SessionStateStore> {
    let state_path = session_state_path()?;
    let legacy_path = session_overrides::overrides_path()?;
    let (store, migrated) = load_state_store_from_paths(&state_path, Some(&legacy_path))?;
    if migrated {
        save_state_store_to_path(&state_path, &store)?;
    }
    Ok(store)
}

pub fn save_state_store(store: &SessionStateStore) -> Result<()> {
    save_state_store_to_path(&session_state_path()?, store)
}

pub fn get_session_state<'a>(
    store: &'a SessionStateStore,
    provider_id: &str,
    session_id: &str,
) -> Option<&'a LocalSessionState> {
    store
        .sessions
        .get(provider_id)
        .and_then(|sessions| sessions.get(session_id))
}

pub fn resolve_session_state(
    store: &SessionStateStore,
    provider_id: &str,
    session_id: &str,
    workspace_dir: Option<&str>,
) -> ResolvedLocalSessionState {
    let Some(state) = get_session_state(store, provider_id, session_id) else {
        return ResolvedLocalSessionState::default();
    };

    let mut resolved = ResolvedLocalSessionState {
        display_title: state
            .display_title
            .as_deref()
            .map(str::trim)
            .filter(|title| !title.is_empty())
            .map(str::to_string),
        archived: state.archived,
        hidden: state.hidden,
        pinned: state.pinned,
        notes: state.notes.clone(),
        tags: state.tags.clone(),
        preferred_targets: state.preferred_targets.clone(),
    };

    if let Some(workspace_dir) = workspace_dir {
        if let Some(workspace_state) = state
            .workspace_overrides
            .iter()
            .find(|entry| entry.workspace_dir == workspace_dir)
        {
            if let Some(hidden) = workspace_state.hidden {
                resolved.hidden = hidden;
            }
            if let Some(pinned) = workspace_state.pinned {
                resolved.pinned = pinned;
            }
            if let Some(preferred_targets) = workspace_state.preferred_targets.as_ref() {
                resolved.preferred_targets = preferred_targets.clone();
            }
        }
    }

    resolved
}

pub fn set_session_state(state: LocalSessionState) -> Result<()> {
    let mut store = load_state_store()?;
    set_session_state_in_store(&mut store, state);
    save_state_store(&store)
}

pub fn set_display_title(provider_id: &str, session_id: &str, title: &str) -> Result<()> {
    let mut store = load_state_store()?;
    set_display_title_in_store(&mut store, provider_id, session_id, title);
    save_state_store(&store)
}

pub fn update_session_state(
    provider_id: &str,
    session_id: &str,
    update: &SessionLocalStateUpdate,
) -> Result<ResolvedLocalSessionState> {
    let mut store = load_state_store()?;
    update_session_state_in_store(&mut store, provider_id, session_id, update);
    let resolved = resolve_session_state(&store, provider_id, session_id, None);
    save_state_store(&store)?;
    Ok(resolved)
}

pub fn remove_session(provider_id: &str, session_id: &str) -> Result<()> {
    let mut store = load_state_store()?;
    remove_session_in_store(&mut store, provider_id, session_id);
    save_state_store(&store)
}

pub fn set_session_state_in_store(store: &mut SessionStateStore, mut state: LocalSessionState) {
    state.updated_at = Utc::now();
    store.version = current_version();
    store
        .sessions
        .entry(state.locator.provider_id.clone())
        .or_default()
        .insert(state.locator.session_id.clone(), state);
}

pub fn set_display_title_in_store(
    store: &mut SessionStateStore,
    provider_id: &str,
    session_id: &str,
    title: &str,
) {
    let title = title.trim();
    if title.is_empty() {
        return;
    }

    let mut state = get_session_state(store, provider_id, session_id)
        .cloned()
        .unwrap_or_else(|| LocalSessionState {
            locator: SessionLocator {
                provider_id: provider_id.to_string(),
                session_id: session_id.to_string(),
            },
            display_title: None,
            archived: false,
            hidden: false,
            pinned: false,
            notes: None,
            tags: Vec::new(),
            preferred_targets: Vec::new(),
            workspace_overrides: Vec::new(),
            updated_at: Utc::now(),
        });
    state.display_title = Some(title.to_string());
    set_session_state_in_store(store, state);
}

pub fn update_session_state_in_store(
    store: &mut SessionStateStore,
    provider_id: &str,
    session_id: &str,
    update: &SessionLocalStateUpdate,
) {
    let mut state = get_session_state(store, provider_id, session_id)
        .cloned()
        .unwrap_or_else(|| LocalSessionState {
            locator: SessionLocator {
                provider_id: provider_id.to_string(),
                session_id: session_id.to_string(),
            },
            display_title: None,
            archived: false,
            hidden: false,
            pinned: false,
            notes: None,
            tags: Vec::new(),
            preferred_targets: Vec::new(),
            workspace_overrides: Vec::new(),
            updated_at: Utc::now(),
        });

    if let Some(display_title) = &update.display_title {
        state.display_title = display_title
            .as_deref()
            .map(str::trim)
            .filter(|value| !value.is_empty())
            .map(str::to_string);
    }
    if let Some(hidden) = update.hidden {
        state.hidden = hidden;
    }
    if let Some(pinned) = update.pinned {
        state.pinned = pinned;
    }
    if let Some(notes) = &update.notes {
        state.notes = notes
            .as_deref()
            .map(str::trim)
            .filter(|value| !value.is_empty())
            .map(str::to_string);
    }
    if let Some(tags) = &update.tags {
        state.tags = normalize_tags(tags);
    }
    if let Some(preferred_targets) = &update.preferred_targets {
        state.preferred_targets = crate::config::normalize_provider_ids(preferred_targets.clone());
    }

    set_session_state_in_store(store, state);
}

pub fn remove_session_in_store(store: &mut SessionStateStore, provider_id: &str, session_id: &str) {
    let Some(sessions) = store.sessions.get_mut(provider_id) else {
        return;
    };
    sessions.remove(session_id);
    if sessions.is_empty() {
        store.sessions.remove(provider_id);
    }
}

fn load_state_store_from_paths(
    state_path: &Path,
    legacy_overrides_path: Option<&Path>,
) -> Result<(SessionStateStore, bool)> {
    if state_path.exists() {
        let raw = std::fs::read_to_string(state_path)
            .with_context(|| format!("Failed to read session state: {}", state_path.display()))?;
        let store = serde_json::from_str(&raw)
            .with_context(|| format!("Failed to parse session state: {}", state_path.display()))?;
        return Ok((store, false));
    }

    if let Some(legacy_overrides_path) = legacy_overrides_path {
        if legacy_overrides_path.exists() {
            return migrate_legacy_overrides(legacy_overrides_path).map(|store| (store, true));
        }
    }

    Ok((SessionStateStore::default(), false))
}

fn save_state_store_to_path(path: &Path, store: &SessionStateStore) -> Result<()> {
    let dir = path
        .parent()
        .with_context(|| format!("Session state path has no parent: {}", path.display()))?;
    std::fs::create_dir_all(dir).with_context(|| {
        format!(
            "Failed to create session state directory: {}",
            dir.display()
        )
    })?;
    let raw = serde_json::to_string_pretty(store)?;
    atomic_write::write_string_atomic(path, &raw)
        .with_context(|| format!("Failed to write session state: {}", path.display()))
}

fn migrate_legacy_overrides(path: &Path) -> Result<SessionStateStore> {
    let raw = std::fs::read_to_string(path)
        .with_context(|| format!("Failed to read session overrides: {}", path.display()))?;
    let legacy: session_overrides::SessionOverrides = serde_json::from_str(&raw)
        .with_context(|| format!("Failed to parse session overrides: {}", path.display()))?;

    let mut store = SessionStateStore {
        version: current_version(),
        sessions: BTreeMap::new(),
    };
    for (provider_id, sessions) in legacy.sessions {
        let provider_bucket = store.sessions.entry(provider_id.clone()).or_default();
        for (session_id, session) in sessions {
            provider_bucket.insert(
                session_id.clone(),
                LocalSessionState {
                    locator: SessionLocator {
                        provider_id: provider_id.clone(),
                        session_id,
                    },
                    display_title: session.display_title,
                    archived: false,
                    hidden: false,
                    pinned: false,
                    notes: None,
                    tags: Vec::new(),
                    preferred_targets: Vec::new(),
                    workspace_overrides: Vec::<WorkspaceSessionState>::new(),
                    updated_at: chrono::DateTime::from_timestamp_millis(session.updated_at)
                        .unwrap_or_else(Utc::now),
                },
            );
        }
    }

    Ok(store)
}

fn normalize_tags(tags: &[String]) -> Vec<String> {
    let mut normalized = Vec::new();
    for tag in tags {
        let tag = tag.trim();
        if tag.is_empty() || normalized.iter().any(|existing| existing == tag) {
            continue;
        }
        normalized.push(tag.to_string());
    }
    normalized
}

#[cfg(test)]
mod tests {
    use super::*;
    use tempfile::tempdir;

    #[test]
    fn resolves_workspace_specific_overrides() {
        let mut store = SessionStateStore::default();
        set_session_state_in_store(
            &mut store,
            LocalSessionState {
                locator: SessionLocator {
                    provider_id: "codex".to_string(),
                    session_id: "abc".to_string(),
                },
                display_title: Some("Renamed".to_string()),
                archived: false,
                hidden: false,
                pinned: true,
                notes: Some("note".to_string()),
                tags: vec!["tag".to_string()],
                preferred_targets: vec!["claude".to_string()],
                workspace_overrides: vec![WorkspaceSessionState {
                    workspace_dir: "/tmp/project".to_string(),
                    hidden: Some(true),
                    pinned: Some(false),
                    preferred_targets: Some(vec!["codex".to_string(), "kiro".to_string()]),
                }],
                updated_at: Utc::now(),
            },
        );

        let resolved = resolve_session_state(&store, "codex", "abc", Some("/tmp/project"));

        assert_eq!(resolved.display_title.as_deref(), Some("Renamed"));
        assert!(resolved.hidden);
        assert!(!resolved.pinned);
        assert_eq!(resolved.preferred_targets, vec!["codex", "kiro"]);
        assert_eq!(resolved.notes.as_deref(), Some("note"));
        assert_eq!(resolved.tags, vec!["tag"]);
    }

    #[test]
    fn migrates_legacy_display_titles() {
        let dir = tempdir().unwrap();
        let state_path = dir.path().join("session_state.json");
        let legacy_path = dir.path().join("session_overrides.json");
        std::fs::write(
            &legacy_path,
            serde_json::json!({
                "version": 1,
                "sessions": {
                    "codex": {
                        "abc": {
                            "display_title": "Renamed",
                            "updated_at": 1_700_000_000_000_i64
                        }
                    }
                }
            })
            .to_string(),
        )
        .unwrap();

        let (store, migrated) =
            load_state_store_from_paths(&state_path, Some(&legacy_path)).unwrap();
        let resolved = resolve_session_state(&store, "codex", "abc", None);

        assert!(migrated);
        assert_eq!(resolved.display_title.as_deref(), Some("Renamed"));
    }

    #[test]
    fn removes_empty_provider_bucket_after_session_cleanup() {
        let mut store = SessionStateStore::default();
        set_display_title_in_store(&mut store, "codex", "abc", "Renamed");

        remove_session_in_store(&mut store, "codex", "abc");

        assert!(resolve_session_state(&store, "codex", "abc", None)
            .display_title
            .is_none());
        assert!(!store.sessions.contains_key("codex"));
    }

    #[test]
    fn applies_local_state_update_fields() {
        let mut store = SessionStateStore::default();
        update_session_state_in_store(
            &mut store,
            "codex",
            "abc",
            &SessionLocalStateUpdate {
                display_title: Some(Some(" Renamed ".to_string())),
                hidden: Some(true),
                pinned: Some(true),
                notes: Some(Some(" note ".to_string())),
                tags: Some(vec!["a".to_string(), "a".to_string(), " ".to_string()]),
                preferred_targets: Some(vec![
                    "claude".to_string(),
                    "claude".to_string(),
                    "missing".to_string(),
                ]),
            },
        );

        let resolved = resolve_session_state(&store, "codex", "abc", None);
        assert_eq!(resolved.display_title.as_deref(), Some("Renamed"));
        assert!(resolved.hidden);
        assert!(resolved.pinned);
        assert_eq!(resolved.notes.as_deref(), Some("note"));
        assert_eq!(resolved.tags, vec!["a"]);
        assert_eq!(resolved.preferred_targets, vec!["claude"]);
    }
}
