import type { ChatRequest, ChatResponse, ChatStreamEvent, ModelClient } from './types.js';
import { withRetry } from './retry.js';
import type { RetryConfig } from './retry.js';
import { logger } from '../logger.js';

export type ModelTier = 'fast' | 'default' | 'complex' | 'local';

export interface ModelRouterConfig {
  default: ModelClient;
  fast?: ModelClient;
  complex?: ModelClient;
  local?: ModelClient;
  fallbackChain: ModelClient[];
  tierFallbacks?: Map<ModelTier, ModelClient[]>;
  retryConfig?: RetryConfig;
  labels?: Partial<Record<ModelTier, string>>;
  onTierChange?: (tier: ModelTier) => void;
}

export class ModelRouter implements ModelClient {
  private clients: Map<ModelTier, ModelClient>;
  private labels: Map<ModelTier, string>;
  private defaultClient: ModelClient;
  private fallbackChain: ModelClient[];
  private tierFallbacks: Map<ModelTier, ModelClient[]>;
  private currentTier: ModelTier = 'default';
  private localProviderName?: string;
  private retryConfig?: RetryConfig;
  private tierChangeListeners: Array<(tier: ModelTier) => void> = [];
  private strictTiers: Set<ModelTier> = new Set();

  constructor(config: ModelRouterConfig) {
    this.clients = new Map();
    this.labels = new Map();
    this.defaultClient = config.default;
    this.fallbackChain = config.fallbackChain;
    this.tierFallbacks = config.tierFallbacks ?? new Map();
    this.retryConfig = config.retryConfig;

    if (config.onTierChange) {
      this.tierChangeListeners.push(config.onTierChange);
    }

    this.clients.set('default', config.default);
    if (config.fast) {this.clients.set('fast', config.fast);}
    if (config.complex) {this.clients.set('complex', config.complex);}
    if (config.local) {this.clients.set('local', config.local);}

    if (config.labels) {
      for (const tier of ['fast', 'default', 'complex', 'local'] as ModelTier[]) {
        if (config.labels[tier]) {
          this.labels.set(tier, config.labels[tier]);
        }
      }
    }
  }

  setTier(tier: ModelTier): boolean {
    if (this.clients.has(tier)) {
      this.currentTier = tier;
      for (const listener of this.tierChangeListeners) {
        listener(tier);
      }
      return true;
    }
    return false;
  }

  setOnTierChange(callback: (tier: ModelTier) => void): void {
    this.tierChangeListeners = [callback];
  }

  addOnTierChange(callback: (tier: ModelTier) => void): void {
    this.tierChangeListeners.push(callback);
  }

  getTier(): ModelTier {
    return this.currentTier;
  }

  getAvailableTiers(): ModelTier[] {
    return Array.from(this.clients.keys());
  }

  async chat(request: ChatRequest, tier?: ModelTier): Promise<ChatResponse> {
    const useTier = tier ?? this.currentTier;
    const primaryClient = this.clients.get(useTier) ?? this.defaultClient;
    const errors: Error[] = [];

    // Try primary client (with retry if configured)
    try {
      if (this.retryConfig) {
        return await withRetry(() => primaryClient.chat(request), this.retryConfig, 'primary model');
      }
      return await primaryClient.chat(request);
    } catch (error) {
      errors.push(error instanceof Error ? error : new Error(String(error)));
      logger.debug(`Primary model failed: ${errors[0].message}`);
    }

    if (this.strictTiers.has(useTier)) {
      throw errors[0];
    }

    // Try tier-specific fallbacks first
    const tierFallbackList = this.tierFallbacks.get(useTier) ?? [];
    for (let i = 0; i < tierFallbackList.length; i++) {
      try {
        const reason = `Primary model failed (${errors[0].message}), using tier fallback #${i + 1}`;
        logger.debug(reason);
        const response = await tierFallbackList[i].chat(request);
        return { ...response, fallback: true, fallbackReason: reason };
      } catch (error) {
        errors.push(error instanceof Error ? error : new Error(String(error)));
        logger.debug(`Tier fallback #${i + 1} failed: ${errors[errors.length - 1].message}`);
      }
    }

    // Then try global fallback chain
    for (let i = 0; i < this.fallbackChain.length; i++) {
      const fallbackClient = this.fallbackChain[i];
      try {
        const reason = `Primary model failed (${errors[0].message}), using global fallback #${i + 1}`;
        logger.debug(reason);
        const response = await fallbackClient.chat(request);
        return { ...response, fallback: true, fallbackReason: reason };
      } catch (error) {
        errors.push(error instanceof Error ? error : new Error(String(error)));
        logger.debug(`Global fallback #${i + 1} failed: ${errors[errors.length - 1].message}`);
      }
    }

    throw new Error(`All model providers failed: ${errors.map(e => e.message).join(', ')}`);
  }

