from __future__ import annotations

import asyncio
import os
import time
from pathlib import Path

import pytest

from core.chain.block_import import _theta_to_target, compute_header_hash
from core.types.block import Block
from core.utils.hash import ZERO32
from p2p.deps import AsyncP2PDeps, P2PDeps
from p2p.node.p2p_service import P2PService, _PeerState
from p2p.tests import free_port, tcp_multiaddr
from p2p.wire.encoding import encode_payload
from p2p.wire.message_ids import MsgID
from p2p.wire.messages import Blocks

GENESIS_PATH = Path(__file__).resolve().parents[2] / "core" / "genesis" / "genesis.json"

os.environ.setdefault("ANIMICA_P2P_DISABLE_DEFAULT_SEEDS", "1")


def _make_deps(tmp_path: Path, name: str) -> tuple[P2PDeps, AsyncP2PDeps]:
    db_path = tmp_path / f"{name}.db"
    sync_deps = P2PDeps.open(f"sqlite:///{db_path}", str(GENESIS_PATH))
    return sync_deps, AsyncP2PDeps(sync_deps)


def _mine_blocks(sync_deps: P2PDeps, count: int) -> None:
    height, head_hash = sync_deps.head()
    header = sync_deps.header_by_hash(head_hash) if head_hash else None
    if header is None:
        header = sync_deps.header_by_number(0)
    assert header is not None

    timestamp = int(getattr(header, "timestamp", 0))
    for _ in range(count):
        timestamp += 1
        target = _theta_to_target(int(getattr(header, "thetaMicro", 0)))
        child = None
        for nonce in range(0, 10000):
            candidate = header.build_child(
                timestamp=timestamp,
                state_root=header.stateRoot,
                txs_root=ZERO32,
                receipts_root=ZERO32,
                proofs_root=ZERO32,
                da_root=ZERO32,
                nonce=nonce,
                extra=b"",
            )
            if int.from_bytes(compute_header_hash(candidate), "big") <= target:
                child = candidate
                break
        assert child is not None, "failed to mine test block"
        block = Block(header=child, txs=(), proofs=(), receipts=None)
        ok, reason = sync_deps.import_block(block)
        assert ok, reason
        header = child


async def _wait_for_height(
    deps: AsyncP2PDeps, height: int, timeout: float = 20.0
) -> bool:
    deadline = time.time() + timeout
    while time.time() < deadline:
        cur, _ = await deps.head()
        if cur >= height:
            return True
        await asyncio.sleep(0.2)
    return False


async def _wait_for_peers(node: P2PService, count: int, timeout: float = 10.0) -> bool:
    deadline = time.time() + timeout
    while time.time() < deadline:
        if node.peer_count() >= count:
            return True
        await asyncio.sleep(0.1)
    return False


async def _wait_for_header_responses(
    node: P2PService, timeout: float = 10.0
) -> bool:
    deadline = time.time() + timeout
    while time.time() < deadline:
        status = node.sync_status_snapshot()
        if (
            status.last_header_response_count > 0
            or status.headers_accepted_total > 0
            or status.best_header_height > 0
        ):
            return True
        await asyncio.sleep(0.2)
    return False


@pytest.mark.asyncio
async def test_two_nodes_sync_from_genesis(tmp_path: Path) -> None:
    deps_a_sync, deps_a = _make_deps(tmp_path, "node_a_two")
    deps_b_sync, deps_b = _make_deps(tmp_path, "node_b_two")

    port_a = free_port()
    port_b = free_port()

    addr_a = tcp_multiaddr(port_a)
    addr_b = tcp_multiaddr(port_b)

    node_a = P2PService(
        listen_addrs=[addr_a],
        seeds=[],
        chain_id=deps_a_sync.chain_id,
        deps=deps_a,
        peerstore_path=str(tmp_path / "node_a_two" / "p2p"),
    )
    node_b = P2PService(
        listen_addrs=[addr_b],
        seeds=[addr_a],
        chain_id=deps_b_sync.chain_id,
        deps=deps_b,
        peerstore_path=str(tmp_path / "node_b_two" / "p2p"),
    )

    await node_a.start()
    await node_b.start()
    try:
        _mine_blocks(deps_a_sync, 3)
        await node_b.dial(addr_a)
        assert await _wait_for_peers(node_b, 1)
        assert await _wait_for_height(deps_b, 3, timeout=20.0)
        assert await _wait_for_header_responses(node_b, timeout=10.0)
        status = node_b.sync_status_snapshot()
        assert status.best_header_height > 0
    finally:
        await node_a.stop()
        await node_b.stop()


