import {
  Address,
  encodeAbiParameters,
  Hex,
  hexToBigInt,
  isAddressEqual,
  keccak256,
  zeroAddress,
  zeroHash,
} from 'viem';
import {
  PERMIT2_MAP,
  TOKEN_DISTRIBUTOR,
  TokenDistributeSubmitRequest,
  TokenDistributeDepositRequest,
  TokenDistributorAbi,
} from '@prex0/prex-structs';
import {
  DistributeRequest,
  DistributionRequestStatus,
  PagingOptions,
  PrexUser,
  TokenDistributeRequestEntity,
} from '../types';
import { PrexSigner } from '../core/sign';
import { PrexApiService } from '../api';
import { PrexClient } from '../prex-client';
import {
  encryptByTmpSecret,
  generateKey,
  generateSecret,
  decryptByTmpSecret,
  generateShortCode,
} from '../utils/tmp-secret';
import { privateKeyToAccount, signMessage } from 'viem/accounts';
import {
  getDistributeRequest,
  getDistributionInfoMap,
} from '../evm-api/get-distribute-request';
import { normalizeErrorFn, PrexSDKError } from '../errors';
import { PrexStorage } from '../storage/PrexStorage';
import { DistributeActionInterface } from '../interfaces/prex-client-interface';
import { queryTokenDistributeRequests } from '../graph/distributions';
import { privateKeyToAddress } from 'viem/accounts';
import {
  DistributionSecret,
  parseDistributionSecret,
  serializeDistributionSecret,
} from '../utils/distribution-secrets';
import { retryFetch } from '../utils/retry-fetch';
import { hexToBase64 } from '../utils/base64';

const DEADLINE_UNTIL = 24 * 60 * 60;

type DistributionSecretParams = DistributionSecret & {
  distributionId: Hex;
  recipient?: Address;
};

export interface RegisterServerGeneratedSecretInput {
  distributionId: string;
  secret: string;
  encodedRequest: Hex;
  chainId: number;
  shortCode: string;
}

export interface ServerGeneratedSecretInput {
  distributionId: string;
  recipient: Address;
  coordinate: Hex;
  deadline: string;
  shortCode: string;
}

export class DistributeAction implements DistributeActionInterface {
  constructor(
    private client: PrexClient,
    private apiService: PrexApiService,
    private storage: PrexStorage,
    private user?: PrexUser,
    private signer?: PrexSigner
  ) {}

  async submit({
    token,
    amount,
    amountPerWithdrawal,
    maxAmountPerAddress,
    expiry,
    coolTime,
    name,
    coordinate,
    sender: senderParam,
    registerPath,
  }: {
    token: Address;
    amount: bigint;
    amountPerWithdrawal: bigint;
    maxAmountPerAddress?: bigint;
    expiry: bigint;
    coolTime?: bigint;
    name: string;
    coordinate?: Hex;
    enableEncryption?: boolean;
    sender?: Address;
    registerPath?: string;
  }) {
    if (!this.user || !this.signer) {
      throw new PrexSDKError('not_initialized', 'User not initialized');
    }

    const sender = senderParam || this.user.address;

    const chainId = this.signer.chainId;

    const { secret, publicKey } = generateSecret();

    const deadline = this.getDeadline();
    const finalCoolTime =
      coolTime !== undefined ? coolTime : 60n * 60n * 24n * 365n;
    const finalMaxAmountPerAddress =
      maxAmountPerAddress !== undefined
        ? maxAmountPerAddress
        : amountPerWithdrawal;

    const submitRequest = new TokenDistributeSubmitRequest(
      {
        token,
        amount,
        publicKey,
        amountPerWithdrawal,
        expiry,
        name,
        dispatcher: TOKEN_DISTRIBUTOR,
        sender: sender,
        nonce: await this.client.updatePermit2Nonce(TOKEN_DISTRIBUTOR, sender),
        deadline,
        cooltime: finalCoolTime,
        maxAmountPerAddress: finalMaxAmountPerAddress,
        additionalValidator: zeroAddress,
        additionalData: '0x',
        coordinate: coordinate
          ? await encryptByTmpSecret(secret, coordinate)
          : zeroHash,
      },
      chainId,
      PERMIT2_MAP[chainId]
    );

    const { domain, types, message, primaryType } = submitRequest.permitData();

    const signature = await this.signer.signTypedData(
      {
        domain,
        types,
        message,
        primaryType,
      },
      sender
    );

    const id = submitRequest.hash().toLowerCase();

    // encode the secret to base64
    const encodedSecret = serializeDistributionSecret({
      type: 'secret',
      secret,
    });

    await this.setSecret(id, encodedSecret);

    const shortCode = DistributeAction.generateShortCode();

    await this.setShortCode(id, shortCode);

    if (registerPath) {
      // if registerPath is provided, we need to register the secret to server
      const { success } = await this._callAPIRegisterRealTimeSecret(
        registerPath,
        {
          distributionId: id,
          secret: encodedSecret,
          encodedRequest: submitRequest.serialize(),
          shortCode,
          chainId,
        }
      );

      if (!success) {
        throw new Error('Failed to register real time secret to server');
      }
    }

    await this.apiService.distributeSubmit(
      submitRequest.serialize(),
      signature
    );

    await this.client.fetchBalance(token);

    return {
      id,
      secret: encodedSecret,
      shortCode,
    };
  }

