import { Pool } from "pg";
import {
  TrackedAddressRecord,
  UserMapping,
  WalletMappingInput,
  WebhookShardMetadata,
} from "../types";
import { IMappingStore, ITrackingStore } from "./StoreContracts";

interface WalletRow {
  user_id: number;
  chain_type: string;
  address: string;
  wallet_id: string | null;
  derivation_index: number | null;
  organization_id: string | null;
  branch_id: string | null;
  terminal_id: string | null;
  payment_id: string | null;
  supported_asset_id: string | null;
  network: string | null;
}

interface TrackedAddressRow {
  normalized_address: string;
  user_id: number;
  chain_type: string;
  address: string;
  wallet_id: string | null;
  shard_key: string;
  webhook_id: string;
  created_at: Date;
  updated_at: Date;
}

interface ShardRow {
  shard_key: string;
  webhook_id: string;
  chain_type: string;
  capacity: number;
  label: string | null;
  last_synced_at: Date | null;
  addresses: string[] | null;
  address_count: number;
}

export class PostgresMappingStore implements IMappingStore {
  constructor(private readonly pool: Pool) {}

  async getByUserId(userId: number): Promise<UserMapping | null> {
    const result = await this.pool.query<WalletRow>(
      `SELECT user_id, chain_type, address, wallet_id, derivation_index, organization_id, branch_id, terminal_id, payment_id, supported_asset_id, network
       FROM user_wallet_mappings
       WHERE user_id = $1
       ORDER BY chain_type ASC`,
      [userId]
    );

    if (result.rowCount === 0) {
      return null;
    }

    const first = result.rows[0];
    return {
      user_id: first.user_id,
      wallets: result.rows
        .filter((row) => row.chain_type && row.address)
        .map((row) => ({
          chain_type: row.chain_type,
          address: row.address,
          wallet_id: row.wallet_id,
          derivation_index: row.derivation_index,
          organization_id: row.organization_id,
          branch_id: row.branch_id,
          terminal_id: row.terminal_id,
          payment_id: row.payment_id,
          supported_asset_id: row.supported_asset_id,
          network: row.network,
        })),
    };
  }

  async upsertUserMapping(userId: number, wallets: WalletMappingInput[]): Promise<UserMapping> {
    const client = await this.pool.connect();
    try {
      await client.query("BEGIN");
      for (const wallet of wallets) {
        await client.query(
          `INSERT INTO user_wallet_mappings (
             user_id, chain_type, address, normalized_address, wallet_id, derivation_index,
             organization_id, branch_id, terminal_id, payment_id, supported_asset_id, network
           )
           VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12)
           ON CONFLICT (normalized_address)
           DO UPDATE SET address = EXCLUDED.address,
                         normalized_address = EXCLUDED.normalized_address,
                         user_id = EXCLUDED.user_id,
                         chain_type = EXCLUDED.chain_type,
                         wallet_id = EXCLUDED.wallet_id,
                         derivation_index = EXCLUDED.derivation_index,
                         organization_id = EXCLUDED.organization_id,
                         branch_id = EXCLUDED.branch_id,
                         terminal_id = EXCLUDED.terminal_id,
                         payment_id = EXCLUDED.payment_id,
                         supported_asset_id = EXCLUDED.supported_asset_id,
                         network = EXCLUDED.network,
                         updated_at = NOW()`,
          [
            userId,
            wallet.chain_type,
            wallet.address,
            this.normalizeAddress(wallet.address),
            wallet.wallet_id || null,
            wallet.derivation_index || null,
            wallet.organization_id || null,
            wallet.branch_id || null,
            wallet.terminal_id || null,
            wallet.payment_id || null,
            wallet.supported_asset_id || null,
            wallet.network || null,
          ]
        );
      }
      await client.query("COMMIT");
    } catch (error) {
      await client.query("ROLLBACK");
      throw error;
    } finally {
      client.release();
    }

    return (await this.getByUserId(userId)) as UserMapping;
  }