  async *chatStream(request: ChatRequest, tier?: ModelTier): AsyncIterable<ChatStreamEvent> {
    const useTier = tier ?? this.currentTier;
    const primaryClient = this.clients.get(useTier) ?? this.defaultClient;
    let primaryError: string | undefined;

    if (primaryClient.chatStream) {
      let hasError = false;
      for await (const event of primaryClient.chatStream(request)) {
        if (event.type === 'error') {
          hasError = true;
          primaryError = event.error?.message ?? 'Unknown error';
          logger.debug(`Primary stream failed: ${primaryError}`);
          break;
        }
        yield event;
      }

      if (!hasError) {return;}
    } else {
      primaryError = 'Primary client does not support streaming';
    }

    if (this.strictTiers.has(useTier)) {
      yield { type: 'error', error: new Error(primaryError ?? 'Primary model failed') };
      return;
    }

    // Try tier-specific fallbacks first
    const tierFallbackList = this.tierFallbacks.get(useTier) ?? [];
    for (let i = 0; i < tierFallbackList.length; i++) {
      const fallbackClient = tierFallbackList[i];
      if (!fallbackClient.chatStream) {continue;}

      const reason = `Primary model failed (${primaryError}), using tier fallback #${i + 1}`;
      logger.debug(reason);
      yield { type: 'fallback_warning', fallbackReason: reason };

      let hasError = false;
      for await (const event of fallbackClient.chatStream(request)) {
        if (event.type === 'error') {
          hasError = true;
          logger.debug(`Tier fallback stream #${i + 1} failed: ${event.error?.message}`);
          break;
        }
        yield event;
      }

      if (!hasError) {return;}
    }

    // Then try global fallback chain
    for (let i = 0; i < this.fallbackChain.length; i++) {
      const fallbackClient = this.fallbackChain[i];
      if (!fallbackClient.chatStream) {continue;}

      const reason = `Primary model failed (${primaryError}), using global fallback #${i + 1}`;
      logger.debug(reason);
      yield { type: 'fallback_warning', fallbackReason: reason };

      let hasError = false;
      for await (const event of fallbackClient.chatStream(request)) {
        if (event.type === 'error') {
          hasError = true;
          logger.debug(`Global fallback stream #${i + 1} failed: ${event.error?.message}`);
          break;
        }
        yield event;
      }

      if (!hasError) {return;}
    }

    yield { type: 'error', error: new Error('All streaming providers failed') };
  }

  getClient(tier: ModelTier): ModelClient | undefined {
    return this.clients.get(tier);
  }

  setLocalClient(client: ModelClient, providerName: string): void {
    this.clients.set('local', client);
    this.localProviderName = providerName;
  }

  getLocalProviderName(): string | undefined {
    return this.localProviderName;
  }

  setClient(tier: ModelTier, client: ModelClient, label: string): void {
    this.clients.set(tier, client);
    this.labels.set(tier, label);
  }

  setTierStrict(tier: ModelTier, strict: boolean): void {
    if (strict) {
      this.strictTiers.add(tier);
      return;
    }
    this.strictTiers.delete(tier);
  }

  isTierStrict(tier: ModelTier): boolean {
    return this.strictTiers.has(tier);
  }

  getLabel(tier: ModelTier): string {
    return this.labels.get(tier) ?? 'unknown';
  }

  getAllLabels(): Record<string, string> {
    const result: Record<string, string> = {};
    for (const tier of this.labels.keys()) {
      result[tier] = this.labels.get(tier) ?? 'unknown';
    }
    return result;
  }
}
