// SPDX-License-Identifier: MIT
pragma solidity 0.8.23;

import "@openzeppelin/contracts/security/ReentrancyGuard.sol";
import "@openzeppelin/contracts/token/ERC20/IERC20.sol";
import "@openzeppelin/contracts/utils/cryptography/ECDSA.sol";

/**
 * @title AereStateChannels — single-contract bidirectional payment channels
 *                            with HTLC support (Connext/Hop-style)
 * @notice Two parties (A, B) lock ERC-20 funds into a channel. They
 *         transact off-chain by exchanging EIP-712 signed channel states
 *         {channelId, nonce, balanceA, balanceB, htlcLocked, lockHash, lockTimeout}.
 *         When done, EITHER party can close the channel:
 *           - Cooperative: both sign the latest state, instant settlement.
 *           - Challenge: one party submits the latest signed state and
 *             starts a 24h challenge window. Counterparty can submit a
 *             newer-nonce state to override. After window expires,
 *             anyone can call `finalize` to release funds.
 *
 *         HTLC payments inside the channel are settled by revealing a
 *         hash preimage on-chain (`settleHtlc`) within the lockTimeout.
 *         If timeout expires unrevealed, the lock returns to the funder.
 *
 *         IMMUTABILITY:
 *           - DOMAIN_SEPARATOR pinned at deploy
 *           - CHALLENGE_WINDOW immutable at deploy
 *           - No admin, no owner, no upgrade path
 *
 *         OUT OF SCOPE for V1:
 *           - Multi-hop routing (would need a Lightning-style HTLC graph)
 *           - Watchtowers (off-chain monitoring services)
 *           - Probabilistic micropayments
 *         These belong in Phase 2 contracts that USE this as the
 *         settlement primitive.
 *
 * @dev    Nonce is monotonic per channel. Challenge with a higher nonce
 *         wins; ties + lower nonces are rejected.
 */
