import { type IdlTypes, eventDiscriminator } from '@coral-xyz/anchor'
import {
  TOKEN_2022_PROGRAM_ID,
  TOKEN_PROGRAM_ID,
  getAssociatedTokenAddressSync,
} from '@solana/spl-token'
import {
  type AccountInfo,
  type AddressLookupTableAccount,
  type Connection,
  type Signer,
  type SimulateTransactionConfig,
  type Transaction,
  type TransactionInstruction,
  ComputeBudgetProgram,
  PublicKey,
  SendTransactionError,
  TransactionMessage,
  VersionedTransaction,
} from '@solana/web3.js'
import { dataLength, dataSlice, hexlify } from 'ethers'

import {
  CCIPSolanaComputeUnitsExceededError,
  CCIPTokenMintInvalidError,
  CCIPTokenMintNotFoundError,
  CCIPTransactionNotFinalizedError,
} from '../errors/index.ts'
import type { ChainLog, WithLogger } from '../types.ts'
import { bigIntReplacer, getDataBytes, isBase64, sleep } from '../utils.ts'
import type { IDL as BASE_TOKEN_POOL_IDL } from './idl/1.6.0/BASE_TOKEN_POOL.ts'
import type { UnsignedSolanaTx, Wallet } from './types.ts'
import type { RateLimiterState } from '../chain.ts'

/**
 * Result of resolving an Associated Token Account for a given mint and owner.
 */
export type ResolvedATA = {
  /** The derived ATA address */
  ata: PublicKey
  /** The token program that owns the mint (SPL Token or Token-2022) */
  tokenProgram: PublicKey
  /** The raw mint account info */
  mintInfo: AccountInfo<Buffer>
}

/**
 * Resolves the Associated Token Account (ATA) for a given mint and owner.
 * Automatically detects the correct token program (SPL Token vs Token-2022).
 *
 * @param connection - Solana connection instance
 * @param mint - Token mint address
 * @param owner - Owner's wallet address
 * @returns ResolvedATA with ata, tokenProgram, and mintInfo
 * @throws CCIPTokenMintNotFoundError If the mint account doesn't exist
 * @throws CCIPTokenMintInvalidError If the mint exists but isn't a valid SPL token
 *
 * @example
 * ```typescript
 * const { ata, tokenProgram } = await resolveATA(connection, mintPubkey, ownerPubkey)
 * ```
 */
export async function resolveATA(
  connection: Connection,
  mint: PublicKey,
  owner: PublicKey,
): Promise<ResolvedATA> {
  const mintInfo = await connection.getAccountInfo(mint)
  if (!mintInfo) {
    throw new CCIPTokenMintNotFoundError(mint.toBase58())
  }

  // Validate the mint is owned by a valid token program
  const isValidTokenProgram =
    mintInfo.owner.equals(TOKEN_PROGRAM_ID) || mintInfo.owner.equals(TOKEN_2022_PROGRAM_ID)

  if (!isValidTokenProgram) {
    throw new CCIPTokenMintInvalidError(mint.toBase58(), mintInfo.owner.toBase58(), [
      TOKEN_PROGRAM_ID.toBase58(),
      TOKEN_2022_PROGRAM_ID.toBase58(),
    ])
  }

  // Allow PDAs as owners (for program vaults, etc.)
  const ata = getAssociatedTokenAddressSync(mint, owner, true, mintInfo.owner)
  return {
    ata,
    tokenProgram: mintInfo.owner,
    mintInfo,
  }
}

/**
 * Generates a hex-encoded discriminator for a Solana event.
 * @param eventName - Name of the event.
 * @returns Hex-encoded discriminator string.
 */
export function hexDiscriminator(eventName: string): string {
  return hexlify(eventDiscriminator(eventName))
}

/**
 * Waits for a Solana transaction to reach finalized status.
 * @param connection - Solana connection instance.
 * @param signature - Transaction signature to wait for.
 * @param intervalMs - Polling interval in milliseconds.
 * @param maxAttempts - Maximum polling attempts before timeout.
 */
export async function waitForFinalization(
  connection: Connection,
  signature: string,
  intervalMs = 500,
  maxAttempts = 500,
): Promise<void> {
  for (let attempt = 0; attempt < maxAttempts; attempt++) {
    const status = await connection.getSignatureStatuses([signature])
    const info = status.value[0]

    if (info?.confirmationStatus === 'finalized') {
      return
    }
    await sleep(intervalMs)
  }

  throw new CCIPTransactionNotFinalizedError(signature)
}