@pytest.mark.asyncio
async def test_follower_syncs_block_2_then_block_3_from_leader(tmp_path: Path) -> None:
    deps_leader_sync, deps_leader = _make_deps(tmp_path, "leader")
    deps_follower_sync, deps_follower = _make_deps(tmp_path, "follower")

    port_leader = free_port()
    port_follower = free_port()
    addr_leader = tcp_multiaddr(port_leader)
    addr_follower = tcp_multiaddr(port_follower)

    node_leader = P2PService(
        listen_addrs=[addr_leader],
        seeds=[],
        chain_id=deps_leader_sync.chain_id,
        deps=deps_leader,
        peerstore_path=str(tmp_path / "leader" / "p2p"),
    )
    node_follower = P2PService(
        listen_addrs=[addr_follower],
        seeds=[addr_leader],
        chain_id=deps_follower_sync.chain_id,
        deps=deps_follower,
        peerstore_path=str(tmp_path / "follower" / "p2p"),
    )

    try:
        _mine_blocks(deps_leader_sync, 3)
        _mine_blocks(deps_follower_sync, 1)
    except AssertionError as exc:
        if "No module named 'animica'" in str(exc):
            pytest.skip("sync block import runtime dependency missing in this environment")
        raise
    assert deps_follower_sync.head()[0] == 1
    assert deps_leader_sync.head()[0] == 3

    await node_leader.start()
    await node_follower.start()
    observed_target = False
    try:
        await node_follower.dial(addr_leader)
        assert await _wait_for_peers(node_follower, 1)

        deadline = time.time() + 20.0
        while time.time() < deadline:
            status = node_follower.sync_status_snapshot()
            if status.target_height is not None and int(status.target_height) >= 3:
                observed_target = True
            if deps_follower_sync.head()[0] >= 3:
                break
            await asyncio.sleep(0.2)

        assert deps_follower_sync.head()[0] >= 3
        assert observed_target
        status = node_follower.sync_status_snapshot(refresh=True)
        assert status.head_height >= 3
    finally:
        await node_follower.stop()
        await node_leader.stop()


@pytest.mark.asyncio
async def test_sync_recovers_under_deterministic_packet_loss(tmp_path: Path) -> None:
    deps_leader_sync, deps_leader = _make_deps(tmp_path, "leader_lossy")
    deps_follower_sync, deps_follower = _make_deps(tmp_path, "follower_lossy")

    port_leader = free_port()
    port_follower = free_port()
    addr_leader = tcp_multiaddr(port_leader)
    addr_follower = tcp_multiaddr(port_follower)

    node_leader = P2PService(
        listen_addrs=[addr_leader],
        seeds=[],
        chain_id=deps_leader_sync.chain_id,
        deps=deps_leader,
        peerstore_path=str(tmp_path / "leader_lossy" / "p2p"),
    )
    node_follower = P2PService(
        listen_addrs=[addr_follower],
        seeds=[addr_leader],
        chain_id=deps_follower_sync.chain_id,
        deps=deps_follower,
        peerstore_path=str(tmp_path / "follower_lossy" / "p2p"),
    )

    # Keep retries fast for deterministic chaos timing in CI.
    node_follower._sync_request_timeout = 1.0
    node_follower._sync_header_watchdog_timeout = 1.0
    node_follower._sync_stall_timeout = 6.0

    drop_budget: dict[MsgID, int] = {MsgID.HEADERS: 3, MsgID.BLOCKS: 1}
    delayed_budget: dict[MsgID, int] = {MsgID.HEADERS: 2, MsgID.BLOCKS: 1}
    dropped_counts: dict[MsgID, int] = {MsgID.HEADERS: 0, MsgID.BLOCKS: 0}
    original_send = node_leader._send

    async def _flaky_send(peer: _PeerState, msg_id: MsgID, payload: object) -> None:
        if msg_id in (MsgID.HEADERS, MsgID.BLOCKS):
            budget = drop_budget.get(msg_id, 0)
            if budget > 0:
                drop_budget[msg_id] = budget - 1
                dropped_counts[msg_id] = int(dropped_counts.get(msg_id, 0)) + 1
                return
            delay = delayed_budget.get(msg_id, 0)
            if delay > 0:
                delayed_budget[msg_id] = delay - 1
                await asyncio.sleep(0.05)
        await original_send(peer, msg_id, payload)

    node_leader._send = _flaky_send  # type: ignore[assignment]

    _mine_blocks(deps_leader_sync, 6)
    await node_leader.start()
    await node_follower.start()
    try:
        await node_follower.dial(addr_leader)
        assert await _wait_for_peers(node_follower, 1)
        assert await _wait_for_height(deps_follower, 6, timeout=35.0)

        status = node_follower.sync_status_snapshot(refresh=True)
        timeout_total = int(node_follower._stats.get("headers_req_timeout", 0)) + int(
            node_follower._stats.get("blocks_req_timeout", 0)
        )
        assert int(dropped_counts.get(MsgID.HEADERS, 0)) >= 1
        assert int(dropped_counts.get(MsgID.BLOCKS, 0)) >= 1
        assert timeout_total >= 1
        assert status.stall_reason is None
        assert status.fatal_error is None
        assert status.head_height >= 6
    finally:
        await node_follower.stop()
        await node_leader.stop()


