use egui_nav::ReturnType;
use enostr::{Filter, NoteId, Pubkey};
use hashbrown::HashMap;
use nostrdb::{Ndb, Subscription};
use notedeck::{
    Accounts, FullHistoryConfig, ScopedSubApi, ScopedSubIdentity, SubConfig, SubKey, SubOwnerKey,
};

use crate::scoped_sub_owner_keys::thread_scope_owner_key;
use crate::timeline::{
    sub::{ndb_sub, ndb_unsub},
    ThreadSelection,
};

type RootNoteId = NoteId;

// column id
type MetaId = usize;

/// Outcome of removing local thread subscriptions for a close action.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
enum UnsubscribeOutcome {
    /// Local NDB sub(s) were removed, but the scope still has stack entries so the
    /// remote scoped-sub owner should remain.
    KeepOwner,
    /// The thread scope was fully removed and the remote scoped-sub owner should
    /// be released using the returned root note id.
    DropOwner(RootNoteId),
}

/// Thread subscription manager keyed by account and column scope.
///
/// Each opened thread scope installs one local NostrDB sub plus one scoped
/// remote sub owner. Closing scope releases owner and tears down local state.
#[derive(Default)]
pub struct ThreadSubs {
    /// Per-account thread subscription bookkeeping.
    by_account: HashMap<Pubkey, AccountThreadSubs>,
}

#[derive(Default)]
struct AccountThreadSubs {
    scopes: HashMap<MetaId, Vec<Scope>>,
}

struct Scope {
    root_id: NoteId,
    stack: Vec<Sub>,
}

struct Sub {
    _selected_id: NoteId,
    sub: Subscription,
    // Keep local filters alive for the full subscription lifetime. Thread
    // filters use custom callbacks and can crash if dropped early.
    _filters: Vec<Filter>,
}

impl ThreadSubs {
    #[allow(clippy::too_many_arguments)]
    pub fn subscribe(
        &mut self,
        ndb: &mut Ndb,
        scoped_subs: &mut ScopedSubApi<'_, '_>,
        meta_id: usize,
        id: &ThreadSelection,
        local_sub_filter: Vec<Filter>,
        new_scope: bool,
        remote_sub_filter: Vec<Filter>,
        remote_history_filter: Vec<Filter>,
    ) {
        let account_pk = scoped_subs.selected_account_pubkey();
        let account_subs = self.by_account.entry(account_pk).or_default();
        let cur_scopes = account_subs.scopes.entry(meta_id).or_default();
        let added_local = if new_scope || cur_scopes.is_empty() {
            local_sub_new_scope(
                ndb,
                scoped_subs,
                account_pk,
                meta_id,
                id,
                local_sub_filter,
                remote_sub_filter,
                remote_history_filter,
                cur_scopes,
            )
        } else {
            let cur_scope = cur_scopes.last_mut().expect("can't be empty");
            sub_current_scope(ndb, id, local_sub_filter, cur_scope)
        };

        if added_local {
            tracing::debug!(
                "Sub stats: account={:?}, num locals: {}",
                account_pk,
                account_subs.scopes.len(),
            );
        }
    }

    pub fn unsubscribe(
        &mut self,
        ndb: &mut Ndb,
        scoped_subs: &mut ScopedSubApi<'_, '_>,
        meta_id: usize,
        id: &ThreadSelection,
        return_type: ReturnType,
    ) {
        let account_pk = scoped_subs.selected_account_pubkey();
        let (owner_to_drop, remove_account_entry) = {
            let Some(account_subs) = self.by_account.get_mut(&account_pk) else {
                return;
            };

            let Some(scopes) = account_subs.scopes.get_mut(&meta_id) else {
                return;
            };

            let scope_depth = scopes.len().saturating_sub(1);
            let Some(unsub_outcome) = (match return_type {
                ReturnType::Drag => unsubscribe_drag(scopes, ndb, id),
                ReturnType::Click => unsubscribe_click(scopes, ndb, id),
            }) else {
                return;
            };

            if scopes.is_empty() {
                account_subs.scopes.remove(&meta_id);
            }

            tracing::debug!(
                "unsub stats: account={:?}, num locals: {}, released owner: {}",
                account_pk,
                account_subs.scopes.len(),
                matches!(unsub_outcome, UnsubscribeOutcome::DropOwner(_)),
            );

            (
                match unsub_outcome {
                    UnsubscribeOutcome::KeepOwner => None,
                    UnsubscribeOutcome::DropOwner(root_id) => Some(thread_scope_owner_key(
                        account_pk,
                        meta_id,
                        &root_id,
                        scope_depth,
                    )),
                },
                account_subs.scopes.is_empty(),
            )
        };

        if remove_account_entry {
            self.by_account.remove(&account_pk);
        }

        if let Some(owner) = owner_to_drop {
            let _ = scoped_subs.drop_owner(owner);
        }
    }