/**
 * Converts a camelCase string to snake_case.
 * @param str - String to convert.
 * @returns snake_case formatted string.
 */
export function camelToSnakeCase(str: string): string {
  return str
    .replace(/([A-Z]+)([A-Z][a-z]|$)/g, (_, p1: string, p2: string) => {
      if (p2) {
        return `_${p1.slice(0, -1).toLowerCase()}_${p2.toLowerCase()}`
      }
      return `_${p1.toLowerCase()}`
    })
    .replace(/([a-z])([A-Z])/g, '$1_$2')
    .toLowerCase()
    .replace(/^_/, '')
}

type ParsedLog = Pick<ChainLog, 'topics' | 'index' | 'address' | 'data'> & {
  data: string
  level: number
}

/**
 * Utility function to parse Solana logs with proper address and topic extraction.
 *
 * Solana logs are structured as a stack-based execution trace:
 * - "Program <address> invoke [<depth>]" - Program call starts
 * - "Program log: <data>" - Program emitted a log message
 * - "Program data: <base64>" - Program emitted structured data (Anchor events)
 * - "Program <address> success/failed" - Program call ends
 *
 * This function:
 * 1. Tracks the program call stack to determine which program emitted each log
 * 2. Extracts the first 8 bytes from base64 "Program data:" logs as topics (event discriminants)
 * 3. Converts logs to EVM-compatible ChainLog format for CCIP compatibility
 * 4. Returns ALL logs from the transaction - filtering should be done by the caller
 *
 * @param logs - Array of logMessages from Solana transaction
 * @returns Array of parsed log objects from all programs in the transaction
 */
export function parseSolanaLogs(logs: readonly string[]): ParsedLog[] {
  const results: ReturnType<typeof parseSolanaLogs> = []
  const programStack: string[] = []

  for (const [i, log] of logs.entries()) {
    // Track program calls and returns to maintain the address stack
    const matchInvoke = log.match(/^Program (\w+) invoke\b/)
    const matchReturn = !matchInvoke && log.match(/^Program (\w+) (success|failed)\b/)
    const matchLog = !matchInvoke && !matchReturn && log.match(/^Program (log|data): /)
    if (matchInvoke) {
      programStack.push(matchInvoke[1]!)
    } else if (matchReturn) {
      // Pop from stack when program returns
      programStack.pop()
    } else if (matchLog) {
      // Extract the actual log data
      const logData = log.slice(matchLog[0].length)
      const currentProgram = programStack[programStack.length - 1]!
      let topics: string[] = []

      if (log.startsWith('Program data: ')) {
        try {
          // Try to decode base64 and extract first 8 bytes as topic/discriminant
          const buffer = getDataBytes(logData)
          if (dataLength(buffer) >= 8) {
            topics = [dataSlice(buffer, 0, 8)]
          }
        } catch {
          // If base64 decoding fails, leave topics empty
        }
      }
      // For regular log messages, use the current program on stack
      results.push({
        topics,
        index: i,
        address: currentProgram,
        data: logData,
        level: programStack.length,
      })
    }
  }

  return results
}

/**
 * Extracts error information from Solana transaction logs.
 * @param logs_ - Raw log strings or parsed log objects (may include `tx` field with error info).
 * @returns Parsed error info with program and error details.
 */