  async getByTrackedAddress(address: string): Promise<{ userId: number; wallet: WalletMappingInput } | null> {
    const result = await this.pool.query<WalletRow>(
      `SELECT user_id, chain_type, address, wallet_id, derivation_index, organization_id, branch_id, terminal_id, payment_id, supported_asset_id, network
       FROM user_wallet_mappings
       WHERE normalized_address = $1
       LIMIT 1`,
      [this.normalizeAddress(address)]
    );

    const row = result.rows[0];
    if (!row) {
      return null;
    }

    return {
      userId: row.user_id,
        wallet: {
          chain_type: row.chain_type,
          address: row.address,
          wallet_id: row.wallet_id,
          derivation_index: row.derivation_index,
          organization_id: row.organization_id,
          branch_id: row.branch_id,
          terminal_id: row.terminal_id,
          payment_id: row.payment_id,
          supported_asset_id: row.supported_asset_id,
          network: row.network,
        },
      };
  }

  async getWalletsByChainType(chainType: string): Promise<Array<{ user_id: number; wallet: WalletMappingInput }>> {
    const result = await this.pool.query<WalletRow>(
      `SELECT user_id, chain_type, address, wallet_id, derivation_index, organization_id, branch_id, terminal_id, payment_id, supported_asset_id, network
       FROM user_wallet_mappings
       WHERE LOWER(chain_type) = $1
       ORDER BY user_id ASC`,
      [String(chainType || "").trim().toLowerCase()]
    );

    return result.rows.map((row) => ({
      user_id: row.user_id,
      wallet: {
        chain_type: row.chain_type,
        address: row.address,
        wallet_id: row.wallet_id,
        derivation_index: row.derivation_index,
        organization_id: row.organization_id,
        branch_id: row.branch_id,
        terminal_id: row.terminal_id,
        payment_id: row.payment_id,
        supported_asset_id: row.supported_asset_id,
        network: row.network,
      },
    }));
  }

  private normalizeAddress(address: string): string {
    const trimmed = address.trim();
    return trimmed.startsWith("0x") || trimmed.startsWith("0X") ? trimmed.toLowerCase() : trimmed;
  }
}

export class PostgresTrackingStore implements ITrackingStore {
  constructor(private readonly pool: Pool) {}

  async getAllTrackedAddresses(): Promise<Array<{ chainType: string; address: string }>> {
    const result = await this.pool.query<{ chain_type: string; normalized_address: string }>(
      `SELECT uwm.chain_type, ta.normalized_address
       FROM tracked_addresses ta
       INNER JOIN user_wallet_mappings uwm ON uwm.normalized_address = ta.normalized_address
       ORDER BY uwm.chain_type ASC, ta.normalized_address ASC`
    );

    return result.rows.map((row) => ({
      chainType: row.chain_type,
      address: row.normalized_address,
    }));
  }

  async getAllShards(): Promise<WebhookShardMetadata[]> {
    const result = await this.pool.query<ShardRow>(
      `SELECT s.shard_key,
              s.webhook_id,
              s.chain_type,
              s.capacity,
              s.label,
              s.last_synced_at,
              COALESCE(array_agg(ta.normalized_address ORDER BY ta.normalized_address)
                FILTER (WHERE ta.normalized_address IS NOT NULL), '{}') AS addresses,
              COUNT(ta.normalized_address)::int AS address_count
       FROM tracking_shards s
       LEFT JOIN tracked_addresses ta ON ta.shard_key = s.shard_key
       GROUP BY s.shard_key, s.webhook_id, s.chain_type, s.capacity, s.label, s.last_synced_at
       ORDER BY s.shard_key ASC`
    );

    return result.rows.map((row) => this.rowToShard(row));
  }

  async getShardByWebhookId(webhookId: string): Promise<WebhookShardMetadata | null> {
    const result = await this.pool.query<ShardRow>(
      `SELECT s.shard_key,
              s.webhook_id,
              s.chain_type,
              s.capacity,
              s.label,
              s.last_synced_at,
              COALESCE(array_agg(ta.normalized_address ORDER BY ta.normalized_address)
                FILTER (WHERE ta.normalized_address IS NOT NULL), '{}') AS addresses,
              COUNT(ta.normalized_address)::int AS address_count
       FROM tracking_shards s
       LEFT JOIN tracked_addresses ta ON ta.shard_key = s.shard_key
       WHERE s.webhook_id = $1
       GROUP BY s.shard_key, s.webhook_id, s.chain_type, s.capacity, s.label, s.last_synced_at
       LIMIT 1`,
      [webhookId]
    );

    return result.rows[0] ? this.rowToShard(result.rows[0]) : null;
  }

