import { Injectable, Logger, type OnModuleInit } from '@nestjs/common';
import { ModelRegistry, AuthStorage } from '@mariozechner/pi-coding-agent';
import type { Model, Api } from '@mariozechner/pi-ai';
import type { ModelInfo, ProviderInfo, CustomProviderConfig } from '@mosaic/types';

@Injectable()
export class ProviderService implements OnModuleInit {
  private readonly logger = new Logger(ProviderService.name);
  private registry!: ModelRegistry;

  async onModuleInit(): Promise<void> {
    const authStorage = AuthStorage.create();
    this.registry = new ModelRegistry(authStorage);

    this.registerOllamaProvider();
    this.registerCustomProviders();

    const available = this.registry.getAvailable();
    this.logger.log(`Providers initialized: ${available.length} models available`);
  }

  getRegistry(): ModelRegistry {
    return this.registry;
  }

  findModel(provider: string, modelId: string): Model<Api> | undefined {
    return this.registry.find(provider, modelId);
  }

  getDefaultModel(): Model<Api> | undefined {
    const available = this.registry.getAvailable();
    return available[0];
  }

  listProviders(): ProviderInfo[] {
    const allModels = this.registry.getAll();
    const availableModels = this.registry.getAvailable();
    const availableIds = new Set(availableModels.map((m) => `${m.provider}:${m.id}`));

    const providerMap = new Map<string, ProviderInfo>();

    for (const model of allModels) {
      let info = providerMap.get(model.provider);
      if (!info) {
        info = {
          id: model.provider,
          name: model.provider,
          available: false,
          models: [],
        };
        providerMap.set(model.provider, info);
      }

      const isAvailable = availableIds.has(`${model.provider}:${model.id}`);
      if (isAvailable) info.available = true;

      info.models.push(this.toModelInfo(model));
    }

    return Array.from(providerMap.values());
  }

  listAvailableModels(): ModelInfo[] {
    return this.registry.getAvailable().map((m) => this.toModelInfo(m));
  }

  registerCustomProvider(config: CustomProviderConfig): void {
    this.registry.registerProvider(config.id, {
      baseUrl: config.baseUrl,
      apiKey: config.apiKey,
      models: config.models.map((m) => ({
        id: m.id,
        name: m.name,
        reasoning: m.reasoning ?? false,
        input: ['text'] as ('text' | 'image')[],
        cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
        contextWindow: m.contextWindow ?? 4096,
        maxTokens: m.maxTokens ?? 4096,
      })),
    });

    this.logger.log(`Registered custom provider: ${config.id} (${config.models.length} models)`);
  }

  private registerOllamaProvider(): void {
    const ollamaUrl = process.env['OLLAMA_BASE_URL'] ?? process.env['OLLAMA_HOST'];
    if (!ollamaUrl) return;

    const modelsEnv = process.env['OLLAMA_MODELS'] ?? 'llama3.2,codellama,mistral';
    const modelIds = modelsEnv
      .split(',')
      .map((m) => m.trim())
      .filter(Boolean);

    this.registerCustomProvider({
      id: 'ollama',
      name: 'Ollama',
      baseUrl: `${ollamaUrl}/v1`,
      models: modelIds.map((id) => ({
        id,
        name: id,
        reasoning: false,
        contextWindow: 8192,
        maxTokens: 4096,
      })),
    });

    this.logger.log(
      `Ollama provider registered at ${ollamaUrl} with models: ${modelIds.join(', ')}`,
    );
  }

  private registerCustomProviders(): void {
    const customJson = process.env['MOSAIC_CUSTOM_PROVIDERS'];
    if (!customJson) return;

    try {
      const configs = JSON.parse(customJson) as CustomProviderConfig[];
      for (const config of configs) {
        this.registerCustomProvider(config);
      }
    } catch (err) {
      this.logger.error('Failed to parse MOSAIC_CUSTOM_PROVIDERS', String(err));
    }
  }

  private toModelInfo(model: Model<Api>): ModelInfo {
    return {
      id: model.id,
      provider: model.provider,
      name: model.name,
      reasoning: model.reasoning,
      contextWindow: model.contextWindow,
      maxTokens: model.maxTokens,
      inputTypes: model.input,
      cost: model.cost,
    };
  }
}