import { Address, Hex, WalletClient, custom, serializeSignature } from 'viem';
import { encodeSignatureBytes } from '@prex0/prex-structs';
import { sign, signMessage } from 'viem/accounts';
import { PrexUser, SubKeyMessage } from '../types';
import {
  createReplaySafeHash,
  PasskeyProvider,
} from '../providers/passkey-provider';
import { ProviderInterface } from '../providers/base-provider';
import { createWalletClient, encodeAbiParameters, getAddress } from 'viem';
import { keccak256 } from 'viem/utils';
import { arbitrum, arbitrumSepolia, base, optimism } from 'viem/chains';
import { PrexApiService } from '../api';

export class PrexSigner {
  constructor(
    public provider: ProviderInterface,
    public walletClient: WalletClient,
    public chainId: number,
    public myAddress: Address,
    public index: number,
    public passkeyIds?: string[]
  ) {}

  static fromPrexWallet(
    apiService: PrexApiService,
    wallet: PrexUser,
    provider?: ProviderInterface
  ) {
    const finalProvider =
      provider ||
      new PasskeyProvider(
        apiService,
        wallet.address,
        wallet.ownerIndex,
        wallet.passkeys.map((passkey) => passkey.id)
      );

    const walletClient = createWalletClient({
      chain: getChain(apiService.chainId),
      transport: custom(finalProvider),
    });

    return new PrexSigner(
      finalProvider,
      walletClient,
      apiService.chainId,
      wallet.address,
      wallet.ownerIndex,
      wallet.passkeys.map((passkey) => passkey.id)
    );
  }

  async signReplaySafeHash(hash: Hex, from?: Address) {
    const messageHash = createReplaySafeHash(
      this.chainId,
      hash,
      from || this.myAddress
    );

    // TODO: use provider
    const signature = await this.provider.request({
      method: 'eth_sign',
      params: [from || this.myAddress, messageHash],
    });

    return signature as Hex;
  }

  async signTypedData(
    params: {
      types: any;
      domain: any;
      message: any;
      primaryType: string;
    },
    from?: Address
  ) {
    return await this.walletClient.signTypedData({
      account: from || this.myAddress,
      types: params.types,
      domain: params.domain,
      message: params.message,
      primaryType: params.primaryType,
    });
  }

  async signHash(hash: Hex, from?: Address) {
    const signature = await this.provider.request({
      method: 'eth_sign',
      params: [from || this.myAddress, hash],
    });

    return signature as Hex;
  }

  async signHashByPrivateKey(hash: Hex, privateKey: Hex) {
    const signature = await sign({
      privateKey,
      hash,
    });

    return encodeSignatureBytes(
      BigInt(this.index),
      serializeSignature(signature)
    );
  }

  async signWithSubKey({
    extraHash,
    subKey,
    from,
  }: {
    extraHash: Hex;
    subKey: {
      subKey: string;
      subSecret: string;
      isNew: boolean;
    };
    from?: Address;
  }): Promise<{
    signature: Hex;
    message: SubKeyMessage;
    keyType: 'main' | 'sub';
  }> {
    const now = Math.floor(Date.now() / 1000);

    const message = {
      issuedAt: now,
      expiredAt: now + 60 * 60 * 24,
      newPublicKey: subKey.subKey,
      extraHash: extraHash,
    };

    const hash = hashSubKeyMessage(message);

    if (subKey.isNew) {
      return {
        signature: await this.signReplaySafeHash(hash, from || this.myAddress),
        message,
        keyType: 'main',
      };
    }

    return {
      signature: await signMessage({
        message: {
          raw: hash,
        },
        privateKey: subKey.subSecret as Hex,
      }),
      message,
      keyType: 'sub',
    };
  }
}

function getChain(chainId: number) {
  switch (chainId) {
    case arbitrumSepolia.id:
      return arbitrumSepolia;
    case base.id:
      return base;
    case optimism.id:
      return optimism;
    default:
      return arbitrum;
  }
}

function hashSubKeyMessage(message: SubKeyMessage) {
  return keccak256(
    encodeAbiParameters(
      [
        {
          type: 'string',
          name: 'prefix',
        },
        {
          type: 'uint256',
          name: 'issued_at',
        },
        {
          type: 'uint256',
          name: 'expired_at',
        },
        {
          type: 'address',
          name: 'new_public_key',
        },
        {
          type: 'bytes32',
          name: 'extra_hash',
        },
      ],
      [
        'prex_subkey',
        BigInt(message.issuedAt),
        BigInt(message.expiredAt),
        getAddress(message.newPublicKey),
        message.extraHash,
      ]
    )
  );
}