// Modern Router implementation from Welshman repository
// Source: https://github.com/coracle-social/welshman/tree/master/packages/util/src/Router.ts

import {first, splitAt, sortBy, uniq, shuffle, pushToMapKey} from '@welshman/lib'
import type {TrustedEvent} from '@welshman/util'
import {getAddress, isReplaceable} from '@welshman/util'
import {isShareableRelayUrl} from '@welshman/util'
import {Address, COMMUNITY, ROOM} from '@welshman/util'

// Address utility functions that were removed in v0.4.2
const decodeAddress = (a: string, relays: string[] = []): Address => {
  return Address.from(a, relays);
}

const isCommunityAddress = (a: Address): boolean => {
  return a.kind === COMMUNITY; // 34550
}

const isGroupAddress = (a: Address): boolean => {
  return a.kind === ROOM; // 35834 (NIP-29 groups)
}

const isContextAddress = (a: Address): boolean => {
  return a.kind === COMMUNITY || a.kind === ROOM;
}

// Simplified Tag class for compatibility
class Tag {
  constructor(private tag: string[]) {}

  static from(tagArray: string[]): Tag {
    return new Tag(tagArray);
  }

  valueOf(): string[] {
    return this.tag;
  }

  value(): string {
    return this.tag[1] || '';
  }
}

// Simplified Tags class for compatibility
class Tags {
  constructor(private tags: string[][]) {}

  static fromEvent(event: TrustedEvent): Tags {
    return new Tags(event.tags || []);
  }

  static wrap(tags: string[][]): Tags {
    return new Tags(tags);
  }

  context(): Tags {
    // Return address tags that are communities or groups
    const contextTags = this.tags.filter(t => {
      if (t[0] === 'a' && t[1]) {
        try {
          const addr = decodeAddress(t[1]);
          return isContextAddress(addr);
        } catch {
          return false;
        }
      }
      return false;
    });
    return new Tags(contextTags);
  }

  communities(): Tags {
    const communityTags = this.tags.filter(t => {
      if (t[0] === 'a' && t[1]) {
        try {
          const addr = decodeAddress(t[1]);
          return isCommunityAddress(addr);
        } catch {
          return false;
        }
      }
      return false;
    });
    return new Tags(communityTags);
  }

  groups(): Tags {
    const groupTags = this.tags.filter(t => {
      if (t[0] === 'a' && t[1]) {
        try {
          const addr = decodeAddress(t[1]);
          return isGroupAddress(addr);
        } catch {
          return false;
        }
      }
      return false;
    });
    return new Tags(groupTags);
  }

  ancestors(): {
    mentions: { values: () => { valueOf: () => string[] }, relays: () => { valueOf: () => string[] } },
    replies: { values: () => { valueOf: () => string[] }, relays: () => { valueOf: () => string[] } },
    roots: { values: () => { valueOf: () => string[] }, relays: () => { valueOf: () => string[] } }
  } {
    const mentions: string[] = [];
    const replies: string[] = [];
    const roots: string[] = [];
    const mentionRelays: string[] = [];
    const replyRelays: string[] = [];
    const rootRelays: string[] = [];

    for (const tag of this.tags) {
      if (tag[0] === 'e' && tag[1]) {
        const relay = tag[2] || '';
        const marker = tag[3] || '';
        
        if (marker === 'mention') {
          mentions.push(tag[1]);
          if (relay) mentionRelays.push(relay);
        } else if (marker === 'reply' || (!marker && replies.length === 0)) {
          replies.push(tag[1]);
          if (relay) replyRelays.push(relay);
        } else if (marker === 'root') {
          roots.push(tag[1]);
          if (relay) rootRelays.push(relay);
        }
      }
    }

    return {
      mentions: {
        values: () => ({ valueOf: () => mentions }),
        relays: () => ({ valueOf: () => mentionRelays })
      },
      replies: {
        values: () => ({ valueOf: () => replies }),
        relays: () => ({ valueOf: () => replyRelays })
      },
      roots: {
        values: () => ({ valueOf: () => roots }),
        relays: () => ({ valueOf: () => rootRelays })
      }
    };
  }

  values(key?: string): { valueOf: () => string[] } {
    const values = this.tags
      .filter(t => !key || t[0] === key)
      .map(t => t[1])
      .filter(Boolean);
    return { valueOf: () => values };
  }

  whereKey(key: string): Tags {
    return new Tags(this.tags.filter(t => t[0] === key));
  }

  mapTo(fn: (tag: Tag) => unknown): { valueOf: () => unknown[] } {
    const mapped = this.tags.map(t => fn(new Tag(t)));
    return { valueOf: () => mapped };
  }