    pub fn get_local(&self, account_pk: &Pubkey, meta_id: usize) -> Option<&Subscription> {
        self.by_account
            .get(account_pk)?
            .scopes
            .get(&meta_id)
            .and_then(|s| s.last())
            .and_then(|s| s.stack.last())
            .map(|s| &s.sub)
    }

    pub fn get_local_for_selected<'a>(
        &'a self,
        accounts: &Accounts,
        meta_id: usize,
    ) -> Option<&'a Subscription> {
        self.get_local(accounts.selected_account_pubkey(), meta_id)
    }
}

fn unsubscribe_drag(
    scopes: &mut Vec<Scope>,
    ndb: &mut Ndb,
    id: &ThreadSelection,
) -> Option<UnsubscribeOutcome> {
    let Some(scope) = scopes.last_mut() else {
        tracing::error!("called drag unsubscribe but there aren't any scopes left");
        return None;
    };

    let Some(cur_sub) = scope.stack.pop() else {
        tracing::error!("expected a scope to be left");
        return None;
    };

    log_scope_root_mismatch(scope, id);

    if !ndb_unsub(ndb, cur_sub.sub, id) {
        // Keep local bookkeeping aligned with NDB when unsubscribe fails.
        scope.stack.push(cur_sub);
        return None;
    }

    if scope.stack.is_empty() {
        let removed_scope = scopes.pop().expect("checked empty above");
        return Some(UnsubscribeOutcome::DropOwner(removed_scope.root_id));
    }

    Some(UnsubscribeOutcome::KeepOwner)
}

fn unsubscribe_click(
    scopes: &mut Vec<Scope>,
    ndb: &mut Ndb,
    id: &ThreadSelection,
) -> Option<UnsubscribeOutcome> {
    let Some(mut scope) = scopes.pop() else {
        tracing::error!("called unsubscribe but there aren't any scopes left");
        return None;
    };

    log_scope_root_mismatch(&scope, id);
    while let Some(sub) = scope.stack.pop() {
        if ndb_unsub(ndb, sub.sub, id) {
            continue;
        }

        // Partial rollback: restore the failed local sub (and any remaining ones)
        // to thread bookkeeping and keep the remote owner alive.
        scope.stack.push(sub);
        scopes.push(scope);
        return None;
    }
    Some(UnsubscribeOutcome::DropOwner(scope.root_id))
}

fn log_scope_root_mismatch(scope: &Scope, id: &ThreadSelection) {
    if scope.root_id.bytes() != id.root_id.bytes() {
        tracing::error!(
            "Somehow the current scope's root is not equal to the selected note's root. scope's root: {:?}, thread's root: {:?}",
            scope.root_id.hex(),
            id.root_id.bytes()
        );
    }
}

#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
enum ThreadScopedSub {
    RepliesByRoot,
}

fn thread_remote_sub_key(root_id: &RootNoteId) -> SubKey {
    SubKey::builder(ThreadScopedSub::RepliesByRoot)
        .with(*root_id.bytes())
        .finish()
}

fn sub_current_scope(
    ndb: &mut Ndb,
    selection: &ThreadSelection,
    local_sub_filter: Vec<Filter>,
    cur_scope: &mut Scope,
) -> bool {
    if selection.root_id.bytes() != cur_scope.root_id.bytes() {
        tracing::error!(
            "Somehow the current scope's root is not equal to the selected note's root"
        );
    }

    if let Some(sub) = ndb_sub(ndb, &local_sub_filter, selection) {
        cur_scope.stack.push(Sub {
            _selected_id: NoteId::new(*selection.selected_or_root()),
            sub,
            _filters: local_sub_filter,
        });
        return true;
    }

    false
}

fn sub_remote(
    scoped_subs: &mut ScopedSubApi<'_, '_>,
    owner: SubOwnerKey,
    key: SubKey,
    filter: Vec<Filter>,
    history_filter: Vec<Filter>,
    id: impl std::fmt::Debug,
) {
    tracing::debug!("Remote subscribe for {:?}", id);

    let identity = ScopedSubIdentity::account(owner, key);
    let config = thread_remote_sub_config(filter, history_filter);
    let _ = scoped_subs.ensure_sub(identity, config);
}

fn thread_remote_sub_config(filter: Vec<Filter>, history_filter: Vec<Filter>) -> SubConfig {
    SubConfig::live(filter)
        .full_history(FullHistoryConfig::new(history_filter))
        .build()
}

