import * as anchor from "@coral-xyz/anchor";
import { Program } from "@coral-xyz/anchor";
import {
  ComputeBudgetProgram,
  Ed25519Program,
  PublicKey,
  SYSVAR_INSTRUCTIONS_PUBKEY,
  Transaction,
} from "@solana/web3.js";
import { createHash } from "crypto";
import { expect } from "chai";
import { Shine } from "../target/types/shine";

const MAGIC = Buffer.from("SHiNE", "utf8");
const FORMAT_MAJOR = 1;
const FORMAT_MINOR = 0;
const ZERO_HASH = Buffer.alloc(32, 0);
const LAST_BLOCK_STATE_PREFIX = Buffer.from("SHiNE_LAST_BLOCK", "utf8");
const BLOCK_TYPE_ROOT_KEY = 1;
const BLOCK_TYPE_DEVICE_KEY = 2;
const BLOCK_TYPE_BLOCKCHAIN_REGISTRY = 3;
const BLOCK_TYPE_SERVER_PROFILE = 30;
const BLOCK_TYPE_ACCESS_SERVERS = 40;
const BLOCK_TYPE_TRUSTED_STATE = 50;
const BLOCK_TYPE_SESSIONS = 55;
const BLOCK_VERSION_0 = 0;
const BLOCKCHAIN_TYPE_MAIN_USER = 1;
const SESSIONS_MODE_MIXED = 1;

const LIMIT_STEP = 10_000n;
const START_BONUS_LIMIT = 100_000n;
const USERS_ECONOMY_CONFIG_SEED = "shine_users_economy_config";
const SHINE_PAYMENTS_PROGRAM_ID = new PublicKey(
  "c4yTa4JT9EtQDCBX9LmWFK6T2gp4JGsuymFbom2EudW"
);
const SHINE_PAYMENTS_INFLOW_VAULT_SEED = "shine_payments_inflow_vault";

type SessionRecord = {
  sessionType: number;
  sessionVersion: number;
  sessionName: string;
  sessionPubKey: PublicKey;
};

type MutableFields = {
  deviceKey: PublicKey;
  blockchainPublicKey: PublicKey;
  blockchainName: string;
  usedBytes: bigint;
  lastBlockNumber: number;
  lastBlockHash: Buffer;
  lastBlockSignature: Buffer;
  arweaveTxId: string;
  isServer: boolean;
  addressFormatType: number;
  addressFormatVersion: number;
  serverAddress: string;
  syncServers: string[];
  accessServers: string[];
  sessionsMode: number;
  sessions: SessionRecord[];
  trustedCount: number;
};

type UnsignedRecord = {
  createdAtMs: bigint;
  updatedAtMs: bigint;
  recordNumber: number;
  prevRecordHash: Buffer;
  login: string;
  rootKey: PublicKey;
  deviceKey: PublicKey;
  blockchain: {
    blockchainType: number;
    blockchainName: string;
    blockchainPublicKey: PublicKey;
    paidLimitBytes: bigint;
    usedBytes: bigint;
    lastBlockNumber: number;
    lastBlockHash: Buffer;
    lastBlockSignature: Buffer;
    arweaveTxId: string;
  };
  isServer: boolean;
  addressFormatType: number;
  addressFormatVersion: number;
  serverAddress: string;
  syncServers: string[];
  accessServers: string[];
  sessionsMode: number;
  sessions: SessionRecord[];
  trustedCount: number;
};

function u32le(v: number): Buffer {
  const b = Buffer.alloc(4);
  b.writeUInt32LE(v, 0);
  return b;
}

function u64le(v: bigint): Buffer {
  const b = Buffer.alloc(8);
  b.writeBigUInt64LE(v, 0);
  return b;
}

function strBytes(value: string): Buffer {
  const bytes = Buffer.from(value, "utf8");
  return Buffer.concat([Buffer.from([bytes.length]), bytes]);
}