  exists(): boolean {
    return this.tags.length > 0;
  }
}

// Use TrustedEvent as the equivalent of Rumor from the original implementation
type Rumor = TrustedEvent

// Export RelayMode so it can be imported from this file
export enum RelayMode {
  Read = "read",
  Write = "write",
}

export type RouterOptions = {
  /**
   * Retrieves the user's public key.
   * @returns The user's public key as a string, or null if not available.
   */
  getUserPubkey: () => string | null

  /**
   * Retrieves the group relays for the specified address.
   * @param address - The address to retrieve group relays for.
   * @returns An array of group relay URLs as strings.
   */
  getGroupRelays: (address: string) => string[]

  /**
   * Retrieves the community relays for the specified address.
   * @param address - The address to retrieve community relays for.
   * @returns An array of community relay URLs as strings.
   */
  getCommunityRelays: (address: string) => string[]

  /**
   * Retrieves the relays for the specified public key and mode.
   * @param pubkey - The public key to retrieve relays for.
   * @param mode - The relay mode (optional).
   * @returns An array of relay URLs as strings.
   */
  getPubkeyRelays: (pubkey: string, mode?: RelayMode) => string[]

  /**
   * Retrieves the static relays. These are used as fallbacks.
   * @returns An array of relay URLs as strings.
   */
  getStaticRelays: () => string[]

  /**
   * Retrieves relays likely to return results for kind 0, 3, and 10002.
   * @returns An array of relay URLs as strings.
   */
  getIndexerRelays: () => string[]

  /**
   * Retrieves relays likely to support NIP-50 search.
   * @returns An array of relay URLs as strings.
   */
  getSearchRelays: () => string[]

  /**
   * Retrieves the quality of the specified relay.
   * @param url - The URL of the relay to retrieve quality for.
   * @returns The quality of the relay as a number between 0 and 1 inclusive.
   */
  getRelayQuality: (url: string) => number

  /**
   * Retrieves the redundancy setting, which is how many relays to use per selection value.
   * @returns The redundancy setting as a number.
   */
  getRedundancy: () => number

  /**
   * Retrieves the limit setting, which is the maximum number of relays that should be
   * returned from getUrls and getSelections.
   * @returns The limit setting as a number.
   */
  getLimit: () => number
}

export type ValuesByRelay = Map<string, string[]>

export type RelayValues = {
  relay: string
  values: string[]
}

export type ValueRelays = {
  value: string
  relays: string[]
}

export type FallbackPolicy = (count: number, limit: number) => number

export class Router {
  constructor(readonly options: RouterOptions) {}

  // Utilities derived from options

  getPubkeySelection = (pubkey: string, mode?: RelayMode) =>
    this.selection(pubkey, this.options.getPubkeyRelays(pubkey, mode))

  getPubkeySelections = (pubkeys: string[], mode?: RelayMode) =>
    pubkeys.map(pubkey => this.getPubkeySelection(pubkey, mode))

  getUserSelections = (mode?: RelayMode) => {
    const userPubkey = this.options.getUserPubkey()
    if (userPubkey) {
      return this.getPubkeySelections([userPubkey], mode)
    } else {
      // If no user is logged in, use static relays
      const staticRelays = this.options.getStaticRelays()
      return [this.selection('anonymous', staticRelays)]
    }
  }

  getContextSelections = (tags: Tags): ValueRelays[] => {
    const communitySelections = tags.communities().mapTo((t: Tag) => this.selection(t.value(), this.options.getCommunityRelays(t.value()))).valueOf() as ValueRelays[]
    const groupSelections = tags.groups().mapTo((t: Tag) => this.selection(t.value(), this.options.getGroupRelays(t.value()))).valueOf() as ValueRelays[]
    
    return [
      ...communitySelections,
      ...groupSelections,
    ]
  }

  // Utilities for creating ValueRelays

  selection = (value: string, relays: Iterable<string>) => ({value, relays: Array.from(relays)})

  selections = (values: string[], relays: string[]) =>
    values.map(value => this.selection(value, relays))

  forceValue = (value: string, selections: ValueRelays[]) =>
    selections.map(({relays}) => this.selection(value, relays))

  // Utilities for processing hints

  relaySelectionsFromMap = (valuesByRelay: ValuesByRelay) =>
    sortBy(
      ({values}) => -values.length,
      Array.from(valuesByRelay)
        .map(([relay, values]: [string, string[]]) => ({relay, values: uniq(values)}))
    )

  scoreRelaySelection = ({values, relay}: RelayValues) =>
    values.length * this.options.getRelayQuality(relay)