export function getErrorFromLogs(
  logs_:
    | readonly string[]
    | readonly Pick<ChainLog, 'address' | 'index' | 'data' | 'topics' | 'tx'>[]
    | null,
): { program: string; [k: string]: string } | undefined {
  if (!logs_?.length) return
  let logs
  if (logs_.every((l) => typeof l === 'string')) logs = parseSolanaLogs(logs_)
  else logs = logs_

  const lastLog = logs[logs.length - 1]!
  // collect all logs from the last program execution (the one which failed)
  const lastProgramLogs = logs
    .reduceRight(
      (acc, l) =>
        // if acc is empty (i.e. on last log), or it is emitted by the same program and not a Program data:
        !acc.length || (l.address === acc[0]!.address && !l.topics.length) ? [l, ...acc] : acc,
      [] as Pick<ChainLog, 'address' | 'index' | 'data'>[],
    )
    .filter(({ data }) => !isBase64(data))
    .map(({ data }) => data as string)
    .reduceRight((acc, l) => {
      l = l.replace(/ (with message|thrown in) /, ' $1: ')
      if (l.endsWith(':') && acc.length) l = `${l} ${acc.shift()!}` // cosmetic: join lines ending in ':' with next
      try {
        // convert number[]s (common in solana logs) into slightly more readable 0x-bytearrays
        l = l.replace(/\[(\d{1,3}, ){3,}\d+\]/g, (m) =>
          hexlify(
            new Uint8Array(
              m
                .substring(1, m.length - 1)
                .split(', ')
                .map((x) => +x),
            ),
          ),
        )
      } catch {
        // ignore
      }
      const L = l.replace(/\.$/, '').split(/(?<=\w)\. /g)
      acc.unshift(...L)
      return acc
    }, [] as string[])

  const res: { program: string; [k: string]: string } = {
    program: lastLog.address,
  }
  if (lastProgramLogs.every((l) => l.match(/\w: /))) {
    Object.assign(
      res,
      Object.fromEntries(
        lastProgramLogs.map((l) => [
          l.substring(0, l.indexOf(': ')),
          l.substring(l.indexOf(': ') + 2),
        ]),
      ),
    )
  } else {
    res['error'] = lastProgramLogs.join('\n')
  }
  if (!!logs[0] && 'tx' in logs[0] && !!logs[0].tx?.error)
    Object.assign(
      res,
      Object.fromEntries(
        Object.entries(logs[0].tx.error as Record<string, [number, string]>).map(
          ([k, [i, e]]) => [`${k}[${i}]`, e] as const,
        ),
      ),
    )
  return res
}

/**
 * Simulates a Solana transaction to estimate compute units.
 * @param params - Simulation parameters including connection and payer.
 * @returns Simulation result with estimated compute units.
 */
export async function simulateTransaction(
  { connection, logger = console }: { connection: Connection } & WithLogger,
  {
    payerKey,
    computeUnitsOverride,
    ...rest
  }: {
    payerKey: PublicKey
    computeUnitsOverride?: number
    addressLookupTableAccounts?: AddressLookupTableAccount[]
  } & ({ instructions: TransactionInstruction[] } | { tx: Transaction | VersionedTransaction }),
) {
  // Add max compute units for simulation
  const maxComputeUnits = 1_400_000
  const recentBlockhash = '11111111111111111111111111111112'
  const computeBudgetIx = ComputeBudgetProgram.setComputeUnitLimit({
    units: computeUnitsOverride || maxComputeUnits,
  })

  let tx: VersionedTransaction
  if (!('tx' in rest)) {
    // Create message with compute budget instruction
    const message = new TransactionMessage({
      payerKey,
      recentBlockhash,
      instructions: [computeBudgetIx, ...rest.instructions],
    })

    const messageV0 = message.compileToV0Message(rest.addressLookupTableAccounts)
    tx = new VersionedTransaction(messageV0)
  } else if (!('version' in rest.tx)) {
    // Create message with compute budget instruction
    const message = new TransactionMessage({
      payerKey,
      recentBlockhash,
      instructions: [computeBudgetIx, ...rest.tx.instructions],
    })

    const messageV0 = message.compileToV0Message(rest.addressLookupTableAccounts)
    tx = new VersionedTransaction(messageV0)
  } else {
    tx = rest.tx
  }

  const config: SimulateTransactionConfig = {
    commitment: 'confirmed',
    replaceRecentBlockhash: true,
    sigVerify: false,
  }

  const result = await connection.simulateTransaction(tx, config)

  logger.debug('Simulation results:', {
    logs: result.value.logs,
    unitsConsumed: result.value.unitsConsumed,
    returnData: result.value.returnData,
    err: result.value.err,
  })
  if (result.value.err) {
    // same error sendTransaction sends, to be catched up
    throw new SendTransactionError({
      action: 'simulate',
      signature: '',
      transactionMessage: JSON.stringify(result.value.err, bigIntReplacer),
      logs: result.value.logs!,
    })
  }

  return result.value
}

/**
 * Used as `provider` in anchor's `Program` constructor, to support `.view()` simulations
 * without * requiring a full AnchorProvider with wallet
 * @param ctx - Context object containing connection and logger
 * @param feePayer - Fee payer for the simulated transaction
 * @returns Value returned by the simulated method
 */