  async getAddress(address: string): Promise<TrackedAddressRecord | null> {
    const result = await this.pool.query<TrackedAddressRow>(
      `SELECT ta.normalized_address,
              uwm.user_id,
              uwm.chain_type,
              uwm.address,
              uwm.wallet_id,
              ta.shard_key,
              ts.webhook_id,
              ta.created_at,
              ta.updated_at
       FROM tracked_addresses ta
       INNER JOIN user_wallet_mappings uwm ON uwm.normalized_address = ta.normalized_address
       INNER JOIN tracking_shards ts ON ts.shard_key = ta.shard_key
       WHERE ta.normalized_address = $1
       LIMIT 1`,
      [this.normalizeAddress(address)]
    );

    const row = result.rows[0];
    if (!row) {
      return null;
    }

    return this.rowToTrackedAddress(row);
  }

  async getShardsByChainType(chainType: string): Promise<WebhookShardMetadata[]> {
    const result = await this.pool.query<ShardRow>(
      `SELECT s.shard_key,
              s.webhook_id,
              s.chain_type,
              s.capacity,
              s.label,
              s.last_synced_at,
              COALESCE(array_agg(ta.normalized_address ORDER BY ta.normalized_address)
                FILTER (WHERE ta.normalized_address IS NOT NULL), '{}') AS addresses,
              COUNT(ta.normalized_address)::int AS address_count
       FROM tracking_shards s
       LEFT JOIN tracked_addresses ta ON ta.shard_key = s.shard_key
       WHERE s.chain_type = $1
       GROUP BY s.shard_key, s.webhook_id, s.chain_type, s.capacity, s.label, s.last_synced_at
       ORDER BY COUNT(ta.normalized_address)::int ASC, s.shard_key ASC`,
      [chainType]
    );

    return result.rows.map((row) => this.rowToShard(row));
  }

  async upsertShards(shards: WebhookShardMetadata[]): Promise<void> {
    const client = await this.pool.connect();
    try {
      await client.query("BEGIN");
      for (const shard of shards) {
        await client.query(
          `INSERT INTO tracking_shards (shard_key, webhook_id, chain_type, capacity, label, last_synced_at)
           VALUES ($1, $2, $3, $4, $5, $6)
           ON CONFLICT (shard_key)
           DO UPDATE SET webhook_id = EXCLUDED.webhook_id,
                         chain_type = EXCLUDED.chain_type,
                         capacity = EXCLUDED.capacity,
                         label = EXCLUDED.label,
                         last_synced_at = COALESCE(tracking_shards.last_synced_at, EXCLUDED.last_synced_at)`,
          [shard.shard_key, shard.webhook_id, shard.chain_type, shard.capacity, shard.label || null, shard.last_synced_at]
        );
      }
      await client.query("COMMIT");
    } catch (error) {
      await client.query("ROLLBACK");
      throw error;
    } finally {
      client.release();
    }
  }

  async registerAssignments(assignments: TrackedAddressRecord[]): Promise<void> {
    if (assignments.length === 0) {
      return;
    }

    const client = await this.pool.connect();
    try {
      await client.query("BEGIN");
      for (const assignment of assignments) {
        await client.query(
          `INSERT INTO tracked_addresses
            (normalized_address, shard_key, created_at, updated_at)
           VALUES ($1, $2, $3, $4)
           ON CONFLICT (normalized_address)
           DO UPDATE SET shard_key = EXCLUDED.shard_key,
                         updated_at = EXCLUDED.updated_at`,
          [
            assignment.normalized_address,
            assignment.shard_key,
            assignment.created_at,
            assignment.updated_at,
          ]
        );
      }
      await client.query("COMMIT");
    } catch (error) {
      await client.query("ROLLBACK");
      throw error;
    } finally {
      client.release();
    }
  }

  async unregisterAddresses(addresses: string[]): Promise<void> {
    if (addresses.length === 0) {
      return;
    }

    const normalizedAddresses = Array.from(new Set(
      addresses
        .map((address) => this.normalizeAddress(address))
        .filter((address) => address.length > 0)
    ));

    if (normalizedAddresses.length === 0) {
      return;
    }

    await this.pool.query(
      `DELETE FROM tracked_addresses
       WHERE normalized_address = ANY($1::text[])`,
      [normalizedAddresses]
    );
  }

  private rowToTrackedAddress(row: TrackedAddressRow): TrackedAddressRecord {
    return {
      address: row.address,
      normalized_address: row.normalized_address,
      user_id: row.user_id,
      chain_type: row.chain_type,
      wallet_id: row.wallet_id,
      webhook_id: row.webhook_id,
      shard_key: row.shard_key,
      created_at: row.created_at.toISOString(),
      updated_at: row.updated_at.toISOString(),
    };
  }

