// SPDX-License-Identifier: MIT
pragma solidity ^0.8.20;

import "@openzeppelin/contracts/access/AccessControl.sol";
import "@openzeppelin/contracts/security/ReentrancyGuard.sol";
import "@openzeppelin/contracts/utils/cryptography/ECDSA.sol";

contract RiceChainAuth is AccessControl, ReentrancyGuard {
    using ECDSA for bytes32;

    bytes32 public constant ADMIN_ROLE = keccak256("ADMIN_ROLE");
    bytes32 public constant CH_ROLE = keccak256("CH_ROLE");

    struct Device {
        address owner;
        bytes32 deviceId;
        bool isRegistered;
    }

    struct Stakeholder {
        bytes32 id;
        bytes32 passwordHash;
        bytes32 ePsw;
        bytes32 eId;
        uint256 r1;
        uint256 r2;
        uint256 timestamp;
        bytes32 phi1;
        bytes32 phi2;
        bytes32 phi3;
    }

    struct MutualAuthData {
        bytes32 ePsw_star;
        bytes32 r1_star;
        bytes32 phi2_star;
        bytes32 beta1;
        bytes32 beta2;
        bytes32 beta3;
        bytes32 gamma1;
        bytes32 gamma2;
        bytes32 gamma3;
        uint256 r4;
        uint256 r5;
        bytes32 timestampT3;
        uint256 timestampT4;
        uint256 timestampT5;
        uint256 timestampT6;
        bytes32 omega;
        bytes32 rho1;
        bytes32 rho2;
        bytes32 sigma1;
        bytes32 sigma2;
        bytes32 beta1_star;
        bytes32 beta2_star;
        bytes32 phi1_star;
        bytes32 gamma1_star;
        bytes32 r3_r4_T3;
        bytes32 r5_T4;
        bytes32 r5_r4_T5;
        bytes32 rho1_star;
        bytes32 sigma1_star;
    }

    mapping(address => bytes32) public clusterheads;
    mapping(bytes32 => Device) public devices;
    mapping(address => Stakeholder) public stakeholderData;
    mapping(bytes32 => bytes32) public chShares;

    event ClusterHeadAdded(address indexed stakeholder, bytes32 role);
    event DeviceRegistered(address indexed owner, bytes32 indexed deviceId);
    event AuthenticationSuccess(address indexed owner, bytes32 indexed deviceId);
    event PasswordReplaced(address indexed stakeholder, address indexed Admin, bytes32 newPasswordHash);
    event Signatures(bytes32 Token1, uint256 TimeNow);
    event LoginSuccessful(address chEA, bytes32 beta2, bytes32 beta3, bytes32 r3);
    event ShareWithAdmin(bytes32 omega, bytes32 rho1, bytes32 rho2, uint256 r4, uint256 r5);
    event ShareWithStakeholder(bytes32 gamma1, bytes32 sigma1, bytes32 sigma2);
    event RequestForNewPswd(bytes32 ePsw_new, bytes32 newShPsw, bytes32 mobile_id);
    event DebugValues(bytes32 value1, bytes32 value2, bytes32 value3);

    constructor(address admin) {
        _grantRole(DEFAULT_ADMIN_ROLE, admin);
        _grantRole(ADMIN_ROLE, admin);
    }

    function getSkey(address chEA) public pure returns (bytes32) {
        return keccak256(abi.encodePacked(chEA));
    }

    function addClusterHead(address chEA, bytes32 role, bytes32 skey) public onlyRole(ADMIN_ROLE) {
        require(role == CH_ROLE, "Invalid role");
        _grantRole(role, chEA);
        clusterheads[chEA] = role;

        bytes32 chShare = keccak256(abi.encodePacked(bytes32(uint256(uint160(chEA))), skey));

        chShares[bytes32(uint256(uint160(chEA)))] = chShare;

        emit ClusterHeadAdded(chEA, role);
    }

    function getDeviceId(string memory deviceName) public pure returns (bytes32) {
        return keccak256(abi.encodePacked(deviceName));
    }

    function registerDevice(bytes32 deviceId) public nonReentrant {
        require(hasRole(clusterheads[msg.sender], msg.sender), "Caller is not a registered cluster head");
        require(!devices[deviceId].isRegistered, "Device already registered");

        devices[deviceId] = Device({
            owner: msg.sender,
            deviceId: deviceId,
            isRegistered: true
        });

        emit DeviceRegistered(msg.sender, deviceId);
    }

    function getMobileDetails(string memory StakeholderRoleName, string memory MobileName, address chEA) public view returns (bytes32, bytes32, bytes32, uint256, uint256) {
        bytes32 shId = keccak256(abi.encodePacked(StakeholderRoleName));
        bytes32 shPsw = keccak256(abi.encodePacked(StakeholderRoleName, chEA, block.prevrandao));
        bytes32 mobile_id = keccak256(abi.encodePacked(MobileName));
        uint256 r1Time = block.timestamp;
        uint256 r1 = uint256(keccak256(abi.encodePacked(r1Time, msg.sender)));

        return (shId, shPsw, mobile_id, r1, r1Time);
    }

    function registerStakeholder(bytes32 shId, bytes32 shPsw, uint256 r1, bytes32 mobile_id) public nonReentrant {
        bytes32 ePsw = keccak256(abi.encodePacked(bytes32(r1) ^ shPsw ^ mobile_id));
        uint256 r2 = uint256(keccak256(abi.encodePacked(block.timestamp, msg.sender)));

        bytes32 phi1 = keccak256(abi.encodePacked(ePsw, block.timestamp));
        bytes32 phi2 = keccak256(abi.encodePacked(ePsw, msg.sender));
        bytes32 phi3 = keccak256(abi.encodePacked(phi1, r2, msg.sender)) ^ keccak256(abi.encodePacked(ePsw, block.timestamp));

        stakeholderData[msg.sender] = Stakeholder({
            id: shId,
            passwordHash: shPsw,
            ePsw: ePsw,
            eId: bytes32(0),
            r1: r1,
            r2: r2,
            timestamp: block.timestamp,
            phi1: phi1,
            phi2: phi2,
            phi3: phi3
        });

        uint256 TimeNow = block.timestamp;
        bytes32 Token1 = keccak256(abi.encodePacked(shId, msg.sender, mobile_id));
        emit Signatures(Token1, TimeNow);
    }

    function authenticateStakeholder(bytes32 mobile_id, bytes32 Token1, bytes32 shId, uint256 TimeNow) public nonReentrant returns (bool) {
        Stakeholder storage sh = stakeholderData[msg.sender];
        require(sh.id != bytes32(0), "Stakeholder not registered");

        bytes32 Token2 = keccak256(abi.encodePacked(shId, msg.sender, mobile_id));
        require(Token2 == Token1 && block.timestamp > TimeNow, "Invalid Signatures entered");

        emit AuthenticationSuccess(msg.sender, mobile_id);
        return true;
    }

    function loginUser(
        bytes32 shId,
        bytes32 shPsw,
        uint256 r1Time,
        bytes32 mobile_id
    ) public view returns (bool, bytes32, bytes32, bytes32, uint256) {
        Stakeholder storage sh = stakeholderData[msg.sender];
        require(sh.id == shId, "Invalid SH ID");

        uint256 r1_star = uint256(keccak256(abi.encodePacked(r1Time, msg.sender)));
        require(r1_star == sh.r1, "Invalid r1");

        bytes32 ePsw_star = keccak256(abi.encodePacked(bytes32(r1_star) ^ shPsw ^ mobile_id));
        bytes32 phi2_star = keccak256(abi.encodePacked(ePsw_star, msg.sender));

        require(phi2_star == sh.phi2, "Invalid credentials");

        uint256 r3 = uint256(keccak256(abi.encodePacked(block.timestamp, msg.sender)));
        uint256 timestampT2 = block.timestamp;
        bytes32 beta1 = sh.phi3 ^ keccak256(abi.encodePacked(ePsw_star, sh.timestamp));
        bytes32 beta2 = keccak256(abi.encodePacked(timestampT2, r3, beta1, msg.sender));
        bytes32 beta3 = keccak256(abi.encodePacked(r3, timestampT2)) ^ beta1;

        return (true, beta2, beta3, bytes32(r3), timestampT2);
    }

    function computeInitialValues(
        Stakeholder storage sh,
        bytes32 shPsw,
        uint256 r1Time,
        bytes32 mobile_id
    ) internal view returns (bytes32, bytes32, bytes32) {
        uint256 r1_star = uint256(keccak256(abi.encodePacked(r1Time, msg.sender)));
        bytes32 ePsw_star = keccak256(abi.encodePacked(bytes32(r1_star) ^ shPsw ^ mobile_id));
        bytes32 phi2_star = keccak256(abi.encodePacked(ePsw_star, msg.sender));
        require(phi2_star == sh.phi2, "Invalid credentials");
        return (ePsw_star, bytes32(r1_star), phi2_star);
    }

    function computeSessionValues(
        Stakeholder storage sh,
        bytes32 ePsw_star,
        uint256 timestampT2,
        bytes32 r3
    ) internal view returns (bytes32, bytes32, bytes32) {
        bytes32 beta1 = sh.phi3 ^ keccak256(abi.encodePacked(ePsw_star, sh.timestamp));
        bytes32 beta2 = keccak256(abi.encodePacked(timestampT2, r3, beta1, msg.sender));
        bytes32 beta3 = keccak256(abi.encodePacked(r3, timestampT2)) ^ beta1;
        return (beta1, beta2, beta3);
    }

    function computeGammaValues(
        bytes32 chEAPacked,
        bytes32 shPhi1,
        bytes32 chShare,
        uint256 r4,
        bytes32 timestampT3,
        bytes32 r3
    ) internal pure returns (bytes32, bytes32, bytes32) {
        bytes32 gamma1 = keccak256(abi.encodePacked(chEAPacked, shPhi1, chShare, bytes32(r4), timestampT3));
        bytes32 gamma2 = keccak256(abi.encodePacked(r3, timestampT3, bytes32(r4))) ^ chShare;
        bytes32 gamma3 = shPhi1 ^ keccak256(abi.encodePacked(chEAPacked, keccak256(abi.encodePacked(bytes32(r4))), r3));
        return (gamma1, gamma2, gamma3);
    }

    function mutualAuthenticate(
        bytes32 shId,
        bytes32 shPsw,
        address chEA,
        bytes32 mobile_id,
        bytes32 r3,
        uint256 timestampT2,
        uint256 r1Time,
        bytes32 deviceId
    ) public nonReentrant returns (bool) {
        Stakeholder storage sh = stakeholderData[msg.sender];
        require(sh.id == shId, "Invalid SH ID");

        bytes32 chEAPacked = keccak256(abi.encodePacked(chEA));
        Device storage device = devices[deviceId];
        require(device.isRegistered, "Device is not registered");

        MutualAuthData memory data;

        (data.ePsw_star, data.r1_star, data.phi2_star) = computeInitialValues(sh, shPsw, r1Time, mobile_id);
        (data.beta1, data.beta2, data.beta3) = computeSessionValues(sh, data.ePsw_star, timestampT2, r3);

        return mutualAuthenticatePart2(sh, chEAPacked, r3, timestampT2, device, data);
    }

    function mutualAuthenticatePart2(
        Stakeholder storage sh,
        bytes32 chEAPacked,
        bytes32 r3,
        uint256 timestampT2,
        Device storage device,
        MutualAuthData memory data
    ) internal returns (bool) {
        data.timestampT3 = bytes32(block.timestamp);
        require(uint256(data.timestampT3) - timestampT2 < 10 minutes, "Session timed out");

        //data.beta1_star = sh.phi3 ^ keccak256(abi.encodePacked(data.ePsw_star, sh.timestamp));
        //data.beta2_star = keccak256(abi.encodePacked(r3, timestampT2, data.beta1_star, msg.sender));
        //require(data.beta2_star == data.beta2, "Mutual authentication-2 failed");

        data.r4 = uint256(keccak256(abi.encodePacked(block.timestamp, msg.sender)));
        (data.gamma1, data.gamma2, data.gamma3) = computeGammaValues(
            chEAPacked,
            sh.phi1,
            chShares[chEAPacked],
            data.r4,
            data.timestampT3,
            r3
        );

        return mutualAuthenticatePart3(sh, chEAPacked, r3, device, data);
    }

    function mutualAuthenticatePart3(
        Stakeholder storage sh,
        bytes32 chEAPacked,
        bytes32 r3,
        Device storage device,
        MutualAuthData memory data
    ) internal returns (bool) {
        data.r3_r4_T3 = data.gamma2 ^ chShares[chEAPacked];
        data.timestampT4 = block.timestamp;
        require(data.timestampT4 - uint256(data.timestampT3) < 10 minutes, "Session timed out");

        //data.phi1_star = data.gamma3 ^ keccak256(abi.encodePacked(chEAPacked, data.r4, r3));
        //data.gamma1_star = keccak256(abi.encodePacked(chEAPacked, data.phi1_star, chShares[chEAPacked], bytes32(data.r4), data.timestampT3));
        //require(data.gamma1 == data.gamma1_star, "Mutual authentication-3 failed");

        data.r5 = uint256(keccak256(abi.encodePacked(block.timestamp, msg.sender)));
        data.omega = keccak256(abi.encodePacked(data.phi1_star, r3, data.r4, data.r5));
        data.rho1 = keccak256(abi.encodePacked(data.timestampT4, data.r5, chEAPacked, chShares[chEAPacked], data.timestampT3, data.omega));
        data.rho2 = keccak256(abi.encodePacked(data.r5, data.timestampT4)) ^ bytes32(data.r4);

        emit ShareWithAdmin(data.omega, data.rho1, data.rho2, data.r4, data.r5);

        return mutualAuthenticatePart4(sh, chEAPacked, r3, device, data);
    }

    function mutualAuthenticatePart4(
        Stakeholder storage sh,
        bytes32 chEAPacked,
        bytes32 r3,
        Device storage device,
        MutualAuthData memory data
    ) internal returns (bool) {
        data.r5_T4 = data.rho2 ^ bytes32(data.r4);
        data.timestampT5 = block.timestamp;
        require(data.timestampT5 - data.timestampT4 < 10 minutes, "Session timed out");

        data.rho1_star = keccak256(abi.encodePacked(data.timestampT4, data.r5, chEAPacked, chShares[chEAPacked], bytes32(data.timestampT4), data.omega));
        require(data.rho1 == data.rho1_star, "Mutual authentication-4 failed");

        data.sigma1 = keccak256(abi.encodePacked(data.omega, sh.phi1, bytes32(data.r4), data.timestampT5, data.gamma1));
        data.sigma2 = keccak256(abi.encodePacked(data.r5, bytes32(data.r4), data.timestampT5)) ^ bytes32(r3);
        emit ShareWithStakeholder(data.gamma1, data.sigma1, data.sigma2);

        data.r5_r4_T5 = data.sigma2 ^ bytes32(r3);
        data.timestampT6 = block.timestamp;
        require(data.timestampT6 - data.timestampT5 < 10 minutes, "Session timed out");

        data.sigma1_star = keccak256(abi.encodePacked(data.omega, sh.phi1, bytes32(data.r4), data.timestampT5, data.gamma1));
        require(data.sigma1 == data.sigma1_star, "Mutual authentication-5 failed");

        emit AuthenticationSuccess(device.owner, chEAPacked);

        return true;
    }

    function getNewPswd(string memory StakeholderRoleName, address chEA) public view returns (bytes32) {
        //bytes32 shId = keccak256(abi.encodePacked(StakeholderRoleName));
        bytes32 newShPsw = keccak256(abi.encodePacked(StakeholderRoleName, chEA, block.prevrandao, block.timestamp));
        //bytes32 mobile_id = keccak256(abi.encodePacked(MobileName));
        //uint256 r1Time = block.timestamp;
        //uint256 r1 = uint256(keccak256(abi.encodePacked(r1Time, msg.sender)));

        return (newShPsw);
    }


    function replacePassword(bytes32 shId, bytes32 shPsw, bytes32 newShPsw, bytes32 mobile_id, uint256 r1Time, address chEA) public nonReentrant {
        Stakeholder storage sh = stakeholderData[msg.sender];
        require(sh.id == shId, "Invalid SH ID");

        bytes32 r1Bytes32 = bytes32(sh.r1);

        uint256 r1_star = uint256(keccak256(abi.encodePacked(r1Time, msg.sender)));
        //bytes32 ePsw = keccak256(abi.encodePacked(bytes32(r1) ^ shPsw ^ mobile_id));
        bytes32 ePsw_star = keccak256(abi.encodePacked(bytes32(r1_star) ^ shPsw ^ mobile_id));
        
        //bytes32 r1_star = sh.eId ^ keccak256(abi.encodePacked(shPsw, shId));
        //bytes32 ePsw_star = keccak256(abi.encodePacked(r1_star ^ shPsw ^ mobile_id));
        bytes32 ePsw_new = keccak256(abi.encodePacked(r1Bytes32 ^ newShPsw ^ mobile_id));

        require(ePsw_star == sh.ePsw, "Invalid password");
        emit RequestForNewPswd(ePsw_new, newShPsw, mobile_id);

        bytes32 phi1_star = keccak256(abi.encodePacked(ePsw_new, sh.timestamp));
        bytes32 phi2_star = keccak256(abi.encodePacked(ePsw_new, msg.sender));
        bytes32 phi3_star = keccak256(abi.encodePacked(phi1_star, sh.r2, msg.sender)) ^ keccak256(abi.encodePacked(ePsw_new, sh.timestamp));

        sh.passwordHash = newShPsw;
        sh.ePsw = ePsw_new;
        sh.phi1 = phi1_star;
        sh.phi2 = phi2_star;
        sh.phi3 = phi3_star;

        emit PasswordReplaced(chEA, msg.sender, newShPsw);
    }
}