use super::{ChainPos, Contract, Slot};

use alloy_sol_types::SolEvent;

use futures::FutureExt as _;

use lyquor_api::anyhow::{self, Context as _Context};

use lyquor_eth::eth;

use lyquor_jsonrpc::client::ClientHandle;

use lyquor_jsonrpc::types::{

    Block, BlockNumber, EthGetBlockByNumber, EthGetBlockByNumberResp, EthGetLogs, EthGetLogsResp, EthNewHeadUpdate,

    EthSubscribe,

};

use lyquor_primitives::{Address, LyquidID, U256, alloy_primitives, debug_struct_name};

use std::collections::HashMap;

use std::time::Duration;

use tokio::sync::mpsc;

use tokio::sync::mpsc::error::TrySendError;

use tokio::time::MissedTickBehavior;

use tokio_util::sync::CancellationToken;

use tokio_util::task::TaskTracker;

use tokio_util::time::FutureExt as _;

use tracing::Instrument;

enum WatchedLogKind {

    Slot,

}

impl WatchedLogKind {

    fn from_topic0(topic0: Option<alloy_primitives::B256>) -> Option<Self> {

        match topic0 {

            Some(t) if t == eth::Slot::SIGNATURE_HASH => Some(Self::Slot),

            _ => None,

        }

    }

    fn topic0_filter() -> Vec<alloy_primitives::B256> {

        vec![eth::Slot::SIGNATURE_HASH]

    }

}

#[derive(Default)]

struct WatchedLogsBatch {

    slots: HashMap<Address, Vec<Slot>>,

}

// Wakeups are edge-triggered: one queued wake is enough to force another sync pass.

const SLOT_CRAWLER_WAKE_CAPACITY: usize = 1;

struct SlotCrawlerWorker {

    client: ClientHandle,

    contracts: HashMap<LyquidID, Contract>,

    blocks_per_request: usize,

    finality_tag: BlockNumber,

}

enum ContractCommand {

    Upsert(Contract),

    RemoveLyquid(LyquidID),

}

impl SlotCrawlerWorker {

    async fn run(mut self, mut wake_rx: mpsc::Receiver<()>, mut contract_rx: mpsc::UnboundedReceiver<ContractCommand>) {

        loop {

            tokio::select! {

                Some(cmd) = contract_rx.recv() => {

                    self.handle_contract_command(cmd);

                }

                Some(()) = wake_rx.recv() => {}

                else => break,

            }

            while let Ok(cmd) = contract_rx.try_recv() {

                self.handle_contract_command(cmd);

            }

            while let Ok(()) = wake_rx.try_recv() {}

            if let Err(e) = Self::sync(

                &mut self.contracts,

                self.finality_tag,

                self.blocks_per_request,

                self.client.clone(),

            )

            .await

            {

                tracing::error!("SlotCrawler: Error during sync: {e:?}");

            }

        }

    }

    fn handle_contract_command(&mut self, cmd: ContractCommand) {

        match cmd {

            ContractCommand::Upsert(contract) => {

                tracing::debug!(

                    "SlotCrawler: Added contract {} for lyquid {}.",

                    contract.addr,

                    contract.id

                );

                self.contracts.insert(contract.id, contract);

            }

            ContractCommand::RemoveLyquid(id) => {

                if self.contracts.remove(&id).is_some() {

                    tracing::debug!("SlotCrawler: Removed contract watch for lyquid {id}.");

                }

            }

        }

    }

    fn extract_watched_logs(resp: EthGetLogsResp) -> WatchedLogsBatch {

        let mut out = WatchedLogsBatch::default();

        for e in resp.0.into_iter() {

            match WatchedLogKind::from_topic0(e.topics.first().copied()) {

                Some(WatchedLogKind::Slot) => {

                    // From https://docs.rs/alloy/latest/alloy/dyn_abi/index.html:

                    // We strongly recommend using the static encoder/decoder when possible. The

                    // dynamic encoder/decoder is significantly more expensive, especially for complex

                    // types. It is also significantly more error prone, as the mapping between

                    // solidity types and rust types is not enforced by the compiler.

                    let rlp_decoded = alloy_primitives::LogData::new_unchecked(e.topics, e.data);

                    let le = match eth::Slot::decode_log_data(&rlp_decoded).ok() {

                        Some(le) => le,

                        None => continue,

                    };

                    let mut slot = Slot::from(le);

                    let block_number = match e.block_number {

                        BlockNumber::Quantity(q) => q,

                        _ => unreachable!(),

                    };

                    let bn: u64 = block_number.to();

                    let ln: u64 = e.log_index.to();

                    if ln >> 32 != 0 {

                        tracing::warn!("Skipped a block that's too big.");

                        continue;

                    } // block number uses 44 bits

                    slot.pos = ChainPos::new(bn, ln as u32);

                    tracing::debug!("SlotCrawler: extracted slot={slot:?}.");

                    out.slots.entry(e.address).or_default().push(slot);

                }

                None => {}

            }

        }

        out

    }