@pytest.mark.asyncio
async def test_three_nodes_converge_headers_first(tmp_path: Path) -> None:
    deps_a_sync, deps_a = _make_deps(tmp_path, "node_a")
    deps_b_sync, deps_b = _make_deps(tmp_path, "node_b")
    deps_c_sync, deps_c = _make_deps(tmp_path, "node_c")

    port_a = free_port()
    port_b = free_port()
    port_c = free_port()

    addr_a = tcp_multiaddr(port_a)
    addr_b = tcp_multiaddr(port_b)

    node_a = P2PService(
        listen_addrs=[addr_a],
        seeds=[],
        chain_id=deps_a_sync.chain_id,
        deps=deps_a,
        peerstore_path=str(tmp_path / "node_a" / "p2p"),
    )
    node_b = P2PService(
        listen_addrs=[addr_b],
        seeds=[addr_a],
        chain_id=deps_b_sync.chain_id,
        deps=deps_b,
        peerstore_path=str(tmp_path / "node_b" / "p2p"),
    )
    node_c = P2PService(
        listen_addrs=[tcp_multiaddr(port_c)],
        seeds=[addr_b],
        chain_id=deps_c_sync.chain_id,
        deps=deps_c,
        peerstore_path=str(tmp_path / "node_c" / "p2p"),
    )

    await node_a.start()
    try:
        _mine_blocks(deps_a_sync, 50)

        await node_b.start()
        await node_b.dial(addr_a)
        assert await _wait_for_peers(node_b, 1)
        assert await _wait_for_height(deps_b, 50, timeout=30.0)

        await node_c.start()
        await node_c.dial(addr_b)
        assert await _wait_for_peers(node_c, 1)
        assert await _wait_for_height(deps_c, 50, timeout=30.0)

        _mine_blocks(deps_a_sync, 5)
        assert await _wait_for_height(deps_b, 55, timeout=30.0)
        assert await _wait_for_height(deps_c, 55, timeout=30.0)

        _, head_hash_a = deps_a_sync.head()
        _, head_hash_b = deps_b_sync.head()
        _, head_hash_c = deps_c_sync.head()
        assert head_hash_a == head_hash_b == head_hash_c
    finally:
        await node_a.stop()
        await node_b.stop()
        await node_c.stop()


@pytest.mark.asyncio
async def test_invalid_block_penalizes_peer_and_sync_continues(
    tmp_path: Path,
) -> None:
    deps_a_sync, deps_a = _make_deps(tmp_path, "node_a_bad")
    deps_b_sync, deps_b = _make_deps(tmp_path, "node_b_bad")

    port_a = free_port()
    port_b = free_port()

    addr_a = tcp_multiaddr(port_a)
    addr_b = tcp_multiaddr(port_b)

    node_a = P2PService(
        listen_addrs=[addr_a],
        seeds=[],
        chain_id=deps_a_sync.chain_id,
        deps=deps_a,
        peerstore_path=str(tmp_path / "node_a_bad" / "p2p"),
    )
    node_b = P2PService(
        listen_addrs=[addr_b],
        seeds=[addr_a],
        chain_id=deps_b_sync.chain_id,
        deps=deps_b,
        peerstore_path=str(tmp_path / "node_b_bad" / "p2p"),
    )

    await node_a.start()
    await node_b.start()
    try:
        _mine_blocks(deps_a_sync, 5)
        await node_b.dial(addr_a)
        assert await _wait_for_peers(node_b, 1)
        assert await _wait_for_height(deps_b, 5, timeout=20.0)

        bad_peer = _PeerState(
            session_id="bad-peer",
            remote="bad-peer:0",
            direction="inbound",
            conn=None,
            stream=None,
            framer=None,
            write_lock=asyncio.Lock(),
        )
        payload = encode_payload(Blocks(blocks=[b"not-a-cbor-block"]))
        await node_b._handle_blocks(bad_peer, payload)

        penalties = node_b.sync_status_snapshot().peer_penalties
        assert (
            penalties.get("bad-peer", 0) >= 1
            or penalties.get("bad-peer:0", 0) >= 1
        )

        _mine_blocks(deps_a_sync, 2)
        await node_b.force_sync()
        if not await _wait_for_height(deps_b, 7, timeout=20.0):
            await node_b.force_sync()
        assert await _wait_for_height(deps_b, 7, timeout=30.0)
    finally:
        await node_a.stop()
        await node_b.stop()