export function simulationProvider(
  ctx: { connection: Connection } & WithLogger,
  feePayer: PublicKey = new PublicKey('11111111111111111111111111111112'),
) {
  return {
    connection: ctx.connection,
    wallet: {
      publicKey: feePayer,
    },
    simulate: async (tx: Transaction | VersionedTransaction, _signers?: Signer[]) =>
      simulateTransaction(ctx, {
        payerKey: feePayer,
        tx,
      }),
  }
}

/**
 * Sign, simulate, send and confirm as many instructions as possible on each transaction
 * @param ctx - Context object containing connection and logger
 * @param wallet - Wallet to sign and pay for txs
 * @param unsignedTx - instructions to sign and send
 *   - instructions - Instructions to send; they may not fit all in a single transaction,
 *       in which case they will be split into multiple transactions
 *   - mainIndex - Index of the main instruction
 *   - lookupTables - lookupTables to be used for main instruction
 * @param computeUnits - max computeUnits limit to be used for main instruction
 * @returns - signature of successful transaction including main instruction
 * @throws {@link CCIPSolanaComputeUnitsExceededError} if simulation exceeds compute units limit
 */
export async function simulateAndSendTxs(
  ctx: { connection: Connection } & WithLogger,
  wallet: Wallet,
  { instructions, mainIndex, lookupTables }: Omit<UnsignedSolanaTx, 'family'>,
  computeUnits?: number,
): Promise<string> {
  const { connection } = ctx
  let mainHash: string
  for (
    let [start, end] = [0, instructions.length];
    start < instructions.length;
    [start, end] = [end, instructions.length]
  ) {
    let computeUnitLimit, lastErr, addressLookupTableAccounts, ixs, includesMain
    do {
      ixs = instructions.slice(start, end)
      includesMain = mainIndex != null && start <= mainIndex && mainIndex < end
      addressLookupTableAccounts = includesMain ? lookupTables : undefined

      try {
        const simulated =
          (
            await simulateTransaction(ctx, {
              payerKey: wallet.publicKey,
              instructions: ixs,
              addressLookupTableAccounts,
            })
          ).unitsConsumed || 0

        if (simulated <= 200000) {
          computeUnitLimit = undefined
        } else if (!includesMain || computeUnits == null || simulated <= computeUnits) {
          computeUnitLimit = Math.ceil(simulated * 1.1)
        } else {
          throw new CCIPSolanaComputeUnitsExceededError(simulated, computeUnits)
        }
        break
      } catch (err) {
        lastErr = err
        end-- // truncate until finding a slice which fits (both computeUnits and tx size limits)
      }
    } while (end > start)
    if (end <= start) throw lastErr

    const blockhash = await connection.getLatestBlockhash('confirmed')
    const txMsg = new TransactionMessage({
      payerKey: wallet.publicKey,
      recentBlockhash: blockhash.blockhash,
      instructions: [
        ...(computeUnitLimit
          ? [ComputeBudgetProgram.setComputeUnitLimit({ units: computeUnitLimit })]
          : []),
        ...ixs,
      ],
    })
    const messageV0 = txMsg.compileToV0Message(addressLookupTableAccounts)
    const tx = new VersionedTransaction(messageV0)

    const signed = await wallet.signTransaction(tx)
    const signature = await connection.sendTransaction(signed)
    await connection.confirmTransaction({ signature, ...blockhash }, 'confirmed')
    if (includesMain) mainHash = signature
  }
  return mainHash!
}

/**
 * Convert TokenPool's rate limit to RateLimiterState object.
 * @param input - On-chain rate limiter bucket from the TokenPool IDL.
 * @returns RateLimiterState with capacity, rate, and current tokens, or null if disabled.
 */
export function convertRateLimiter(
  input: IdlTypes<typeof BASE_TOKEN_POOL_IDL>['BaseChain']['inboundRateLimit'],
): RateLimiterState {
  if (!input.cfg.enabled) return null
  const tokens = BigInt(input.tokens.toString())
  const out: RateLimiterState = {
    capacity: BigInt(input.cfg.capacity.toString()),
    rate: BigInt(input.cfg.rate.toString()),
    get tokens() {
      const cur =
        tokens + this.rate * BigInt(Math.floor(Date.now() / 1000) - input.lastUpdated.toNumber())
      if (cur < this.capacity) return cur
      else return this.capacity
    },
  }
  return out
}