  sortRelaySelections = (relaySelections: RelayValues[]) => {
    const scores = new Map<string, number>()
    const getScore = (relayValues: RelayValues) => -(scores.get(relayValues.relay) || 0)

    for (const relayValues of relaySelections) {
      scores.set(relayValues.relay, this.scoreRelaySelection(relayValues))
    }

    return sortBy(getScore, relaySelections.filter(getScore))
  }

  // Utilities for creating scenarios

  scenario = (selections: ValueRelays[]) => new RouterScenario(this, selections)

  merge = (scenarios: RouterScenario[]) =>
    this.scenario(scenarios.flatMap((scenario: RouterScenario) => scenario.selections))

  product = (values: string[], relays: string[]) =>
    this.scenario(this.selections(values, relays))

  fromRelays = (relays: string[]) => this.scenario([this.selection("", relays)])

  // Routing scenarios

  User = () => this.scenario(this.getUserSelections())

  ReadRelays = () => this.scenario(this.getUserSelections(RelayMode.Read))

  WriteRelays = () => this.scenario(this.getUserSelections(RelayMode.Write))

  Messages = (pubkeys: string[]) =>
    this.scenario([
      ...this.getUserSelections(),
      ...this.getPubkeySelections(pubkeys),
    ])

  PublishMessage = (pubkey: string) =>
    this.scenario([
      ...this.getUserSelections(RelayMode.Write),
      this.getPubkeySelection(pubkey, RelayMode.Read),
    ]).policy(this.addMinimalFallbacks)

  Event = (event: Rumor) =>
    this.scenario(this.forceValue(event.id, [
      this.getPubkeySelection(event.pubkey, RelayMode.Write),
      ...this.getContextSelections(Tags.fromEvent(event).context()),
    ]))

  EventChildren = (event: Rumor) =>
    this.scenario(this.forceValue(event.id, [
      this.getPubkeySelection(event.pubkey, RelayMode.Read),
      ...this.getContextSelections(Tags.fromEvent(event).context()),
    ]))

  EventAncestors = (event: Rumor, type: "mentions" | "replies" | "roots") => {
    const tags = Tags.fromEvent(event)
    const ancestors = tags.ancestors()[type]
    const pubkeys = tags.whereKey("p").values().valueOf()
    const communities = tags.communities().values().valueOf()
    const groups = tags.groups().values().valueOf()
    const relays = uniq([
      ...this.options.getPubkeyRelays(event.pubkey, RelayMode.Read),
      ...pubkeys.flatMap((k: string) => this.options.getPubkeyRelays(k, RelayMode.Write)),
      ...communities.flatMap((a: string) => this.options.getCommunityRelays(a)),
      ...groups.flatMap((a: string) => this.options.getGroupRelays(a)),
      ...ancestors.relays().valueOf(),
    ])

    return this.product(ancestors.values().valueOf(), relays)
  }

  EventMentions = (event: Rumor) => this.EventAncestors(event, "mentions")

  EventParents = (event: Rumor) => this.EventAncestors(event, "replies")

  EventRoots = (event: Rumor) => this.EventAncestors(event, "roots")

  PublishEvent = (event: Rumor) => {
    const tags = Tags.fromEvent(event)
    const mentions = tags.values("p").valueOf()

    // If we're publishing to private groups, only publish to those groups' relays
    if (tags.groups().exists()) {
      return this
        .scenario(this.getContextSelections(tags.groups()))
        .policy(this.addNoFallbacks)
    }

    return this.scenario(this.forceValue(event.id, [
      this.getPubkeySelection(event.pubkey, RelayMode.Write),
      ...this.getContextSelections(tags.context()),
      ...this.getPubkeySelections(mentions, RelayMode.Read),
    ]))
  }

  FromPubkeys = (pubkeys: string[]) =>
    this.scenario(this.getPubkeySelections(pubkeys, RelayMode.Write))

  ForPubkeys = (pubkeys: string[]) =>
    this.scenario(this.getPubkeySelections(pubkeys, RelayMode.Read))

  WithinGroup = (address: string, _relays?: string) =>
    this
      .scenario(this.getContextSelections(Tags.wrap([["a", address]])))
      .policy(this.addNoFallbacks)

  WithinCommunity = (address: string) =>
    this.scenario(this.getContextSelections(Tags.wrap([["a", address]])))

  WithinContext = (address: string) => {
    if (isGroupAddress(decodeAddress(address))) {
      return this.WithinGroup(address)
    }

    if (isCommunityAddress(decodeAddress(address))) {
      return this.WithinCommunity(address)
    }

    throw new Error(`Unknown context ${address}`)
  }

  WithinMultipleContexts = (addresses: string[]) =>
    this.merge(addresses.map(this.WithinContext))

