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

use crate::OracleInfoFetcher;

use alloy_sol_types::{SolCall, SolEvent, SolValue};

use async_trait::async_trait;

use lyquor_api::actor::FetchOracleInfo;

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

use lyquor_eth::eth;

use lyquor_jsonrpc::client::ClientHandle;

use lyquor_jsonrpc::types::{

    Block, BlockNumber, EthCall, EthCallResp, EthCallTx, EthGetBlockByNumber, EthGetBlockByNumberResp, EthGetLogs,

    EthGetLogsResp, EthNewHeadUpdate, EthSubscribe,

};

use lyquor_primitives::{Address, U256, alloy_primitives, debug_struct_name, oracle::OracleServiceTarget};

use std::collections::{HashMap, hash_map};

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;

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;

pub(super) struct SlotCrawlerRuntime {

    client: ClientHandle,

    contracts: HashMap<Address, Contract>,

    blocks_per_request: usize,

    finality_tag: BlockNumber,

}

impl SlotCrawlerRuntime {

    async fn run(

        mut self, mut wake_rx: mpsc::Receiver<()>, mut contract_rx: mpsc::UnboundedReceiver<(Address, Contract)>,

    ) {

        loop {

            tokio::select! {

                Some((addr, contract)) = contract_rx.recv() => {

                    self.handle_contract(addr, contract);

                }

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

                else => break,

            }

            while let Ok((addr, contract)) = contract_rx.try_recv() {

                self.handle_contract(addr, contract);

            }

            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(&mut self, addr: Address, contract: Contract) {

        tracing::debug!("SlotCrawler: Added contract {addr}.");

        self.contracts.insert(addr, contract);

    }

    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: &[Address], from_block: BlockNumber, to_block: BlockNumber,

    ) -> anyhow::Result<WatchedLogsBatch> {

        let resp: EthGetLogsResp = client

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

                from_block,

                to_block,

                address: batch.to_vec(),

                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<Address, Contract>, resync: &mut bool,

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

    ) {

        'outer: for (addr, slots) in updates {

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

                Some(c) => c,

                None => {

                    tracing::debug!("Contract {addr} was removed.");

                    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 c.lyquid.send(SwitchContract { pos: s.pos }).await.is_err() {

                        tracing::warn!("Sync: LyquidBackend of {addr} already stopped.");

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

                        // is already unreachable.

                        contracts.remove(&addr);

                        continue 'outer;

                    }

                    assert!(c.new_addr.is_none());

                    c.new_addr = Some(new_addr);

                    *resync = true;

                    // Ignore all subsequent slots for this old address.

                    break;

                }

                to_be_filled.push(s);

            }

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

                tracing::error!("Sync: Failed to fill slots from {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.

                contracts.remove(&addr);

            }

        }

    }

    // 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<Address, 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 = Vec::new();

        for (addr, c) in contracts.iter().filter_map(|(addr, c)| {

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

                None

            } else {

                Some((*addr, c))

            }

        }) {

            match c.seq.next_block().await {

                Ok(bn) => {

                    chosen.push((addr, bn));

                }

                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().

                }

            }

        }

        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.1);

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

        let mut next_block_number = None;

        let mut batch = Vec::new();

        let mut resync = false;

        for (addr, block) 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.push(addr);

                    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, from_block, to_block).await?;

                        Self::update_contracts(contracts, &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.push(addr);

        }

        // 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 in batch {

                if let hash_map::Entry::Occupied(mut e) = contracts.entry(addr) {

                    let c = e.get_mut();

                    if let Some(new_addr) = c.new_addr.take() {

                        // This lyquid is switching the contract. So we need to update the address

                        // info in the watch list. And we should *not* advance the next block

                        // because this lyquid needs to check for the slots using the new adddress

                        // in the previous block.

                        let c = e.remove();

                        contracts.insert(new_addr, c);

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

                    } else {

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

                        // have successfully filled all slots up to n.

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

                            tracing::error!("SlotCrawler: Failed to advance fetching progress 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().

                        }

                    }

                }

            }

        }

        if resync {

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

        }

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

        Ok(())

    }

}

#[derive(Clone)]

pub(super) struct SlotCrawler {

    client: ClientHandle,

    wake_tx: mpsc::Sender<()>,

    contract_tx: mpsc::UnboundedSender<(Address, Contract)>,

    task_tracker: TaskTracker,

    finality_tag: BlockNumber,

}

debug_struct_name!(SlotCrawler);

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(

            SlotCrawlerRuntime {

                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 {

            wake_tx,

            contract_tx,

            task_tracker,

            client,

            finality_tag,

        }

    }

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

        self.task_tracker.wait().await;

    }

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

        match self.wake_tx.try_send(()) {

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

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

        }

    }

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

        self.contract_tx

            .send((addr, contract))

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

    }

}