#[allow(clippy::too_many_arguments)]
fn local_sub_new_scope(
    ndb: &mut Ndb,
    scoped_subs: &mut ScopedSubApi<'_, '_>,
    account_pk: Pubkey,
    meta_id: usize,
    id: &ThreadSelection,
    local_sub_filter: Vec<Filter>,
    remote_sub_filter: Vec<Filter>,
    remote_history_filter: Vec<Filter>,
    scopes: &mut Vec<Scope>,
) -> bool {
    let root_id = id.root_id.to_note_id();
    let scope_depth = scopes.len();
    let owner = thread_scope_owner_key(account_pk, meta_id, &root_id, scope_depth);
    tracing::info!(
        "thread sub with owner: pk: {account_pk:?}, col: {meta_id}, rootid: {root_id:?}, depth: {scope_depth}"
    );
    sub_remote(
        scoped_subs,
        owner,
        thread_remote_sub_key(&root_id),
        remote_sub_filter,
        remote_history_filter,
        id,
    );

    let Some(sub) = ndb_sub(ndb, &local_sub_filter, id) else {
        let _ = scoped_subs.drop_owner(owner);
        return false;
    };

    scopes.push(Scope {
        root_id,
        stack: vec![Sub {
            _selected_id: NoteId::new(*id.selected_or_root()),
            sub,
            _filters: local_sub_filter,
        }],
    });

    true
}

#[cfg(test)]
mod tests {
    use super::*;
    use egui::Context;
    use nostrdb::Transaction;
    use notedeck::{Notedeck, RootNoteIdBuf, ScopedSubEoseStatus, ScopedSubIdentity};
    use tempfile::TempDir;

    use crate::timeline::{thread::Threads, ThreadSelection};

    struct ThreadHostHarness {
        _tmp: TempDir,
        ui_ctx: Context,
        notedeck: Notedeck,
        threads: Threads,
    }

    impl ThreadHostHarness {
        fn new() -> Self {
            let tmp = TempDir::new().expect("tmp dir");
            let ui_ctx = Context::default();
            // Force a relay onto the account so opening a thread can install a
            // *live* remote sub. `--testrunner` hands a fresh account an empty
            // bootstrap set, but forced relays (`--relay`) still apply, so this
            // gives the outbox a relay to subscribe on without dialing the real
            // network.
            let notedeck = Notedeck::init(
                &ui_ctx,
                tmp.path(),
                &[
                    "notedeck".to_owned(),
                    "--testrunner".to_owned(),
                    "--relay".to_owned(),
                    "ws://127.0.0.1:9".to_owned(),
                ],
            );

            Self {
                _tmp: tmp,
                ui_ctx,
                notedeck,
                threads: Threads::default(),
            }
        }
    }

    fn thread_selection(tag: u8) -> ThreadSelection {
        ThreadSelection::from_root_id(RootNoteIdBuf::new_unsafe([tag; 32]))
    }

    fn thread_identity(
        account_pk: Pubkey,
        col: usize,
        root_id: &NoteId,
        scope_depth: usize,
    ) -> ScopedSubIdentity {
        ScopedSubIdentity::account(
            thread_scope_owner_key(account_pk, col, root_id, scope_depth),
            thread_remote_sub_key(root_id),
        )
    }

    // Opening a thread stages remote work through `OutboxSessionHandler`; drop
    // can build `WebsocketConn`, whose native transport spawns onto Tokio, so
    // this needs a running Tokio runtime.
    #[tokio::test]
    async fn open_thread_remote_sub_restores_after_account_switch() {
        let mut h = ThreadHostHarness::new();
        let selection = thread_selection(0x33);
        let root_id = selection.root_id.to_note_id();
        let col = 9;

        let mut app_ctx = h.notedeck.app_context(&h.ui_ctx);
        let account_a = *app_ctx.accounts.selected_account_pubkey();
        let account_b = enostr::FullKeypair::generate().to_keypair();
        let account_b_pk = account_b.pubkey;
        let add_response = app_ctx
            .accounts
            .add_account(account_b)
            .expect("add account");
        assert_eq!(add_response.switch_to, account_b_pk);

        {
            let txn = Transaction::new(app_ctx.ndb).expect("txn");
            let mut scoped_subs = app_ctx.remote.scoped_subs(app_ctx.accounts);
            let _ = h.threads.open(
                app_ctx.ndb,
                &txn,
                &mut scoped_subs,
                &selection,
                true,
                col,
                0.0,
            );

            let identity = thread_identity(account_a, col, &root_id, 0);
            assert!(
                matches!(scoped_subs.sub_eose_status(identity), ScopedSubEoseStatus::Live(_)),
                "opening the thread should install one live remote thread sub for the selected account"
            );
        }

        app_ctx.select_account(&account_b_pk);
        assert!(
            h.threads
                .subs
                .get_local_for_selected(app_ctx.accounts, col)
                .is_none(),
            "switching away should leave no selected-account local thread sub"
        );

        app_ctx.select_account(&account_a);
        assert!(
            h.threads
                .subs
                .get_local_for_selected(app_ctx.accounts, col)
                .is_some(),
            "switching back should make the original account's local thread sub visible again"
        );

        let identity = thread_identity(account_a, col, &root_id, 0);
        let scoped_subs = app_ctx.remote.scoped_subs(app_ctx.accounts);
        assert!(
            matches!(scoped_subs.sub_eose_status(identity), ScopedSubEoseStatus::Live(_)),
            "switching back should restore the thread's account-scoped remote sub without reopening it"
        );
    }
}