    async fn fetch_watched_logs_for_batch(

        client: &ClientHandle, batch: impl IntoIterator<Item = &Address>, from_block: BlockNumber,

        to_block: BlockNumber,

    ) -> anyhow::Result<WatchedLogsBatch> {

        let resp: EthGetLogsResp = client

            .request::<EthGetLogs, EthGetLogsResp>(EthGetLogs {

                from_block,

                to_block,

                address: batch.into_iter().copied().collect(),

                topics: vec![WatchedLogKind::topic0_filter()],

            })

            .await

            .with_context(|| format!("Failed to retrieve {}..{}.", from_block, to_block))?;

        Ok(Self::extract_watched_logs(resp))

    }

    async fn update_contracts(

        contracts: &mut HashMap<LyquidID, Contract>, batch_ids_by_addr: &HashMap<Address, LyquidID>, resync: &mut bool,

        updates: impl Iterator<Item = (Address, Vec<Slot>)>,

    ) {

        for (addr, slots) in updates {

            let id = match batch_ids_by_addr.get(&addr).copied() {

                Some(id) => id,

                None => {

                    tracing::debug!("Skipped stale logs for unknown contract address {addr}.");

                    continue;

                }

            };

            let mut remove_self = false;

            {

                let c = match contracts.get_mut(&id) {

                    Some(c) => c,

                    None => {

                        tracing::debug!("Lyquid {id} was removed.");

                        continue;

                    }

                };

                if c.addr != addr {

                    tracing::debug!(

                        "Skipped stale logs for Lyquid {id} because current address {} does not match {}.",

                        c.addr,

                        addr

                    );

                    continue;

                }

                let mut to_be_filled = Vec::new();

                for s in slots.into_iter() {

                    if let crate::NextContractID::Eth(new_addr) = s.switch_contract {

                        if let Err(e) = c.lyquid.switch_contract(s.pos).await {

                            tracing::warn!("Sync: LyquidBackend of {id} is unavailable: {e:?}");

                            // We need to remove it from the watch list because the LyquidBackend

                            // worker is already unreachable.

                            remove_self = true;

                        } else {

                            tracing::debug!("SlotCrawler: Switching contract for {id}: {addr} => {new_addr}.");

                            c.addr = new_addr;

                            *resync = true;

                        }

                        // Ignore all subsequent slots for this old address.

                        break;

                    }

                    to_be_filled.push(s);

                }

                if !remove_self {

                    if let Err(e) = c.seq.fill_slots(to_be_filled).await {

                        tracing::error!("Sync: Failed to fill slots from {id} ({addr}): {e:?}");

                        // We need to remove it from the watch list so the lyquid's next block won't be

                        // advanced because the slots weren't filled properly.

                        remove_self = true;

                    }

                }

            }

            if remove_self {

                contracts.remove(&id);

                continue;

            }

        }

    }

    // This function is the core logic of SlotCrawler.

    //

    // It progressively tries to bring all contracts ("watch list") to be the same, latest block

    // known on the chain. While each query to the chain may bundle multiple contracts, the

    // philosophy is to still treat each contract's sync progress separately. This is good for

    // robust handling of LyquidBackend, because upon any errors during syncing, there is no

    // incorrectly advanced progress for any Lyquids.

    #[tracing::instrument(level = "trace", skip_all)]

