import {
  createWalletClient,
  type Account,
  type Address,
  type Client,
  type ClientConfig,
  type ParseAccount,
  type Transport,
  type WalletActions,
  type WalletRpcSchema,
} from 'viem'
import { addPnsContracts } from '../contracts/addPnsContracts.js'
import type {
  ChainWithBaseContracts,
  ChainWithPns,
  CheckedChainWithPns,
} from '../contracts/consts.js'
import type { Assign, Prettify } from '../types.js'
import { pnsWalletActions, type PnsWalletActions } from './decorators/wallet.js'

export type PnsWalletClientConfig<
  TTransport extends Transport,
  TChain extends ChainWithBaseContracts,
  TAccountOrAddress extends Account | Address | undefined =
    | Account
    | Address
    | undefined,
> = Assign<
  Pick<
    ClientConfig<TTransport, TChain, TAccountOrAddress>,
    'account' | 'chain' | 'key' | 'name' | 'pollingInterval' | 'transport'
  >,
  {
    chain: TChain
  }
>

export type PnsWalletClient<
  TTransport extends Transport = Transport,
  TChain extends ChainWithPns = ChainWithPns,
  TAccount extends Account | undefined = Account | undefined,
> = Prettify<
  Client<
    TTransport,
    TChain,
    TAccount,
    WalletRpcSchema,
    WalletActions<TChain, TAccount> & PnsWalletActions<TChain, TAccount>
  >
>

/**
 * Creates an PNS Wallet Client with a given [Transport](https://viem.sh/docs/clients/intro.html) configured for a [Chain](https://viem.sh/docs/clients/chains.html).
 *
 * @param config - {@link PnsWalletClientConfig}
 * @returns An PNS Wallet Client. {@link PnsWalletClient}
 *
 * @example
 * import { custom } from 'viem'
 * import { mainnet } from 'viem/chains'
 * import { createPnsWalletClient } from '@pnsdomains/pnsjs'
 *
 * const client = createPnsWalletClient({
 *   chain: mainnet,
 *   transport: custom(window.ethereum),
 * })
 */
export const createPnsWalletClient = <
  TTransport extends Transport,
  TChain extends ChainWithBaseContracts,
  TAccountOrAddress extends Account | Address | undefined = undefined,
>({
  account,
  chain,
  key = 'pnsWallet',
  name = 'PNS Wallet Client',
  transport,
  pollingInterval,
}: PnsWalletClientConfig<
  TTransport,
  TChain,
  TAccountOrAddress
>): PnsWalletClient<
  TTransport,
  CheckedChainWithPns<TChain>,
  ParseAccount<TAccountOrAddress>
> => {
  return createWalletClient({
    account,
    chain: addPnsContracts(chain),
    key,
    name,
    pollingInterval,
    transport,
  }).extend(pnsWalletActions)
}