contract AereStateChannels is ReentrancyGuard {
    using ECDSA for bytes32;

    /* ------------------------------- immutable ------------------------------- */

    uint64  public immutable CHALLENGE_WINDOW;       // seconds
    bytes32 public immutable DOMAIN_SEPARATOR;
    bytes32 public constant STATE_TYPEHASH = keccak256(
        "ChannelState(bytes32 channelId,uint256 nonce,uint256 balanceA,uint256 balanceB,uint256 htlcLockedFromA,uint256 htlcLockedFromB,bytes32 lockHash,uint256 lockTimeout)"
    );

    /// @notice Safety-valve deadline for an Open channel that is never co-signed
    ///         closed. After this long past `open()`, with the channel still in
    ///         Status.Open (no cooperative close, no challenge), EITHER party may
    ///         call `closeOpenTimeout` to reclaim their ORIGINAL deposit. This
    ///         removes the "counterparty refuses to ever co-sign" freeze — most
    ///         acute in a one-directional channel (depositB == 0), where B stakes
    ///         nothing and griefing A would otherwise be free and permanent.
    ///         Deliberately generous so it can never pre-empt a channel that is
    ///         still being used off-chain; a party owed more than its deposit by a
    ///         later co-signed state simply submits it via `startChallenge` before
    ///         the timeout, which flips the channel out of Status.Open and disables
    ///         this path entirely.
    uint256 public constant CHANNEL_OPEN_TIMEOUT = 30 days;

    /* --------------------------------- types --------------------------------- */

    enum Status { None, Open, Challenged, Finalised }

    struct Channel {
        Status status;
        address partyA;
        address partyB;
        address token;
        uint256 depositA;
        uint256 depositB;
        // Latest acknowledged state (post-challenge or cooperative-close).
        uint256 latestNonce;
        uint256 latestBalanceA;
        uint256 latestBalanceB;
        uint256 latestHtlcLockedFromA;
        uint256 latestHtlcLockedFromB;
        bytes32 latestLockHash;
        uint256 latestLockTimeout;
        // Challenge window bookkeeping.
        uint64 challengeStartedAt;
        // HTLC settlement.
        bytes32 settledHtlcLockHash;        // last revealed preimage's hash
        bool    htlcResolved;
        // Open-timeout safety valve bookkeeping (appended so the public getter's
        // existing field ordering is preserved).
        uint64  openedAt;                   // block.timestamp at open()
    }

    struct ChannelState {
        bytes32 channelId;
        uint256 nonce;
        uint256 balanceA;
        uint256 balanceB;
        uint256 htlcLockedFromA;
        uint256 htlcLockedFromB;
        bytes32 lockHash;                     // keccak256(preimage); 0 = no active HTLC
        uint256 lockTimeout;                  // unix seconds; 0 if no HTLC
    }

    /* --------------------------------- state -------------------------------- */

    mapping(bytes32 => Channel) public channels;

    /* --------------------------------- events ------------------------------- */

    event Opened(bytes32 indexed channelId, address indexed partyA, address indexed partyB, address token, uint256 depositA, uint256 depositB);
    event CooperativeClosed(bytes32 indexed channelId, uint256 finalBalanceA, uint256 finalBalanceB);
    event ChallengeStarted(bytes32 indexed channelId, address indexed by, uint256 nonce, uint64 deadline);
    event ChallengeOverridden(bytes32 indexed channelId, uint256 newNonce);
    event Finalised(bytes32 indexed channelId, uint256 finalA, uint256 finalB);
    event HtlcSettled(bytes32 indexed channelId, bytes32 indexed lockHash, address indexed beneficiary);
    event HtlcExpired(bytes32 indexed channelId, bytes32 indexed lockHash);
    event Withdrawn(bytes32 indexed channelId, address indexed recipient, uint256 amount);
    event OpenTimeoutClosed(bytes32 indexed channelId, address indexed by, uint256 refundA, uint256 refundB);

    /* --------------------------------- errors ------------------------------- */

    error ZeroAddress();
    error ZeroAmount();
    error ChannelExists();
    error ChannelNotOpen();
    error ChannelNotChallenged();
    error ChannelNotFinalised();
    error NotParticipant();
    error WrongChannelId();
    error BadSignature(string which);
    error NonceNotHigher(uint256 incoming, uint256 latest);
    error WindowOpen();
    error WindowExpired();
    error OpenTimeoutNotReached();
    error BalancesMismatch(uint256 sum, uint256 deposits);
    error HtlcAlreadyResolved();
    error PreimageMismatch();
    error LockExpired();
    error LockNotExpired();
    error TransferFailed();

    /* ----------------------------- constructor ------------------------------ */

    constructor(uint64 challengeWindow) {
        CHALLENGE_WINDOW = challengeWindow;
        DOMAIN_SEPARATOR = keccak256(abi.encode(
            keccak256("EIP712Domain(string name,string version,uint256 chainId,address verifyingContract)"),
            keccak256("AereStateChannels"),
            keccak256("1"),
            block.chainid,
            address(this)
        ));
    }

    /* --------------------------------- open --------------------------------- */

    /// @notice Opens a channel between (partyA, partyB) with deterministic id.
    ///         BOTH parties' tokens must be `transferFrom`-able to this contract.
    ///         Order of (partyA, partyB) matters for channelId derivation.
    /// @dev    The caller MUST be one of the two channel parties. This blocks a
    ///         third party from using A's (or B's) standing ERC-20 approval to
    ///         force their tokens into a channel they never agreed to — which,
    ///         combined with the every-close-needs-both-signatures rule, could
    ///         otherwise strand those tokens until the open-timeout fallback.
    function open(address partyA, address partyB, address token, uint256 depositA, uint256 depositB) external nonReentrant returns (bytes32 channelId) {
        if (partyA == address(0) || partyB == address(0) || token == address(0)) revert ZeroAddress();
        if (partyA == partyB) revert ZeroAddress();
        if (msg.sender != partyA && msg.sender != partyB) revert NotParticipant();
        channelId = keccak256(abi.encode(partyA, partyB, token, block.timestamp, block.number, msg.sender));
        if (channels[channelId].status != Status.None) revert ChannelExists();
        if (depositA + depositB == 0) revert ZeroAmount();

        if (depositA > 0 && !IERC20(token).transferFrom(partyA, address(this), depositA)) revert TransferFailed();
        if (depositB > 0 && !IERC20(token).transferFrom(partyB, address(this), depositB)) revert TransferFailed();

        Channel storage c = channels[channelId];
        c.status = Status.Open;
        c.partyA = partyA;
        c.partyB = partyB;
        c.token = token;
        c.depositA = depositA;
        c.depositB = depositB;
        // Initial state: balances equal deposits, no HTLC, nonce=0.
        c.latestBalanceA = depositA;
        c.latestBalanceB = depositB;
        c.openedAt = uint64(block.timestamp);

        emit Opened(channelId, partyA, partyB, token, depositA, depositB);
    }

    /* --------------------------- cooperative close --------------------------- */

    function cooperativeClose(
        ChannelState calldata state,
        bytes calldata sigA,
        bytes calldata sigB
    ) external nonReentrant {
        Channel storage c = channels[state.channelId];
        if (c.status != Status.Open) revert ChannelNotOpen();
        if (state.channelId == bytes32(0)) revert WrongChannelId();
        _checkBalances(c, state);
        bytes32 digest = _hashState(state);
        if (digest.recover(sigA) != c.partyA) revert BadSignature("A");
        if (digest.recover(sigB) != c.partyB) revert BadSignature("B");

        c.status = Status.Finalised;
        c.latestNonce = state.nonce;
        c.latestBalanceA = state.balanceA + state.htlcLockedFromA + state.htlcLockedFromB; // unlock HTLC to A on coop close
        c.latestBalanceB = state.balanceB; // simplified — see NOTE below
        // NOTE on cooperative close: parties agree on final off-chain.
        // HTLC tokens go back to their funder unless the state already
        // resolved them (state.lockHash == 0). For real production,
        // parties typically settle HTLCs before cooperative close.
        if (state.lockHash != bytes32(0)) {
            // unresolved HTLC at coop close — refund to funder (whoever
            // contributed htlcLockedFromX is credited that X amount).
            c.latestBalanceA = state.balanceA + state.htlcLockedFromA;
            c.latestBalanceB = state.balanceB + state.htlcLockedFromB;
        } else {
            c.latestBalanceA = state.balanceA;
            c.latestBalanceB = state.balanceB;
        }

        emit CooperativeClosed(state.channelId, c.latestBalanceA, c.latestBalanceB);
    }

    /* --------------------------- open-timeout close -------------------------- */

    /// @notice Safety valve against a non-cooperating counterparty. If a channel
    ///         is still Status.Open CHANNEL_OPEN_TIMEOUT after it was opened —
    ///         meaning NO co-signed state was ever brought on-chain (no
    ///         cooperativeClose, no startChallenge) — EITHER party may finalise
    ///         it here and reclaim their ORIGINAL deposit, then `withdraw`.
    ///
    ///         Without this, every close path requires BOTH signatures, so a
    ///         counterparty who simply never co-signs any closing state can
    ///         permanently freeze the other party's funds (costless in a
    ///         one-directional channel where depositB == 0).
    ///
    /// @dev    SECURITY — this can never override a legitimately-signed newer
    ///         state:
    ///           * It only runs while `status == Status.Open`. The instant any
    ///             co-signed state is submitted, the channel leaves Open
    ///             (cooperativeClose -> Finalised, startChallenge -> Challenged),
    ///             so this function reverts with ChannelNotOpen thereafter.
    ///           * While Open, `latestBalanceA/B` still equal the deposits set in
    ///             `open()` (no on-chain state ever mutated them), so refunding
    ///             the deposits is the only provable allocation.
    ///           * A party owed more than its deposit by a later off-chain state
    ///             has the full CHANNEL_OPEN_TIMEOUT window to enforce it via
    ///             `startChallenge`; doing so moves the channel to Challenged and
    ///             permanently disables this path.
    function closeOpenTimeout(bytes32 channelId) external nonReentrant {
        Channel storage c = channels[channelId];
        if (c.status != Status.Open) revert ChannelNotOpen();
        if (msg.sender != c.partyA && msg.sender != c.partyB) revert NotParticipant();
        if (block.timestamp < uint256(c.openedAt) + CHANNEL_OPEN_TIMEOUT) revert OpenTimeoutNotReached();

        // No co-signed state was ever recorded: refund the original deposits.
        // (latestBalanceA/B already equal depositA/depositB from open(); set
        // explicitly for clarity and defence in depth.)
        c.status = Status.Finalised;
        c.latestBalanceA = c.depositA;
        c.latestBalanceB = c.depositB;

        emit OpenTimeoutClosed(channelId, msg.sender, c.depositA, c.depositB);
    }

    /* ------------------------------ challenge close -------------------------- */

    /// @notice Submit the latest signed state — starts the challenge window.
    ///         Counterparty has CHALLENGE_WINDOW seconds to submit a higher
    ///         nonce. After window expires, anyone calls finalize.
    function startChallenge(
        ChannelState calldata state,
        bytes calldata sigA,
        bytes calldata sigB
    ) external nonReentrant {
        Channel storage c = channels[state.channelId];
        if (c.status != Status.Open) revert ChannelNotOpen();
        _checkBalances(c, state);
        bytes32 digest = _hashState(state);
        if (digest.recover(sigA) != c.partyA) revert BadSignature("A");
        if (digest.recover(sigB) != c.partyB) revert BadSignature("B");

        c.status = Status.Challenged;
        c.latestNonce = state.nonce;
        c.latestBalanceA = state.balanceA;
        c.latestBalanceB = state.balanceB;
        c.latestHtlcLockedFromA = state.htlcLockedFromA;
        c.latestHtlcLockedFromB = state.htlcLockedFromB;
        c.latestLockHash = state.lockHash;
        c.latestLockTimeout = state.lockTimeout;
        c.challengeStartedAt = uint64(block.timestamp);

        emit ChallengeStarted(state.channelId, msg.sender, state.nonce, uint64(block.timestamp + CHALLENGE_WINDOW));
    }

    /// @notice Override a pending challenge with a higher-nonce state.
    function overrideChallenge(
        ChannelState calldata state,
        bytes calldata sigA,
        bytes calldata sigB
    ) external nonReentrant {
        Channel storage c = channels[state.channelId];
        if (c.status != Status.Challenged) revert ChannelNotChallenged();
        if (block.timestamp >= uint256(c.challengeStartedAt) + CHALLENGE_WINDOW) revert WindowExpired();
        if (state.nonce <= c.latestNonce) revert NonceNotHigher(state.nonce, c.latestNonce);
        _checkBalances(c, state);
        bytes32 digest = _hashState(state);
        if (digest.recover(sigA) != c.partyA) revert BadSignature("A");
        if (digest.recover(sigB) != c.partyB) revert BadSignature("B");

        c.latestNonce = state.nonce;
        c.latestBalanceA = state.balanceA;
        c.latestBalanceB = state.balanceB;
        c.latestHtlcLockedFromA = state.htlcLockedFromA;
        c.latestHtlcLockedFromB = state.htlcLockedFromB;
        c.latestLockHash = state.lockHash;
        c.latestLockTimeout = state.lockTimeout;

        emit ChallengeOverridden(state.channelId, state.nonce);
    }

    /// @notice After window expires, anyone can finalise. Resolves HTLC:
    ///         if lockHash is set and `settleHtlc` was called within
    ///         lockTimeout, HTLC tokens go to the receiver (partyB if
    ///         funded by A, partyA if funded by B). Otherwise (no
    ///         settle or settle was late), HTLC tokens return to funder.
    function finalize(bytes32 channelId) external nonReentrant {
        Channel storage c = channels[channelId];
        if (c.status != Status.Challenged) revert ChannelNotChallenged();
        if (block.timestamp < uint256(c.challengeStartedAt) + CHALLENGE_WINDOW) revert WindowOpen();

        // Resolve HTLC.
        uint256 htlcToA;
        uint256 htlcToB;
        if (c.latestLockHash != bytes32(0)) {
            // ROUND-3 FIX: previously this check also required
            // `block.timestamp <= latestLockTimeout + CHALLENGE_WINDOW`, which
            // penalised a receiver who legitimately revealed the preimage
            // inside lockTimeout (settleHtlc already enforces that) whenever
            // finalize was called late. settleHtlc is the authoritative
            // timeliness gate; finalize must honour its result regardless of
            // when finalize itself runs.
            bool settled = c.htlcResolved && c.settledHtlcLockHash == c.latestLockHash;
            if (settled) {
                // Funder loses, receiver gains. Convention: A funds → B
                // gets, B funds → A gets.
                htlcToB += c.latestHtlcLockedFromA;
                htlcToA += c.latestHtlcLockedFromB;
            } else {
                // Refund to funder.
                htlcToA += c.latestHtlcLockedFromA;
                htlcToB += c.latestHtlcLockedFromB;
            }
        }
        c.status = Status.Finalised;
        c.latestBalanceA += htlcToA;
        c.latestBalanceB += htlcToB;

        emit Finalised(channelId, c.latestBalanceA, c.latestBalanceB);
    }

    /* ----------------------------- HTLC ----------------------------- */

    /// @notice Anyone with the preimage can settle an HTLC during the
    ///         challenge window. Records on-chain that the lock was opened.
    function settleHtlc(bytes32 channelId, bytes calldata preimage) external {
        Channel storage c = channels[channelId];
        if (c.status != Status.Challenged) revert ChannelNotChallenged();
        if (c.htlcResolved) revert HtlcAlreadyResolved();
        if (block.timestamp > c.latestLockTimeout) revert LockExpired();
        bytes32 hash = keccak256(preimage);
        if (hash != c.latestLockHash) revert PreimageMismatch();
        c.htlcResolved = true;
        c.settledHtlcLockHash = hash;
        emit HtlcSettled(channelId, hash, msg.sender);
    }

    /* ------------------------------ withdraw -------------------------------- */

    function withdraw(bytes32 channelId) external nonReentrant {
        Channel storage c = channels[channelId];
        if (c.status != Status.Finalised) revert ChannelNotFinalised();
        if (msg.sender != c.partyA && msg.sender != c.partyB) revert NotParticipant();
        uint256 amount;
        if (msg.sender == c.partyA) {
            amount = c.latestBalanceA;
            c.latestBalanceA = 0;
        } else {
            amount = c.latestBalanceB;
            c.latestBalanceB = 0;
        }
        if (amount == 0) return;
        if (!IERC20(c.token).transfer(msg.sender, amount)) revert TransferFailed();
        emit Withdrawn(channelId, msg.sender, amount);
    }

    /* ------------------------------- internal ------------------------------- */

    function _hashState(ChannelState calldata s) internal view returns (bytes32) {
        bytes32 structHash = keccak256(abi.encode(
            STATE_TYPEHASH, s.channelId, s.nonce, s.balanceA, s.balanceB,
            s.htlcLockedFromA, s.htlcLockedFromB, s.lockHash, s.lockTimeout
        ));
        return keccak256(abi.encodePacked("\x19\x01", DOMAIN_SEPARATOR, structHash));
    }

    function _checkBalances(Channel storage c, ChannelState calldata s) internal view {
        uint256 total = s.balanceA + s.balanceB + s.htlcLockedFromA + s.htlcLockedFromB;
        uint256 deposits = c.depositA + c.depositB;
        if (total != deposits) revert BalancesMismatch(total, deposits);
    }
}