  async generateSubSecret(id: Hex) {
    const secret = await this.getSecret(id);

    if (!secret) {
      throw new Error('Secret not found');
    }

    const parsedSecret = parseDistributionSecret(secret);

    if (parsedSecret.type !== 'secret') {
      throw new Error('Sub secret already exists');
    }

    const { secret: subSecret, publicKey: subPublicKey } = generateSecret();

    const expiry = await this.getExpiryOfRequest(id);

    const nonce = hexToBigInt(keccak256(subPublicKey));
    const deadline = expiry;

    const subSig = await DistributeAction.generateWithdrawSignature(
      parsedSecret.secret,
      TOKEN_DISTRIBUTOR,
      nonce,
      deadline,
      subPublicKey
    );

    return {
      id,
      secret: serializeDistributionSecret({
        type: 'pre_generated',
        subSig,
        subSecret,
      }),
    };
  }

  async deposit({
    amount,
    requestId,
    sender: senderParam,
  }: {
    amount: bigint;
    requestId: Hex;
    sender?: Address;
  }) {
    if (!this.user || !this.signer) {
      throw new PrexSDKError('not_initialized', 'User not initialized');
    }

    const sender = senderParam || this.user.address;

    const token = await this.getTokenOfRequest(requestId);

    const deadline = this.getDeadline();

    const depositRequest = new TokenDistributeDepositRequest(
      {
        token: token,
        amount,
        dispatcher: TOKEN_DISTRIBUTOR,
        sender: sender,
        nonce: this.client.getPermit2Nonce(),
        deadline,
        requestId,
      },
      this.apiService.chainId,
      PERMIT2_MAP[this.apiService.chainId]
    );

    const { domain, types, message } = depositRequest.permitData();

    const signature = await this.signer.signTypedData(
      {
        domain,
        types,
        message,
        primaryType: 'PermitWitnessTransferFrom',
      },
      sender
    );

    await this.apiService.distributeDeposit(
      depositRequest.serialize(),
      signature
    );

    await this.client.fetchBalance(token);
  }

  async cancel({ requestId, from }: { requestId: Hex; from?: Address }) {
    await this.client.executeOperation(
      {
        address: TOKEN_DISTRIBUTOR,
        abi: TokenDistributorAbi,
        functionName: 'cancel',
        args: [requestId],
      },
      from
    );
  }

  async getTokenOfRequest(id: Hex): Promise<Address> {
    const request = await getDistributeRequest(
      this.client.evmChainClient,
      TOKEN_DISTRIBUTOR,
      id
    );

    return request[4];
  }
  async getExpiryOfRequest(id: Hex): Promise<bigint> {
    const request = await getDistributeRequest(
      this.client.evmChainClient,
      TOKEN_DISTRIBUTOR,
      id
    );

    return request[7];
  }

  async getCoordinateOfRequest(id: Hex): Promise<Hex> {
    const request = await getDistributeRequest(
      this.client.evmChainClient,
      TOKEN_DISTRIBUTOR,
      id
    );

    return request[12];
  }

  async getRequest(
    id: Hex,
    {
      secret,
      coordinate: coordinateFromServer,
    }: {
      secret?: Hex;
      coordinate?: Hex;
    }
  ): Promise<DistributeRequest> {
    const request = await getDistributeRequest(
      this.client.evmChainClient,
      TOKEN_DISTRIBUTOR,
      id
    );

    const infoMap = this.user
      ? await getDistributionInfoMap(
          this.client.evmChainClient,
          TOKEN_DISTRIBUTOR,
          id,
          this.user.address
        )
      : null;

    function getStatus(status: number): DistributionRequestStatus {
      if (status === 1) {
        return 'PENDING';
      } else if (status === 2) {
        return 'COMPLETED';
      } else if (status === 3) {
        return 'CANCELLED';
      }
      return 'EMPTY';
    }

    const encryptedCoordinate = request[12];
    let coordinate: Hex = zeroHash;

    if (encryptedCoordinate !== zeroHash) {
      if (secret) {
        coordinate = await decryptByTmpSecret(secret, encryptedCoordinate);
      } else if (coordinateFromServer) {
        coordinate = coordinateFromServer;
      } else {
        throw new Error('Coordinate not found');
      }
    } else {
      coordinate = zeroHash;
    }

    return {
      id,
      amount: request[0],
      amountPerWithdrawal: request[1],
      cooltime: request[2],
      maxAmountPerAddress: request[3],
      token: request[4],
      publicKey: request[5],
      sender: request[6],
      expiry: request[7],
      status: getStatus(request[8]),
      name: request[9],
      coordinate: coordinate,
      userStatus: infoMap
        ? {
            lastDistributedAt: infoMap[0],
            amount: infoMap[1],
          }
        : undefined,
    };
  }

