/**
 * ASP Client
 *
 * High-level API for interacting with the ASP system.
 * Orchestrates between providers (off-chain storage) and
 * the blockchain (on-chain registry).
 */

import { generateMembershipProof } from './proof.js'
import type { IHashFunction } from './hash/interface.js'
import type {
  ASPClientConfig,
  IASPProvider,
  MembershipProof,
  RegisterASPOptions,
  PublishRootOptions,
} from './types.js'
import { ASP_REGISTRY_HUB_ABI } from '../contracts/abi/ASPRegistryHub.js'

/**
 * ASP Client instance
 */
export class ASPClient {
  private readonly provider: IASPProvider
  private readonly publicClient: any // viem PublicClient
  private readonly registryAddress: `0x${string}`
  private readonly hashFn: IHashFunction | undefined
  private aspId: bigint | undefined

  constructor(config: ASPClientConfig) {
    this.provider = config.provider
    this.publicClient = config.publicClient
    this.registryAddress = config.registryAddress
    this.aspId = config.aspId
    this.hashFn = config.hashFunction
  }

  /**
   * Register a new ASP on-chain
   */
  async register(options: RegisterASPOptions): Promise<bigint> {
    const walletClient = options.walletClient as any

    const hash = await walletClient.writeContract({
      address: this.registryAddress,
      abi: ASP_REGISTRY_HUB_ABI,
      functionName: 'registerASP',
      args: [options.name],
    })

    // Wait for receipt to get the ASP ID from events
    const receipt = await this.publicClient.waitForTransactionReceipt({ hash })

    // Extract ASP ID from ASPRegistered event
    for (const log of receipt.logs) {
      try {
        // ASPRegistered(uint256 indexed aspId, address indexed operator, string name)
        // The aspId is the first indexed topic
        if (log.topics[1]) {
          this.aspId = BigInt(log.topics[1])
          return this.aspId
        }
      } catch {
        continue
      }
    }

    throw new Error('Failed to extract ASP ID from transaction receipt')
  }

  /**
   * Get the current ASP ID
   */
  getASPId(): bigint | undefined {
    return this.aspId
  }

  /**
   * Set the ASP ID (for existing registrations)
   */
  setASPId(aspId: bigint): void {
    this.aspId = aspId
  }

  /**
   * Add a member identity to the tree
   */
  async addMember(identity: bigint): Promise<void> {
    await this.provider.addMember(identity)
  }

  /**
   * Remove a member identity from the tree
   */
  async removeMember(identity: bigint): Promise<void> {
    await this.provider.removeMember(identity)
  }

  /**
   * Get all members
   */
  async getMembers(): Promise<bigint[]> {
    return this.provider.getMembers()
  }

  /**
   * Check if an identity is a member
   */
  async hasMember(identity: bigint): Promise<boolean> {
    return this.provider.hasMember(identity)
  }

  /**
   * Get the current Merkle root from the provider
   */
  async getRoot(): Promise<bigint> {
    return this.provider.getRoot()
  }

  /**
   * Generate a membership proof for an identity
   */
  async generateProof(identity: bigint): Promise<MembershipProof> {
    const members = await this.provider.getMembers()
    return generateMembershipProof(identity, members, this.provider.treeDepth, this.hashFn)
  }

  /**
   * Publish the current Merkle root on-chain
   */
  async publishRoot(options: PublishRootOptions): Promise<`0x${string}`> {
    if (!this.aspId) {
      throw new Error('ASP not registered. Call register() first or set aspId.')
    }

    const root = await this.provider.getRoot()
    const walletClient = options.walletClient as any

    const hash = await walletClient.writeContract({
      address: this.registryAddress,
      abi: ASP_REGISTRY_HUB_ABI,
      functionName: 'updateRoot',
      args: [this.aspId, root],
    })

    return hash as `0x${string}`
  }

  /**
   * Publish a STARK-side Merkle root on-chain.
   *
   * The STARK side maintains a parallel tree (typically over the same
   * membership set) using a STARK-friendly hash (e.g., Poseidon31). The
   * root is supplied by the caller because the SNARK provider doesn't
   * compute it — the caller is expected to manage a `MerkleTree` with a
   * `PoseidonM31` hash and pass its root here.
   *
   * Operators MUST publish in lock-step with the SNARK root so both
   * proof systems stay aligned (an attacker who can update only one of
   * the two roots can keep the other stale and bypass compliance on
   * one path).
   */
  async publishStarkRoot(
    starkRoot: bigint,
    options: PublishRootOptions,
  ): Promise<`0x${string}`> {
    if (!this.aspId) {
      throw new Error('ASP not registered. Call register() first or set aspId.')
    }

    const walletClient = options.walletClient as any

    const hash = await walletClient.writeContract({
      address: this.registryAddress,
      abi: ASP_REGISTRY_HUB_ABI,
      functionName: 'updateStarkRoot',
      args: [this.aspId, starkRoot],
    })

    return hash as `0x${string}`
  }

  /**
   * Check if a root is valid on-chain (current or historical)
   */
  async isValidRoot(root: bigint): Promise<boolean> {
    if (!this.aspId) return false

    return this.publicClient.readContract({
      address: this.registryAddress,
      abi: ASP_REGISTRY_HUB_ABI,
      functionName: 'isValidASPRoot',
      args: [this.aspId, root],
    }) as Promise<boolean>
  }

  /**
   * Check if a STARK-side root is valid on-chain (current or historical).
   */
  async isValidStarkRoot(starkRoot: bigint): Promise<boolean> {
    if (!this.aspId) return false

    return this.publicClient.readContract({
      address: this.registryAddress,
      abi: ASP_REGISTRY_HUB_ABI,
      functionName: 'isValidSTARKRoot',
      args: [this.aspId, starkRoot],
    }) as Promise<boolean>
  }

  /**
   * Get the current on-chain root
   */
  async getCurrentOnChainRoot(): Promise<bigint> {
    if (!this.aspId) return 0n

    return this.publicClient.readContract({
      address: this.registryAddress,
      abi: ASP_REGISTRY_HUB_ABI,
      functionName: 'getCurrentRoot',
      args: [this.aspId],
    }) as Promise<bigint>
  }

  /**
   * Get the current on-chain STARK root.
   */
  async getCurrentOnChainStarkRoot(): Promise<bigint> {
    if (!this.aspId) return 0n

    return this.publicClient.readContract({
      address: this.registryAddress,
      abi: ASP_REGISTRY_HUB_ABI,
      functionName: 'getCurrentStarkRoot',
      args: [this.aspId],
    }) as Promise<bigint>
  }

  /**
   * Get the provider
   */
  getProvider(): IASPProvider {
    return this.provider
  }
}

/**
 * Create a new ASP client
 */
export function createASPClient(config: ASPClientConfig): ASPClient {
  return new ASPClient(config)
}