    async fn sync(

        contracts: &mut HashMap<LyquidID, Contract>, finality_tag: BlockNumber, blocks_per_request: usize,

        client: ClientHandle,

    ) -> anyhow::Result<()> {

        // Choose the contracts that should be considered during this sync.

        let mut chosen = {

            let next_blocks = contracts.iter().filter_map(|(id, c)| {

                if c.seq.flow_control().should_stop_filling() {

                    None

                } else {

                    Some(c.seq.next_block().map(|bn| (bn, *id, c.addr)))

                }

            });

            futures::future::join_all(next_blocks)

                .await

                .into_iter()

                .filter_map(|(res, id, addr)| {

                    res.map_err(|e| tracing::warn!("Sync: Failed to get the next block number for {addr}: {e:?}"))

                        // We don't have to remove this contract from the watch list because it

                        // is not harmful to try again upon next sync().

                        .ok()

                        .map(|bn| (bn, id, addr))

                })

                .collect::<Vec<_>>()

        };

        if chosen.is_empty() {

            tracing::debug!("No contract need to fetch.");

            return Ok(());

        }

        // Let set a fix goal for this sync: up to the `finalized` block for all chosen contracts.

        let finalized: Block = client

            .request::<EthGetBlockByNumber, EthGetBlockByNumberResp>(EthGetBlockByNumber {

                block: finality_tag,

                full_tx: false,

            })

            .in_current_span()

            .await

            .map_err(|e| e.into())

            .and_then(|f: EthGetBlockByNumberResp| f.0.ok_or(anyhow::anyhow!("result is null")))

            .context("Failed obtain the last finalized block.")?;

        let finalized = match finalized.number {

            BlockNumber::Quantity(b) => b,

            _ => {

                return Err(anyhow::anyhow!("Invalid finalized block number."));

            }

        }

        .to::<u64>();

        // Because the next chain position for each chosen contract (lyquid) could be different

        // (some may just be added to the watch list of SlotCrawler), we need to bundle the queries

        // in a way that is:

        //

        // - Not always just for a single address (to reduce the number of queries).

        // - Not always ask for a lot of addreses whose corresponding blocks are already

        // queried (to reduce the redundant data returned by each query).

        //

        // A clean and simple approach is to sort them according to their respective next block to

        // crawl from, and start from the lowest to catch up with the second lowest, then we bundle

        // up the contracts in a rolling style so there is *no* redundant data received.

        //

        // At the end of the rolling, if there is no error or contract switching (which requires a

        // resync), all current contracts are brought to the same next block value.

        chosen.sort_by_key(|k| k.0);

        let sentinel_id = LyquidID::from(Address::ZERO);

        chosen.push((finalized + 1, sentinel_id, Address::ZERO)); // sentinel pos, to flush all batches

        let mut next_block_number = None;

        let mut batch = HashMap::new();

        let mut resync = false;

        for (block, id, addr) in chosen {

            if let Some(mut n) = next_block_number {

                if block == n {

                    // in the same batch (same next block), let's don't query yet

                    batch.insert(addr, id);

                    continue;

                } else {

                    // Now we need to query with the existing items in `batch` because this

                    // candidate's next block is larger.

                    while n < block {

                        // `end` is inclusive

                        let mut end: u64 = n + blocks_per_request as u64 - 1;

                        end = end.min(block - 1);

                        tracing::info!("Fetching [{}..{}] ({} addrs).", n, end, batch.len());

                        let from_block = BlockNumber::Quantity(U256::from(n));

                        let to_block = BlockNumber::Quantity(U256::from(end));

                        let watched =

                            Self::fetch_watched_logs_for_batch(&client, batch.keys(), from_block, to_block).await?;

                        Self::update_contracts(contracts, &batch, &mut resync, watched.slots.into_iter()).await;

                        n = end + 1;

                    }

                    // Now n == block, we can change the next value and continue bundling more addresses.

                }

            }

            next_block_number = Some(block);

            batch.insert(addr, id);

        }

        // Finalizing the sync by setting the next block (telling FCO about the progress) for each

        // Lyquid whose contract has successfully synced to n.

        if let Some(n) = next_block_number {

            assert_eq!(n, finalized + 1);

            for (addr, id) in batch {

                if let Some(c) = contracts.get_mut(&id) {

                    if c.addr != addr {

                        tracing::debug!(

                            "SlotCrawler: Skipping stale batch entry for {id}; current address is {}, expected {}.",

                            c.addr,

                            addr

                        );

                        continue;

                    }

                    if let Err(e) = c.seq.set_next_block(n).await {

                        // Otherwise, we can safetly tell FCO to advance the progress because we

                        // have successfully filled all slots up to n.

                        tracing::error!("SlotCrawler: Failed to advance fetching progress for {id} ({addr}): {e:?}");

                        // We don't have to remove this contract from the watch list because it

                        // is not harmful to try again upon next sync().

                    }

                }

            }

        }

        if resync {

            Box::pin(Self::sync(contracts, finality_tag, blocks_per_request, client)).await?;

        }

        tracing::info!("SlotCrawler: Sync complete.");

        Ok(())

    }

}

// TODO: Drop Clone once the outer backend ownership is explicit enough that only

// SlotCrawlerHandle needs to be cloned.

#[derive(Clone)]

pub(super) struct SlotCrawler {

    handle: SlotCrawlerHandle,

    task_tracker: TaskTracker,

}

debug_struct_name!(SlotCrawler);

#[derive(Clone)]

pub(super) struct SlotCrawlerHandle {

    wake_tx: mpsc::Sender<()>,

    contract_tx: mpsc::UnboundedSender<ContractCommand>,

}

debug_struct_name!(SlotCrawlerHandle);

impl SlotCrawler {