function serializeUnsignedRecord(r: UnsignedRecord): Buffer {
  const out: Buffer[] = [];
  out.push(MAGIC);
  out.push(Buffer.from([FORMAT_MAJOR]));
  out.push(Buffer.from([FORMAT_MINOR]));
  out.push(Buffer.alloc(2, 0));

  out.push(u64le(r.createdAtMs));
  out.push(u64le(r.updatedAtMs));
  out.push(u32le(r.recordNumber));
  out.push(r.prevRecordHash);
  out.push(strBytes(r.login));

  out.push(Buffer.from([r.isServer ? 7 : 6]));
  out.push(Buffer.from([BLOCK_TYPE_ROOT_KEY, BLOCK_VERSION_0]));
  out.push(r.rootKey.toBuffer());
  out.push(Buffer.from([BLOCK_TYPE_DEVICE_KEY, BLOCK_VERSION_0]));
  out.push(r.deviceKey.toBuffer());
  out.push(Buffer.from([BLOCK_TYPE_BLOCKCHAIN_REGISTRY, BLOCK_VERSION_0, 1]));
  out.push(Buffer.from([r.blockchain.blockchainType]));
  out.push(strBytes(r.blockchain.blockchainName));
  out.push(r.blockchain.blockchainPublicKey.toBuffer());
  out.push(u64le(r.blockchain.paidLimitBytes));
  out.push(u64le(r.blockchain.usedBytes));
  out.push(u32le(r.blockchain.lastBlockNumber));
  out.push(r.blockchain.lastBlockHash);
  out.push(r.blockchain.lastBlockSignature);
  if (r.blockchain.arweaveTxId.length === 0) {
    out.push(Buffer.from([0]));
  } else {
    out.push(Buffer.from([1]));
    out.push(strBytes(r.blockchain.arweaveTxId));
  }

  if (r.isServer) {
    out.push(Buffer.from([BLOCK_TYPE_SERVER_PROFILE, BLOCK_VERSION_0, 1]));
    out.push(Buffer.from([r.addressFormatType, r.addressFormatVersion]));
    out.push(strBytes(r.serverAddress));
    out.push(Buffer.from([r.syncServers.length]));
    for (const s of r.syncServers) out.push(strBytes(s));
  }

  out.push(Buffer.from([BLOCK_TYPE_ACCESS_SERVERS, BLOCK_VERSION_0, r.accessServers.length]));
  for (const s of r.accessServers) out.push(strBytes(s));

  out.push(Buffer.from([BLOCK_TYPE_SESSIONS, BLOCK_VERSION_0, r.sessionsMode, r.sessions.length]));
  for (const s of r.sessions) {
    out.push(Buffer.from([s.sessionType, s.sessionVersion]));
    out.push(strBytes(s.sessionName));
    out.push(s.sessionPubKey.toBuffer());
  }

  out.push(Buffer.from([BLOCK_TYPE_TRUSTED_STATE, BLOCK_VERSION_0, r.trustedCount]));

  const unsigned = Buffer.concat(out);
  const recordLen = unsigned.length + 64;
  unsigned.writeUInt16LE(recordLen, 7);
  return unsigned;
}

function sha256(buf: Buffer): Buffer {
  return createHash("sha256").update(buf).digest();
}

function serializeLastBlockState(r: UnsignedRecord): Buffer {
  return Buffer.concat([
    LAST_BLOCK_STATE_PREFIX,
    strBytes(r.login),
    strBytes(r.blockchain.blockchainName),
    u32le(r.blockchain.lastBlockNumber),
    r.blockchain.lastBlockHash,
    u64le(r.blockchain.usedBytes),
  ]);
}

function extractSigFromEdIx(ixData: Buffer): Buffer {
  const signatureOffset = ixData.readUInt16LE(2);
  return ixData.subarray(signatureOffset, signatureOffset + 64);
}