  async prepareSecret(_params: {
    path: string;
    auth?: string;
    requestId: Hex;
    shortCode: string;
    recipient?: Address;
  }) {
    if (!this.user) {
      throw new PrexSDKError('not_initialized', 'User not initialized');
    }

    const coordinate = await this.getCoordinateOfRequest(_params.requestId);
    const expiry = await this.getExpiryOfRequest(_params.requestId);

    try {
      const secret = await this._callAPIGenerateRealTimeSecret(_params.path, {
        distributionId: _params.requestId,
        shortCode: _params.shortCode,
        recipient: _params.recipient || this.user.address,
        coordinate,
        deadline: expiry.toString(),
      }, _params.auth);

      return {
        secret: secret.secret,
      };
    } catch (error) {
      const normalizedError = normalizeErrorFn(
        'Failed to generate real time secret'
      )(error);

      throw new PrexSDKError('generate_secret_failed', normalizedError.message);
    }
  }

  async prepareReceive(params: {
    requestId: Hex;
    secret: string;
    recipient?: Address;
  }) {
    const parsedSecret = parseDistributionSecret(params.secret);

    this._prepareReceive({
      ...parsedSecret,
      distributionId: params.requestId,
      recipient: params.recipient,
    });

    return this.getRequest(params.requestId, {
      secret: parsedSecret.type === 'secret' ? parsedSecret.secret : undefined,
      coordinate:
        parsedSecret.type === 'real_time' ? parsedSecret.coordinate : undefined,
    });
  }

  async _prepareReceive(params: DistributionSecretParams) {
    if (!this.user || !this.signer) {
      throw new PrexSDKError('not_initialized', 'User not initialized');
    }

    const expiry = await this.getExpiryOfRequest(params.distributionId);

    const deadline = expiry;

    const recipient = params.recipient || this.user.address;

    if (params.type === 'pre_generated') {
      const subPublicKey = privateKeyToAddress(params.subSecret);
      const nonce = hexToBigInt(keccak256(subPublicKey));

      const sig = await DistributeAction.generateWithdrawSignature(
        params.subSecret,
        TOKEN_DISTRIBUTOR,
        nonce,
        deadline,
        recipient
      );

      await this.storage.setItemToSessionStorage(
        this.getKeyForParams(params.distributionId),
        {
          sig: params.subSig,
          nonce,
          deadline,
          recipient,
          subPublicKey,
          subSig: sig,
        }
      );
    } else if (params.type === 'real_time') {
      const sig = params.sig;
      const nonce = params.nonce;

      await this.storage.setItemToSessionStorage(
        this.getKeyForParams(params.distributionId),
        {
          sig,
          nonce,
          deadline: expiry,
          recipient,
          subPublicKey: zeroAddress,
          subSig: '0x',
        }
      );
    } else if (params.type === 'secret') {
      const nonce = DistributeAction.generateNonce();
      const sig = await DistributeAction.generateWithdrawSignature(
        params.secret,
        TOKEN_DISTRIBUTOR,
        nonce,
        deadline,
        recipient
      );

      await this.storage.setItemToSessionStorage(
        this.getKeyForParams(params.distributionId),
        {
          sig,
          nonce,
          deadline,
          recipient,
          subPublicKey: zeroAddress,
          subSig: '0x',
        }
      );
    } else {
      throw new Error('Invalid secret type');
    }
  }

  async receive({ requestId }: { requestId: Hex }) {
    if (!this.user || !this.signer) {
      throw new PrexSDKError('not_initialized', 'User not initialized');
    }

    const params = await this.storage.getItemFromSessionStorage<{
      sig: Hex;
      nonce: string;
      deadline: string;
      recipient: Address;
      subPublicKey: Hex;
      subSig: Hex;
    }>(this.getKeyForParams(requestId));

    if (!params) {
      throw new Error('Params not found');
    }

    const { sig, nonce, deadline, recipient, subPublicKey, subSig } = params;

    const token = await this.getTokenOfRequest(requestId);

    await this.apiService.distributeWithdraw({
      requestId,
      sig,
      nonce: nonce,
      deadline: deadline.toString(),
      recipient,
      subPublicKey,
      subSig,
    });

    await this.client.fetchBalance(token);
  }