    pub(super) fn new(

        client: ClientHandle, blocks_per_request: usize, finality_tag: BlockNumber, shutdown: CancellationToken,

    ) -> Self {

        let (wake_tx, wake_rx) = mpsc::channel(SLOT_CRAWLER_WAKE_CAPACITY);

        let (contract_tx, contract_rx) = mpsc::unbounded_channel();

        let task_tracker = TaskTracker::new();

        task_tracker.spawn(

            SlotCrawlerWorker {

                contracts: HashMap::new(),

                blocks_per_request,

                finality_tag,

                client: client.clone(),

            }

            .run(wake_rx, contract_rx)

            .with_cancellation_token_owned(shutdown),

        );

        task_tracker.close();

        Self {

            handle: SlotCrawlerHandle { wake_tx, contract_tx },

            task_tracker,

        }

    }

    pub(super) fn handle(&self) -> SlotCrawlerHandle {

        self.handle.clone()

    }

    pub(super) async fn wait_for_shutdown(&self) {

        self.task_tracker.wait().await;

    }

}

impl SlotCrawlerHandle {

    pub(super) fn wake(&self) -> bool {

        match self.wake_tx.try_send(()) {

            Ok(()) | Err(TrySendError::Full(_)) => true,

            Err(TrySendError::Closed(_)) => false,

        }

    }

    // TODO: consider return a token that automatically removes the contract from the watch list when dropped,

    // to avoid the need of explicit removal and the potential risk of forgetting to remove.

    pub(super) async fn add_contract(&self, contract: Contract) -> anyhow::Result<()> {

        self.contract_tx

            .send(ContractCommand::Upsert(contract))

            .map_err(|_| anyhow::anyhow!("SlotCrawler control loop already stopped"))

    }

    pub(super) fn remove_lyquid(&self, lyquid: LyquidID) -> anyhow::Result<()> {

        self.contract_tx

            .send(ContractCommand::RemoveLyquid(lyquid))

            .map_err(|_| anyhow::anyhow!("SlotCrawler control loop already stopped"))

    }

}

#[derive(Clone)]

pub(super) struct HeadWatcher {

    task_tracker: TaskTracker,

}

impl HeadWatcher {

    pub(super) fn new(client: ClientHandle, slot_crawler: SlotCrawlerHandle, shutdown: CancellationToken) -> Self {

        let task_tracker = TaskTracker::new();

        task_tracker.spawn(Self::head_watcher_loop(client, slot_crawler).with_cancellation_token_owned(shutdown));

        task_tracker.close();

        Self { task_tracker }

    }

    pub(super) async fn wait_for_shutdown(&self) {

        self.task_tracker.wait().await;

    }

    async fn subscribe_to_new_heads(

        client: &ClientHandle,

    ) -> Option<lyquor_jsonrpc::client::Subscription<EthNewHeadUpdate>> {

        match client.subscribe(EthSubscribe("newHeads".to_string())).await {

            Ok(subscription) => {

                tracing::debug!("Subscribed to new head");

                Some(subscription)

            }

            Err(e) => {

                tracing::error!("Failed to subscribe to new head: {e:?}");

                None

            }

        }

    }

    async fn head_watcher_loop(client: ClientHandle, slot_crawler: SlotCrawlerHandle) {

        let mut subscription = Self::subscribe_to_new_heads(&client).await;

        // NOTE: Hardhat Network may silently lose eth_subscribe state after a period of

        // inactivity (e.g., no block mining). Although the WebSocket connection remains

        // open, subscriptions can become stale — resulting in:

        // - No new calls being delivered, even when blocks are mined later

        // - eth_unsubscribe failing with "Subscription not found"

        // This is likely due to internal state cleanup or subscription manager desync in Hardhat's

        // in-memory event system. To work around this, we re-subscribe when activity resumes.

        //

        // This is not an issue for production chains like Ethereum.

        let mut periodic_poll = tokio::time::interval(Duration::from_secs(10));

        periodic_poll.set_missed_tick_behavior(MissedTickBehavior::Delay);

        periodic_poll.tick().await;

        let mut reset = tokio::time::interval(Duration::from_secs(60));

        reset.set_missed_tick_behavior(MissedTickBehavior::Delay);

        reset.tick().await;

        loop {

            tokio::select! {

                _ = periodic_poll.tick() => {

                    if !slot_crawler.wake() {

                        break;

                    }

                }

                _ = reset.tick() => {

                    tracing::debug!("Resetting new head subscription");

                    subscription = Self::subscribe_to_new_heads(&client).await;

                }

                update = async {

                    match &mut subscription {

                        Some(subscription) => subscription.next().await,

                        None => std::future::pending().await,

                    }

                } => match update {

                    Some(msg) => {

                        tracing::debug!("New head: {:?}.", msg);

                        if !slot_crawler.wake() {

                            break;

                        }

                    }

                    None => {

                        tracing::warn!("New head subscription ended; waiting for reset to retry");

                        subscription = None;

                    }

                }

            }

        }

    }

}