#[async_trait]

impl OracleInfoFetcher for SlotCrawler {

    async fn fetch_oracle_info(&self, msg: FetchOracleInfo) -> Option<lyquor_primitives::oracle::OracleEpochInfo> {

        let OracleServiceTarget::EVM { eth_contract, .. } = msg.target.target else {

            return None;

        };

        let topic = alloy_primitives::keccak256(msg.topic.as_bytes());

        fetch_oracle_info(&self.client, self.finality_tag, eth_contract, topic, msg.full_config).await

    }

}

async fn finalized_block_number(client: &ClientHandle, finality_tag: BlockNumber) -> Option<u64> {

    let finalized: Block = client

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

            block: finality_tag,

            full_tx: false,

        })

        .await

        .ok()?

        .0?;

    match finalized.number {

        BlockNumber::Quantity(bn) => Some(bn.to()),

        _ => None,

    }

}

async fn fetch_oracle_info(

    client: &ClientHandle, finality_tag: BlockNumber, eth_contract: Address, topic: alloy_primitives::B256,

    full_config: bool,

) -> Option<lyquor_primitives::oracle::OracleEpochInfo> {

    let block_number = BlockNumber::Quantity(U256::from(finalized_block_number(client, finality_tag).await?));

    let get_epoch_input = eth::ISequenceBackend::getEpochCall { topic }.abi_encode();

    let epoch_ret: EthCallResp = client

        .request(EthCall {

            tx: EthCallTx {

                from: None,

                to: Some(eth_contract),

                gas: None,

                gas_price: None,

                value: None,

                data: Some(get_epoch_input.into()),

            },

            block_number: block_number.clone(),

        })

        .await

        .ok()?;

    let get_hash_input = eth::ISequenceBackend::getConfigHashCall { topic }.abi_encode();

    let hash_ret: EthCallResp = client

        .request(EthCall {

            tx: EthCallTx {

                from: None,

                to: Some(eth_contract),

                gas: None,

                gas_price: None,

                value: None,

                data: Some(get_hash_input.into()),

            },

            block_number: block_number.clone(),

        })

        .await

        .ok()?;

    let get_count_input = eth::ISequenceBackend::getChangeCountCall { topic }.abi_encode();

    let count_ret: EthCallResp = client

        .request(EthCall {

            tx: EthCallTx {

                from: None,

                to: Some(eth_contract),

                gas: None,

                gas_price: None,

                value: None,

                data: Some(get_count_input.into()),

            },

            block_number: block_number.clone(),

        })

        .await

        .ok()?;

    let epoch = eth::ISequenceBackend::getEpochCall::abi_decode_returns(epoch_ret.0.as_ref()).ok()?;

    let config_hash = eth::ISequenceBackend::getConfigHashCall::abi_decode_returns(hash_ret.0.as_ref()).ok()?;

    let change_count = eth::ISequenceBackend::getChangeCountCall::abi_decode_returns(count_ret.0.as_ref()).ok()?;

    let config_hash: lyquor_primitives::HashBytes = config_hash.0.into();

    let config = if full_config && config_hash != [0; 32].into() {

        let get_config_input = eth::ISequenceBackend::getConfigCall { topic }.abi_encode();

        let config_ret: EthCallResp = client

            .request(EthCall {

                tx: EthCallTx {

                    from: None,

                    to: Some(eth_contract),

                    gas: None,

                    gas_price: None,

                    value: None,

                    data: Some(get_config_input.into()),

                },

                block_number,

            })

            .await

            .ok()?;

        let config = eth::ISequenceBackend::getConfigCall::abi_decode_returns(config_ret.0.as_ref()).ok()?;

        let actual_hash: lyquor_primitives::Hash = alloy_primitives::keccak256(&config.abi_encode()).0.into();

        if actual_hash != config_hash.into_inner() {

            return None;

        }

        Some(lyquor_primitives::oracle::OracleConfig {

            committee: config

                .committee

                .into_iter()

                .map(|signer| lyquor_primitives::oracle::OracleSigner {

                    id: signer.id,

                    key: signer.nodeID.as_slice().to_vec().into(),

                })

                .collect(),

            threshold: config.threshold,

        })

    } else {

        None

    };

    Some(lyquor_primitives::oracle::OracleEpochInfo {

        epoch,

        config_hash,

        change_count,

        config,

    })

}

#[derive(Clone)]

pub(super) struct HeadWatcher {

    task_tracker: TaskTracker,

}

impl HeadWatcher {

    pub(super) fn new(client: ClientHandle, slot_crawler: SlotCrawler, 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: SlotCrawler) {

        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;

                    }

                }

            }

        }

    }

}