  Search = (term: string, relays: string[] = []) =>
    this.product([term], uniq(this.options.getSearchRelays().concat(relays)))

  Indexers = (relays: string[] = []) =>
    this.fromRelays(uniq(this.options.getIndexerRelays().concat(relays)))

  // Fallback policies

  addNoFallbacks = (count: number, _redundancy: number) => count

  addMinimalFallbacks = (count: number, _redundancy: number) => Math.max(count, 1)

  addMaximalFallbacks = (count: number, redundancy: number) => redundancy - count

  // Higher level utils that use hints

  tagPubkey = (pubkey: string) =>
    Tag.from(["p", pubkey, this.FromPubkeys([pubkey]).getUrl()])

  tagEventId = (event: Rumor, mark = "") =>
    Tag.from(["e", event.id, this.Event(event).getUrl(), mark, event.pubkey])

  tagEventAddress = (event: Rumor, mark = "") =>
    Tag.from(["a", getAddress(event), this.Event(event).getUrl(), mark, event.pubkey])

  tagEvent = (event: Rumor, mark = "") => {
    const tags = [this.tagEventId(event, mark)]

    if (isReplaceable(event)) {
      tags.push(this.tagEventAddress(event, mark))
    }

    // Convert Tag objects to raw tag arrays
    const rawTags = tags.map(tag => tag.valueOf())
    return new Tags(rawTags)
  }

  address = (event: Rumor) =>
    Address.fromEvent(event, this.Event(event).redundancy(3).getUrls())
}

// Router Scenario

export type RouterScenarioOptions = {
  redundancy?: number
  policy?: FallbackPolicy
  limit?: number
}

export class RouterScenario {
  constructor(readonly router: Router, readonly selections: ValueRelays[], readonly options: RouterScenarioOptions = {}) {}

  clone = (options: RouterScenarioOptions) =>
    new RouterScenario(this.router, this.selections, {...this.options, ...options})

  select = (f: (selection: string) => boolean) =>
    new RouterScenario(this.router, this.selections.filter(({value}) => f(value)), this.options)

  redundancy = (redundancy: number) => this.clone({redundancy})

  policy = (policy: FallbackPolicy) => this.clone({policy})

  limit = (limit: number) => this.clone({limit})

  getRedundancy = () => this.options.redundancy || this.router.options.getRedundancy()

  getPolicy = () => this.options.policy || this.router.addMaximalFallbacks

  getLimit = () => this.options.limit || this.router.options.getLimit()

  getSelections = () => {
    const allValues = new Set()
    const valuesByRelay: ValuesByRelay = new Map()
    for (const {value, relays} of this.selections) {
      allValues.add(value)

      for (const relay of relays) {
        if (isShareableRelayUrl(relay)) {
          pushToMapKey(valuesByRelay, relay, value)
        }
      }
    }

    // Adjust redundancy by limit, since if we're looking for very specific values odds
    // are we're less tolerant of failure. Add more redundancy to fill our relay limit.
    const limit = this.getLimit()
    const redundancy = this.getRedundancy()
    const adjustedRedundancy = Math.max(redundancy, redundancy * (limit / (allValues.size * redundancy)))

    const seen = new Map<string, number>()
    const result: ValuesByRelay = new Map()
    const relaySelections = this.router.relaySelectionsFromMap(valuesByRelay)
    for (const {relay} of this.router.sortRelaySelections(relaySelections)) {
      const values = new Set<string>()
      for (const value of valuesByRelay.get(relay) || []) {
        const timesSeen = seen.get(value) || 0

        if (timesSeen < adjustedRedundancy) {
          seen.set(value, timesSeen + 1)
          values.add(value)
        }
      }

      if (values.size > 0) {
        result.set(relay, Array.from(values))
      }
    }

    const fallbacks = shuffle(this.router.options.getStaticRelays())
    const fallbackPolicy = this.getPolicy()
    for (const {value} of this.selections) {
      const timesSeen = seen.get(value) || 0
      const fallbacksNeeded = fallbackPolicy(timesSeen, adjustedRedundancy)

      if (fallbacksNeeded > 0) {
        for (const relay of fallbacks.slice(0, fallbacksNeeded)) {
          pushToMapKey(result, relay, value)
        }
      }
    }

    const [keep, discard] = splitAt(limit, this.router.relaySelectionsFromMap(result))

    for (const target of keep.slice(0, redundancy)) {
      target.values = uniq(discard.concat(target).flatMap((selection: RelayValues) => selection.values))
    }

    return keep
  }

  getUrls = () => this.getSelections().map((selection: RelayValues) => selection.relay)

  getUrl = () => first(this.getUrls())
}