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

import {Options, Upgrades} from "openzeppelin-foundry-upgrades/Upgrades.sol";
import {console2} from "forge-std/console2.sol";
import {BankV1} from "../src/BankV1.sol";
import {MockUSDC} from "../src/MockUSDC.sol";
import {DemoScript} from "./lib/DemoScript.sol";

contract UpgradeV2 is DemoScript {
    string internal constant UPGRADE_PENDING_MANIFEST_PATH = "deployments/upgrade-pending.json";

    error UnexpectedVersion(uint256 version);

    struct Snapshot {
        address proxy;
        address implementation;
        address owner;
        address asset;
        bool paused;
        uint256[] balances;
        uint256 liabilities;
        uint256 reserves;
        uint256 surplus;
        uint256 deploymentBlock;
        uint256 version;
    }

    function run() external returns (bool upgraded, address implementation) {
        _requireSupportedChain(block.chainid);
        _printEducationalWarning();
        string memory activePath = vm.envOr("UPGRADE_ACTIVE_MANIFEST_PATH", ACTIVE_MANIFEST_PATH);
        return _run(activePath, _manifestPath(UPGRADE_PENDING_MANIFEST_PATH), vm.envAddress("SCRIPT_SENDER"));
    }

    function _run(string memory activePath, string memory pendingPath, address sender)
        internal
        returns (bool upgraded, address implementation)
    {
        address configuredRecipient;
        if (block.chainid == BASE_SEPOLIA_CHAIN_ID) {
            (, configuredRecipient,) = _requireBaseConfig(sender);
        }
        Manifest memory manifest = _readManifest(activePath, true);
        if (block.chainid == BASE_SEPOLIA_CHAIN_ID) {
            _assertAddress("Presenter", manifest.actors[0].address_, sender);
            _assertAddress("Recipient", manifest.actors[1].address_, configuredRecipient);
        }
        BankV1 bank = BankV1(manifest.proxy);

        address actualImplementation = Upgrades.getImplementationAddress(manifest.proxy);
        _assertAddress("implementation", manifest.implementation, actualImplementation);
        _assertAddress("owner", manifest.owner, bank.owner());
        _assertAddress("asset", manifest.token, address(bank.asset()));
        _assertAddress("SCRIPT_SENDER", bank.owner(), sender);

        uint256 version = bank.contractVersion();
        if (version == 2) {
            implementation = actualImplementation;
            _writeNoopMarker(pendingPath, manifest, vm.getNonce(bank.owner()), block.number);
            console2.log("BankV2 already active; no upgrade broadcast.");
            return (false, implementation);
        }
        if (version != 1) revert UnexpectedVersion(version);

        Snapshot memory before_ = _snapshot(manifest);
        Options memory opts;
        opts.referenceContract = "BankV1.sol:BankV1";
        Upgrades.validateUpgrade("BankV2.sol:BankV2", opts);

        if (block.chainid == ANVIL_CHAIN_ID) {
            (uint256 ownerKey, address derivedOwner) = _deriveLocalActor(block.chainid, 0);
            _assertAddress("local owner", sender, derivedOwner);
            vm.startBroadcast(ownerKey);
        } else {
            vm.startBroadcast(sender);
        }
        Upgrades.upgradeProxy(manifest.proxy, "BankV2.sol:BankV2", "", opts);
        vm.stopBroadcast();

        Snapshot memory after_ = _snapshot(manifest);
        implementation = after_.implementation;
        _requireCode("implementation", implementation);
        if (implementation == before_.implementation) {
            revert UnexpectedAddress("implementation changed", before_.implementation, implementation);
        }
        _assertUint("version", 2, after_.version);
        _assertSnapshotUnchanged(before_, after_);
        _writeUpgradeMarker(pendingPath, manifest, before_, implementation);
        return (true, implementation);
    }

    function _snapshot(Manifest memory manifest) internal view returns (Snapshot memory snapshot) {
        BankV1 bank = BankV1(manifest.proxy);
        uint256 reserves = MockUSDC(manifest.token).balanceOf(manifest.proxy);
        uint256 liabilities = bank.totalLiabilities();
        snapshot.proxy = manifest.proxy;
        snapshot.implementation = Upgrades.getImplementationAddress(manifest.proxy);
        snapshot.owner = bank.owner();
        snapshot.asset = address(bank.asset());
        snapshot.paused = bank.paused();
        snapshot.balances = new uint256[](manifest.actors.length);
        for (uint256 i; i < manifest.actors.length; ++i) {
            snapshot.balances[i] = bank.balanceOf(manifest.actors[i].address_);
        }
        snapshot.liabilities = liabilities;
        snapshot.reserves = reserves;
        snapshot.surplus = reserves - liabilities;
        snapshot.deploymentBlock = manifest.deploymentBlock;
        snapshot.version = bank.contractVersion();
    }

    function _assertSnapshotUnchanged(Snapshot memory before_, Snapshot memory after_) internal pure {
        _assertAddress("proxy", before_.proxy, after_.proxy);
        _assertAddress("owner", before_.owner, after_.owner);
        _assertAddress("asset", before_.asset, after_.asset);
        _assertUint("paused", before_.paused ? 1 : 0, after_.paused ? 1 : 0);
        _assertUint("actor count", before_.balances.length, after_.balances.length);
        for (uint256 i; i < before_.balances.length; ++i) {
            _assertUint("actor balance", before_.balances[i], after_.balances[i]);
        }
        _assertUint("liabilities", before_.liabilities, after_.liabilities);
        _assertUint("reserves", before_.reserves, after_.reserves);
        _assertUint("surplus", before_.surplus, after_.surplus);
        _assertUint("deployment block", before_.deploymentBlock, after_.deploymentBlock);
    }

    function _writeNoopMarker(string memory path, Manifest memory manifest, uint256 ownerNonce, uint256 observedBlock)
        private
    {
        vm.writeJson(
            string.concat(
                '{"mode":"noop","chainId":',
                vm.toString(manifest.chainId),
                ',"observedBlock":',
                vm.toString(observedBlock),
                ',"ownerNonce":',
                vm.toString(ownerNonce),
                ',"proxy":"',
                vm.toString(manifest.proxy),
                '","implementation":"',
                vm.toString(manifest.implementation),
                '"}'
            ),
            path
        );
    }

    function _writeUpgradeMarker(
        string memory path,
        Manifest memory manifest,
        Snapshot memory before_,
        address implementation
    ) private {
        string memory balances = "[";
        for (uint256 i; i < manifest.actors.length; ++i) {
            if (i != 0) balances = string.concat(balances, ",");
            balances = string.concat(
                balances,
                '{"address":"',
                vm.toString(manifest.actors[i].address_),
                '","balance":',
                vm.toString(before_.balances[i]),
                "}"
            );
        }
        balances = string.concat(balances, "]");
        string memory snapshot = string.concat('{"proxy":"', vm.toString(before_.proxy), '"');
        snapshot = string.concat(snapshot, ',"implementation":"', vm.toString(before_.implementation), '"');
        snapshot = string.concat(snapshot, ',"owner":"', vm.toString(before_.owner), '"');
        snapshot = string.concat(snapshot, ',"asset":"', vm.toString(before_.asset), '"');
        snapshot = string.concat(snapshot, ',"paused":', vm.toString(before_.paused));
        snapshot = string.concat(snapshot, ',"balances":', balances);
        snapshot = string.concat(snapshot, ',"liabilities":', vm.toString(before_.liabilities));
        snapshot = string.concat(snapshot, ',"reserves":', vm.toString(before_.reserves));
        snapshot = string.concat(snapshot, ',"surplus":', vm.toString(before_.surplus));
        snapshot = string.concat(snapshot, ',"deploymentBlock":', vm.toString(before_.deploymentBlock));
        snapshot = string.concat(snapshot, ',"version":', vm.toString(before_.version), "}");

        string memory json = string.concat('{"mode":"upgrade","network":"', manifest.network, '"');
        json = string.concat(json, ',"chainId":', vm.toString(manifest.chainId));
        json = string.concat(json, ',"token":"', vm.toString(manifest.token), '"');
        json = string.concat(json, ',"proxy":"', vm.toString(manifest.proxy), '"');
        json = string.concat(json, ',"previousImplementation":"', vm.toString(manifest.implementation), '"');
        json = string.concat(json, ',"implementation":"', vm.toString(implementation), '"');
        json = string.concat(json, ',"owner":"', vm.toString(manifest.owner), '"');
        json = string.concat(json, ',"deploymentBlock":', vm.toString(manifest.deploymentBlock));
        json = string.concat(json, ',"snapshot":', snapshot, "}");
        vm.writeJson(json, path);
    }
}