  private rowToShard(row: ShardRow): WebhookShardMetadata {
    return {
      shard_key: row.shard_key,
      webhook_id: row.webhook_id,
      chain_type: row.chain_type,
      capacity: row.capacity,
      address_count: row.address_count,
      addresses: row.addresses || [],
      label: row.label,
      last_synced_at: row.last_synced_at ? row.last_synced_at.toISOString() : null,
    };
  }

  private normalizeAddress(address: string): string {
    const trimmed = address.trim();
    return trimmed.startsWith("0x") || trimmed.startsWith("0X") ? trimmed.toLowerCase() : trimmed;
  }
}

export async function ensurePostgresControlStore(pool: Pool): Promise<void> {
  await pool.query(`
    CREATE TABLE IF NOT EXISTS user_wallet_mappings (
      user_id BIGINT NOT NULL,
      chain_type TEXT NOT NULL,
      address TEXT NOT NULL,
      normalized_address TEXT NOT NULL UNIQUE,
      wallet_id TEXT NULL,
      derivation_index BIGINT NULL,
      organization_id TEXT NULL,
      branch_id TEXT NULL,
      terminal_id TEXT NULL,
      payment_id TEXT NULL,
      supported_asset_id TEXT NULL,
      network TEXT NULL,
      created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
      updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
      PRIMARY KEY (normalized_address)
    )
  `);

  await pool.query(`ALTER TABLE user_wallet_mappings ADD COLUMN IF NOT EXISTS derivation_index BIGINT NULL`);
  await pool.query(`ALTER TABLE user_wallet_mappings ADD COLUMN IF NOT EXISTS organization_id TEXT NULL`);
  await pool.query(`ALTER TABLE user_wallet_mappings ADD COLUMN IF NOT EXISTS branch_id TEXT NULL`);
  await pool.query(`ALTER TABLE user_wallet_mappings ADD COLUMN IF NOT EXISTS terminal_id TEXT NULL`);
  await pool.query(`ALTER TABLE user_wallet_mappings ADD COLUMN IF NOT EXISTS payment_id TEXT NULL`);
  await pool.query(`ALTER TABLE user_wallet_mappings ADD COLUMN IF NOT EXISTS supported_asset_id TEXT NULL`);
  await pool.query(`ALTER TABLE user_wallet_mappings ADD COLUMN IF NOT EXISTS network TEXT NULL`);

  // Older schema versions linked user_wallet_mappings.user_id to a legacy
  // user_mappings table. The application no longer uses that parent table,
  // so we proactively drop the stale foreign key during bootstrap.
  await pool.query(`
    ALTER TABLE user_wallet_mappings
    DROP CONSTRAINT IF EXISTS user_wallet_mappings_user_id_fkey
  `);

  await pool.query(`
    CREATE TABLE IF NOT EXISTS tracking_shards (
      shard_key TEXT PRIMARY KEY,
      webhook_id TEXT NOT NULL,
      chain_type TEXT NOT NULL,
      capacity INTEGER NOT NULL,
      label TEXT NULL,
      last_synced_at TIMESTAMPTZ NULL
    )
  `);

  await pool.query(`
    CREATE TABLE IF NOT EXISTS tracked_addresses (
      normalized_address TEXT PRIMARY KEY,
      shard_key TEXT NOT NULL REFERENCES tracking_shards(shard_key) ON DELETE CASCADE,
      created_at TIMESTAMPTZ NOT NULL,
      updated_at TIMESTAMPTZ NOT NULL
    )
  `);

  await pool.query(`
    DO $$
    BEGIN
      IF NOT EXISTS (
        SELECT 1
        FROM pg_constraint
        WHERE conname = 'tracked_addresses_normalized_address_fkey'
      ) THEN
        ALTER TABLE tracked_addresses
        ADD CONSTRAINT tracked_addresses_normalized_address_fkey
        FOREIGN KEY (normalized_address)
        REFERENCES user_wallet_mappings(normalized_address)
        ON DELETE CASCADE;
      END IF;
    END $$;
  `);

  await pool.query("CREATE INDEX IF NOT EXISTS idx_user_wallet_mappings_chain_type ON user_wallet_mappings (chain_type)");
  await pool.query("CREATE INDEX IF NOT EXISTS idx_tracked_addresses_shard_key ON tracked_addresses (shard_key)");
}