describe("shine_users e2e", () => {
  anchor.setProvider(anchor.AnchorProvider.env());
  const provider = anchor.getProvider() as anchor.AnchorProvider;
  const program = anchor.workspace.shine as Program<Shine>;

  it("registers user and updates balance/server data", async () => {
    const login = `u${Date.now().toString().slice(-10)}`;
    const [userPda] = PublicKey.findProgramAddressSync(
      [Buffer.from("login="), Buffer.from(login, "utf8")],
      program.programId
    );
    const [usersEconomyConfigPda] = PublicKey.findProgramAddressSync(
      [Buffer.from(USERS_ECONOMY_CONFIG_SEED, "utf8")],
      program.programId
    );
    const [inflowVaultPda] = PublicKey.findProgramAddressSync(
      [Buffer.from(SHINE_PAYMENTS_INFLOW_VAULT_SEED, "utf8")],
      SHINE_PAYMENTS_PROGRAM_ID
    );

    const economyAi = await provider.connection.getAccountInfo(usersEconomyConfigPda);
    if (!economyAi) {
      await program.methods
        .initUsersEconomyConfig()
        .accounts({
          signer: provider.wallet.publicKey,
          usersEconomyConfigPda,
        })
        .rpc();
    }

    const root = anchor.web3.Keypair.generate();
    const blockchain = anchor.web3.Keypair.generate();
    const deviceKey = anchor.web3.Keypair.generate().publicKey;
    const blockchainName = `${login}-001`;

    const createFields: MutableFields = {
      deviceKey,
      blockchainPublicKey: blockchain.publicKey,
      blockchainName,
      usedBytes: 0n,
      lastBlockNumber: 0,
      lastBlockHash: ZERO_HASH,
      lastBlockSignature: Buffer.alloc(64, 0),
      arweaveTxId: "",
      isServer: true,
      addressFormatType: 1,
      addressFormatVersion: 0,
      serverAddress: "https://srv-1.local",
      syncServers: ["sync_srv_1", "sync_srv_2"],
      accessServers: ["access_srv_1"],
      sessionsMode: SESSIONS_MODE_MIXED,
      sessions: [],
      trustedCount: 0,
    };

    const createdAtMs = BigInt(Date.now());
    const additionalLimitCreate = 20_000n;
    expect(additionalLimitCreate % LIMIT_STEP).eq(0n);

    const createRecord: UnsignedRecord = {
      createdAtMs,
      updatedAtMs: createdAtMs,
      recordNumber: 0,
      prevRecordHash: ZERO_HASH,
      login,
      rootKey: root.publicKey,
      deviceKey,
      blockchain: {
        blockchainType: BLOCKCHAIN_TYPE_MAIN_USER,
        blockchainName,
        blockchainPublicKey: blockchain.publicKey,
        paidLimitBytes: START_BONUS_LIMIT + additionalLimitCreate,
        usedBytes: createFields.usedBytes,
        lastBlockNumber: createFields.lastBlockNumber,
        lastBlockHash: createFields.lastBlockHash,
        lastBlockSignature: Buffer.alloc(64, 0),
        arweaveTxId: createFields.arweaveTxId,
      },
      isServer: createFields.isServer,
      addressFormatType: createFields.addressFormatType,
      addressFormatVersion: createFields.addressFormatVersion,
      serverAddress: createFields.serverAddress,
      syncServers: createFields.syncServers,
      accessServers: createFields.accessServers,
      sessionsMode: createFields.sessionsMode,
      sessions: createFields.sessions,
      trustedCount: createFields.trustedCount,
    };
    const createLastBlockHash = sha256(serializeLastBlockState(createRecord));
    const createLastBlockEdIx = Ed25519Program.createInstructionWithPrivateKey({
      privateKey: blockchain.secretKey,
      message: createLastBlockHash,
    });
    createRecord.blockchain.lastBlockSignature = extractSigFromEdIx(Buffer.from(createLastBlockEdIx.data));
    createFields.lastBlockSignature = createRecord.blockchain.lastBlockSignature;

    const createUnsigned = serializeUnsignedRecord(createRecord);
    const createHash = sha256(createUnsigned);
    const createEdIx = Ed25519Program.createInstructionWithPrivateKey({
      privateKey: root.secretKey,
      message: createHash,
    });
    const createSig = extractSigFromEdIx(Buffer.from(createEdIx.data));

    const createIx = await program.methods
      .createUserPda({
        login,
        rootKey: root.publicKey,
        createdAtMs: new anchor.BN(createdAtMs.toString()),
        additionalLimit: new anchor.BN(additionalLimitCreate.toString()),
        fields: {
          deviceKey: createFields.deviceKey,
          blockchainPublicKey: createFields.blockchainPublicKey,
          blockchainName: createFields.blockchainName,
          usedBytes: new anchor.BN(createFields.usedBytes.toString()),
          lastBlockNumber: createFields.lastBlockNumber,
          lastBlockHash: createFields.lastBlockHash,
          lastBlockSignature: createFields.lastBlockSignature,
          arweaveTxId: createFields.arweaveTxId,
          isServer: createFields.isServer,
          addressFormatType: createFields.addressFormatType,
          addressFormatVersion: createFields.addressFormatVersion,
          serverAddress: createFields.serverAddress,
          syncServers: createFields.syncServers,
          accessServers: createFields.accessServers,
          sessionsMode: createFields.sessionsMode,
          sessions: createFields.sessions.map((s) => ({
            sessionType: s.sessionType,
            sessionVersion: s.sessionVersion,
            sessionName: s.sessionName,
            sessionPubKey: s.sessionPubKey,
          })),
          trustedCount: createFields.trustedCount,
        },
        signature: createSig,
      })
      .accounts({
        signer: provider.wallet.publicKey,
        userPda,
        inflowVault: inflowVaultPda,
        instructions: SYSVAR_INSTRUCTIONS_PUBKEY,
        usersEconomyConfigPda,
      })
      .instruction();

    await provider.sendAndConfirm(new Transaction().add(createEdIx, createLastBlockEdIx, createIx), []);

    const createAcc = await provider.connection.getAccountInfo(userPda);
    expect(createAcc).not.eq(null);
    expect(createAcc!.owner.toBase58()).eq(program.programId.toBase58());

    const additionalLimitUpdate = 30_000n;
    expect(additionalLimitUpdate % LIMIT_STEP).eq(0n);
    const updatedDeviceKey = anchor.web3.Keypair.generate().publicKey;
    const updateFields: MutableFields = {
      deviceKey: updatedDeviceKey,
      blockchainPublicKey: blockchain.publicKey,
      blockchainName,
      usedBytes: 512n,
      lastBlockNumber: 1,
      lastBlockHash: sha256(Buffer.from("first-shine-block")),
      lastBlockSignature: Buffer.alloc(64, 0),
      arweaveTxId: "",
      isServer: true,
      addressFormatType: 1,
      addressFormatVersion: 0,
      serverAddress: "https://srv-2.local",
      syncServers: ["sync_srv_3"],
      accessServers: ["access_srv_2", "access_srv_3"],
      sessionsMode: SESSIONS_MODE_MIXED,
      sessions: [],
      trustedCount: 0,
    };

    const updateRecord: UnsignedRecord = {
      createdAtMs,
      updatedAtMs: createdAtMs + 1_000n,
      recordNumber: 1,
      prevRecordHash: sha256(createUnsigned),
      login,
      rootKey: root.publicKey,
      deviceKey: updatedDeviceKey,
      blockchain: {
        blockchainType: BLOCKCHAIN_TYPE_MAIN_USER,
        blockchainName,
        blockchainPublicKey: blockchain.publicKey,
        paidLimitBytes: START_BONUS_LIMIT + additionalLimitCreate + additionalLimitUpdate,
        usedBytes: updateFields.usedBytes,
        lastBlockNumber: updateFields.lastBlockNumber,
        lastBlockHash: updateFields.lastBlockHash,
        lastBlockSignature: Buffer.alloc(64, 0),
        arweaveTxId: updateFields.arweaveTxId,
      },
      isServer: updateFields.isServer,
      addressFormatType: updateFields.addressFormatType,
      addressFormatVersion: updateFields.addressFormatVersion,
      serverAddress: updateFields.serverAddress,
      syncServers: updateFields.syncServers,
      accessServers: updateFields.accessServers,
      sessionsMode: updateFields.sessionsMode,
      sessions: updateFields.sessions,
      trustedCount: updateFields.trustedCount,
    };
    const updateLastBlockHash = sha256(serializeLastBlockState(updateRecord));
    const updateLastBlockEdIx = Ed25519Program.createInstructionWithPrivateKey({
      privateKey: blockchain.secretKey,
      message: updateLastBlockHash,
    });
    updateRecord.blockchain.lastBlockSignature = extractSigFromEdIx(Buffer.from(updateLastBlockEdIx.data));
    updateFields.lastBlockSignature = updateRecord.blockchain.lastBlockSignature;

    const updateUnsigned = serializeUnsignedRecord(updateRecord);
    const updateHash = sha256(updateUnsigned);
    const updateEdIx = Ed25519Program.createInstructionWithPrivateKey({
      privateKey: root.secretKey,
      message: updateHash,
    });
    const updateSig = extractSigFromEdIx(Buffer.from(updateEdIx.data));

    const updateIx = await program.methods
      .updateUserPda({
        login,
        rootKey: root.publicKey,
        createdAtMs: new anchor.BN(createdAtMs.toString()),
        updatedAtMs: new anchor.BN((createdAtMs + 1_000n).toString()),
        version: 1,
        prevHash: sha256(createUnsigned),
        additionalLimit: new anchor.BN(additionalLimitUpdate.toString()),
        fields: {
          deviceKey: updateFields.deviceKey,
          blockchainPublicKey: updateFields.blockchainPublicKey,
          blockchainName: updateFields.blockchainName,
          usedBytes: new anchor.BN(updateFields.usedBytes.toString()),
          lastBlockNumber: updateFields.lastBlockNumber,
          lastBlockHash: updateFields.lastBlockHash,
          lastBlockSignature: updateFields.lastBlockSignature,
          arweaveTxId: updateFields.arweaveTxId,
          isServer: updateFields.isServer,
          addressFormatType: updateFields.addressFormatType,
          addressFormatVersion: updateFields.addressFormatVersion,
          serverAddress: updateFields.serverAddress,
          syncServers: updateFields.syncServers,
          accessServers: updateFields.accessServers,
          sessionsMode: updateFields.sessionsMode,
          sessions: updateFields.sessions.map((s) => ({
            sessionType: s.sessionType,
            sessionVersion: s.sessionVersion,
            sessionName: s.sessionName,
            sessionPubKey: s.sessionPubKey,
          })),
          trustedCount: updateFields.trustedCount,
        },
        signature: updateSig,
      })
      .accounts({
        signer: provider.wallet.publicKey,
        userPda,
        inflowVault: inflowVaultPda,
        instructions: SYSVAR_INSTRUCTIONS_PUBKEY,
        usersEconomyConfigPda,
      })
      .instruction();

    const updateComputeIx = ComputeBudgetProgram.setComputeUnitLimit({ units: 800_000 });
    const updateHeapIx = ComputeBudgetProgram.requestHeapFrame({ bytes: 262_144 });
    await provider.sendAndConfirm(
      new Transaction().add(updateComputeIx, updateHeapIx, updateEdIx, updateLastBlockEdIx, updateIx),
      []
    );

    const updatedAcc = await provider.connection.getAccountInfo(userPda);
    expect(updatedAcc).not.eq(null);
    const data = updatedAcc!.data;
    expect(data.subarray(0, 5).toString("utf8")).eq("SHiNE");
    expect(data[5]).eq(FORMAT_MAJOR);
    expect(data[6]).eq(FORMAT_MINOR);
  });
});
