import type { LlmProvider, Tier } from './types.js';

/**
 * Registry for LLM providers. Supports tier-based routing (fast/heavy)
 * with cross-tier fallback, and legacy single-provider mode.
 */
export class ProviderRegistry {
  private providers = new Map<string, LlmProvider>();
  private activeProvider: string | null = null;
  private tierProviders = new Map<Tier, string[]>();
  /** Maps a centralized Llm name → local provider name that can substitute when mcpd is unreachable. */
  private failoverMap = new Map<string, string>();

  register(provider: LlmProvider): void {
    this.providers.set(provider.name, provider);
    if (this.activeProvider === null) {
      this.activeProvider = provider.name;
    }
  }

  unregister(name: string): void {
    this.providers.delete(name);
    if (this.activeProvider === name) {
      const first = this.providers.keys().next();
      this.activeProvider = first.done ? null : first.value;
    }
    // Remove from tier assignments
    for (const [tier, names] of this.tierProviders) {
      const filtered = names.filter((n) => n !== name);
      if (filtered.length === 0) {
        this.tierProviders.delete(tier);
      } else {
        this.tierProviders.set(tier, filtered);
      }
    }
    // Remove from failover map (any entry whose local-provider value points at this name)
    for (const [centralName, localName] of this.failoverMap) {
      if (localName === name) this.failoverMap.delete(centralName);
    }
  }

  /** Mark `localProviderName` as the failover for the centralized Llm named `centralLlmName`. */
  registerFailover(centralLlmName: string, localProviderName: string): void {
    if (!this.providers.has(localProviderName)) {
      throw new Error(`Provider '${localProviderName}' is not registered`);
    }
    this.failoverMap.set(centralLlmName, localProviderName);
  }

  /** Look up the local provider that can substitute for a centralized Llm, if any. */
  getFailoverFor(centralLlmName: string): LlmProvider | null {
    const localName = this.failoverMap.get(centralLlmName);
    if (localName === undefined) return null;
    return this.providers.get(localName) ?? null;
  }

  /** Names of central Llms that have a local failover registered. Used in status output. */
  listFailovers(): Array<{ centralLlmName: string; localProviderName: string }> {
    return [...this.failoverMap.entries()].map(([centralLlmName, localProviderName]) => ({ centralLlmName, localProviderName }));
  }

  setActive(name: string): void {
    if (!this.providers.has(name)) {
      throw new Error(`Provider '${name}' is not registered`);
    }
    this.activeProvider = name;
  }

  getActive(): LlmProvider | null {
    if (this.activeProvider === null) return null;
    return this.providers.get(this.activeProvider) ?? null;
  }

  /** Assign a provider to a tier. Call order = priority within the tier. */
  assignTier(providerName: string, tier: Tier): void {
    if (!this.providers.has(providerName)) {
      throw new Error(`Provider '${providerName}' is not registered`);
    }
    const existing = this.tierProviders.get(tier) ?? [];
    if (!existing.includes(providerName)) {
      this.tierProviders.set(tier, [...existing, providerName]);
    }
  }

  /**
   * Get provider for a specific tier with fallback.
   * Resolution: requested tier → other tier → getActive() (legacy).
   */
  getProvider(tier: Tier): LlmProvider | null {
    const primary = this.firstInTier(tier);
    if (primary) return primary;

    const otherTier: Tier = tier === 'fast' ? 'heavy' : 'fast';
    const fallback = this.firstInTier(otherTier);
    if (fallback) return fallback;

    return this.getActive();
  }

  /** Get provider names assigned to a tier. */
  getTierProviders(tier: Tier): string[] {
    return this.tierProviders.get(tier) ?? [];
  }

  /** Whether any tier assignments exist (vs legacy single-provider mode). */
  hasTierConfig(): boolean {
    return this.tierProviders.size > 0;
  }

  get(name: string): LlmProvider | undefined {
    return this.providers.get(name);
  }

  list(): string[] {
    return [...this.providers.keys()];
  }

  getActiveName(): string | null {
    return this.activeProvider;
  }

  /** Provider info for status display. */
  listProviders(): Array<{ name: string; tiers: Tier[] }> {
    return this.list().map((name) => {
      const tiers: Tier[] = [];
      for (const [tier, names] of this.tierProviders) {
        if (names.includes(name)) tiers.push(tier);
      }
      return { name, tiers };
    });
  }

  /** Eagerly start providers that manage subprocesses (e.g., vLLM). */
  warmupAll(): void {
    for (const provider of this.providers.values()) {
      provider.warmup?.();
    }
  }

  /** Dispose all registered providers that have a dispose method. */
  disposeAll(): void {
    for (const provider of this.providers.values()) {
      provider.dispose?.();
    }
  }

  private firstInTier(tier: Tier): LlmProvider | null {
    const names = this.tierProviders.get(tier);
    if (!names) return null;
    for (const name of names) {
      const provider = this.providers.get(name);
      if (provider) return provider;
    }
    return null;
  }
}