  async queryRequests(
    query?: { user: Address },
    pagingOptions?: PagingOptions
  ) {
    if (!this.user) {
      throw new PrexSDKError('not_initialized', 'User not initialized');
    }

    const requests = await queryTokenDistributeRequests(
      this.apiService,
      query || { user: this.user.address },
      pagingOptions?.offset || 0,
      pagingOptions?.limit || 10
    );

    return Promise.all(
      requests.map(async (request) => ({
        ...request,
        secret: await this.getSecret(request.id),
        shortCode: await this.getShortCode(request.id),
      }))
    ) as Promise<TokenDistributeRequestEntity[]>;
  }

  static generateNonce() {
    const randomKey = generateKey();

    return hexToBigInt(randomKey);
  }

  static generateShortCode() {
    return hexToBase64(generateShortCode());
  }

  static async generateWithdrawSignature(
    secret: Hex,
    dispatcher: Address,
    nonce: bigint,
    deadline: bigint,
    recipient: Address
  ) {
    return await signMessage({
      privateKey: secret,
      message: {
        raw: keccak256(
          encodeAbiParameters(
            [
              { type: 'address' },
              { type: 'uint256' },
              { type: 'uint256' },
              { type: 'address' },
            ],
            [dispatcher, nonce, deadline, recipient]
          )
        ),
      },
    });
  }

  private getDeadline() {
    return BigInt(Math.floor(Date.now() / 1000) + DEADLINE_UNTIL);
  }

  private getKeyForParams(id: Hex) {
    return `distribute-params-${id}`;
  }

  private async getSecret(id: string) {
    return this.storage.getItem<string>('secret-' + id.toLowerCase());
  }

  private async setSecret(id: string, secret: string) {
    this.storage.setItem('secret-' + id.toLowerCase(), secret);
  }

  private async getShortCode(id: string) {
    return this.storage.getItem<string>('short-code-' + id.toLowerCase());
  }

  private async setShortCode(id: string, shortCode: string) {
    this.storage.setItem('short-code-' + id.toLowerCase(), shortCode);
  }

  async _callAPIRegisterRealTimeSecret(
    path: string,
    data: RegisterServerGeneratedSecretInput
  ) {
    const res = await this._post(path, data);

    return res as {
      success: boolean;
    };
  }

  async _callAPIGenerateRealTimeSecret(
    path: string,
    data: ServerGeneratedSecretInput,
    auth?: string
  ) {
    const res = await this._post(path, data, auth);

    return res as {
      secret: Hex;
    };
  }

  async _post(path: string, data: any, auth?: string) {
    const headers = auth ? {
      'Content-Type': 'application/json',
      'Authorization': `Bearer ${auth}`
    } : {
      'Content-Type': 'application/json',
    } as HeadersInit;

    const res = await retryFetch(
      path,
      {
        method: 'POST',
        mode: 'cors',
        headers,
        body: JSON.stringify(data),
      },
      {
        retries: 3,
      }
    );

    if (res.status >= 500 && res.status < 600) {
      throw new Error('Server Error');
    }

    return await res.json();
  }
}

export async function validateDistributionAndSecret(
  registerRequest: RegisterServerGeneratedSecretInput
) {
  const chainId = registerRequest.chainId;

  const request = TokenDistributeSubmitRequest.parse(
    registerRequest.encodedRequest,
    chainId,
    PERMIT2_MAP[chainId]
  );

  const hash = request.hash().toLowerCase();

  if (hash !== registerRequest.distributionId.toLowerCase()) {
    return false;
  }

  const parsedSecret = parseDistributionSecret(registerRequest.secret);

  if (parsedSecret.type !== 'secret') {
    return false;
  }

  const account = privateKeyToAccount(parsedSecret.secret);

  if (!isAddressEqual(account.address, request.params.publicKey)) {
    return false;
  }

  return true;
}

export async function createServerGeneratedSecret(
  secret: string,
  { recipient, coordinate, deadline }: ServerGeneratedSecretInput
) {
  const nonce = DistributeAction.generateNonce();

  const parsedSecret = parseDistributionSecret(secret);

  if (parsedSecret.type !== 'secret') {
    throw new Error('Invalid secret type');
  }

  const sig = await DistributeAction.generateWithdrawSignature(
    parsedSecret.secret,
    TOKEN_DISTRIBUTOR,
    nonce,
    BigInt(deadline),
    recipient
  );

  return serializeDistributionSecret({
    type: 'real_time',
    sig,
    nonce,
    coordinate:
      coordinate === zeroHash
        ? zeroHash
        : await decryptByTmpSecret(parsedSecret.secret, coordinate),
  });
}