import { Injectable, Logger } from '@nestjs/common';
import type { ZodSchema } from 'zod';
import { ModelRouter } from '../model-router';
import { PromptTemplateService } from '../prompts/prompt-template.service';
import { AiCostCalculatorService } from '../usage/ai-cost-calculator.service';
import { AiUsageLogService } from '../usage/ai-usage-log.service';
import { ContentSafetyService } from '../../content-safety/content-safety.service';
import { EventBusService } from '../../../common/event-bus/event-bus.service';
import { BaseDomainEvent } from '../../../common/events/base-domain.event';
import { PrismaService } from '../../../infrastructure/database/prisma.service';
import type { AiProvider } from '../providers/ai-provider.interface';
import type { GatewayRequest, GatewayResponse, ModelTier } from './ai-gateway.types';
import type { StreamChunk } from '../providers/ai-provider.interface';

class AIUsageRecorded extends BaseDomainEvent {
  eventType = 'ai.usage.recorded';
  constructor(public readonly usage: Record<string, any>) { super(); }
}
class ModelFallbackTriggered extends BaseDomainEvent {
  eventType = 'ai.fallback.triggered';
  constructor(public readonly payload: Record<string, any>) { super(); }
}

@Injectable()
export class AiGatewayService {
  private readonly logger = new Logger(AiGatewayService.name);
  private readonly DEFAULT_TIMEOUT_MS = 30_000;

  constructor(
    private readonly modelRouter: ModelRouter,
    private readonly promptTemplate: PromptTemplateService,
    private readonly costCalculator: AiCostCalculatorService,
    private readonly usageLog: AiUsageLogService,
    private readonly providers: Map<string, AiProvider>,
    private readonly prisma: PrismaService,
    private readonly contentSafety?: ContentSafetyService,
    private readonly eventBus?: EventBusService,
  ) {}

  async generate(request: GatewayRequest, timeoutMs = this.DEFAULT_TIMEOUT_MS): Promise<GatewayResponse> {
    const tierConfig = this.modelRouter.resolve(request.tier);

    const prompt = this.promptTemplate.get(request.promptKey, request.promptVersion);
    const messages = [
      { role: 'system' as const, content: this.buildSystemPrompt(prompt.systemPrompt, prompt.outputSchemaDesc) },
      ...request.messages,
    ];

    let lastError: Error | null = null;

    for (let attempt = 0; attempt <= tierConfig.maxRetries; attempt++) {
      const target = attempt === 0 ? tierConfig.preferred : tierConfig.fallback;
      const attemptProvider = this.resolveProviderForTarget(target.provider);
      const controller = new AbortController();
      const timeoutId = setTimeout(() => controller.abort(), timeoutMs);

      try {
        const output = await attemptProvider.generate({
          model: target.model,
          messages,
          temperature: 0.3,
          maxTokens: request.maxTokens ?? 4096,
          responseFormat: { type: 'json_object' },
          signal: controller.signal,
        });

        const safetyCheck = await this.contentSafety?.check(output.rawText, { contentType: 'ai_output' }).catch(() => ({ safe: true }));
        if (safetyCheck && !safetyCheck.safe) throw new Error('AI output blocked by content safety');

        const parsed = this.parseJson(output.rawText, request.outputSchema);
        const estimatedCost = this.costCalculator.calculate(
          target.provider,
          target.model,
          output.usage.inputTokens,
          output.usage.outputTokens,
        );

        this.usageLog.log({
          userId: request.userId,
          feature: request.feature,
          provider: target.provider,
          model: target.model,
          tier: request.tier,
          promptKey: request.promptKey,
          promptVersion: prompt.version,
          inputTokens: output.usage.inputTokens,
          outputTokens: output.usage.outputTokens,
          estimatedCost,
          latencyMs: output.latencyMs,
          success: true,
        }).catch(() => {});

        // Publish cost event for Quota/Cost module
        try { this.eventBus?.publish(new AIUsageRecorded({
          userId: request.userId,
          feature: request.feature,
          provider: target.provider,
          model: target.model,
          inputTokens: output.usage.inputTokens,
          outputTokens: output.usage.outputTokens,
          estimatedCost,
          timestamp: new Date().toISOString(),
        })); } catch {}

        clearTimeout(timeoutId);
        return {
          parsed,
          usage: {
            provider: target.provider,
            model: target.model,
            inputTokens: output.usage.inputTokens,
            outputTokens: output.usage.outputTokens,
            estimatedCost,
            latencyMs: output.latencyMs,
          },
        };
      } catch (error) {
        clearTimeout(timeoutId);
        lastError = error as Error;
        this.logger.warn(
          `AI attempt ${attempt + 1}/${tierConfig.maxRetries + 1} failed (${target.provider}/${target.model}): ${lastError.message}`,
        );

        // Record fallback when switching from preferred to fallback
        if (attempt === 0 && tierConfig.maxRetries > 0) {
          const fb = tierConfig.fallback;
          this.prisma.fallbackEvent.create({
            data: {
              tier: request.tier,
              taskType: request.feature,
              fromProvider: tierConfig.preferred.provider,
              fromModel: tierConfig.preferred.model,
              toProvider: fb.provider,
              toModel: fb.model,
              errorMessage: lastError.message?.slice(0, 500),
            },
          }).catch(() => {});

          try {
            this.eventBus?.publish(new ModelFallbackTriggered({
              tier: request.tier,
              fromProvider: tierConfig.preferred.provider,
              fromModel: tierConfig.preferred.model,
              toProvider: fb.provider,
              toModel: fb.model,
              errorMessage: lastError.message?.slice(0, 200),
            }));
          } catch {}
        }
      }
    }

    this.usageLog.log({
      userId: request.userId,
      feature: request.feature,
      provider: tierConfig.preferred.provider,
      model: tierConfig.preferred.model,
      tier: request.tier,
      promptKey: request.promptKey,
      promptVersion: prompt.version,
      inputTokens: 0,
      outputTokens: 0,
      estimatedCost: 0,
      latencyMs: 0,
      success: false,
      errorMessage: lastError?.message,
    }).catch(() => {});

    throw new Error(`All AI attempts failed: ${lastError?.message}`);
  }

  private resolveProvider(tier: ModelTier): AiProvider {
    return this.resolveProviderForTarget(this.modelRouter.resolve(tier).preferred.provider);
  }

  private resolveProviderForTarget(providerName: string): AiProvider {
    const provider = this.providers.get(providerName);
    if (!provider) {
      throw new Error(`AI provider not found: ${providerName}`);
    }
    return provider;
  }

  private parseJson(raw: string, schema?: ZodSchema): Record<string, any> {
    if (!schema) return {};

    // Layer 1: direct parse
    try {
      return schema.parse(JSON.parse(raw)) as Record<string, any>;
    } catch {
      // Layer 2: extract from markdown code fences
      const fenced = raw.match(/```(?:json)?\s*\n?([\s\S]*?)```/);
      if (fenced) {
        try {
          return schema.parse(JSON.parse(fenced[1])) as Record<string, any>;
        } catch {
          // fall through
        }
      }

      // Layer 3: extract first JSON object from text
      const objMatch = raw.match(/\{[\s\S]*\}/);
      if (objMatch) {
        return schema.parse(JSON.parse(objMatch[0])) as Record<string, any>;
      }

      throw new Error('No valid JSON found in AI response');
    }
  }

  async *generateStream(request: GatewayRequest, timeoutMs = this.DEFAULT_TIMEOUT_MS): AsyncGenerator<StreamChunk, void, undefined> {
    const tierConfig = this.modelRouter.resolve(request.tier);
    const prompt = this.promptTemplate.get(request.promptKey, request.promptVersion);
    const messages = [
      { role: 'system' as const, content: prompt.systemPrompt },
      ...request.messages,
    ];

    const target = tierConfig.preferred;
    const provider = this.resolveProviderForTarget(target.provider);

    const controller = new AbortController();
    const timeoutId = setTimeout(() => controller.abort(), timeoutMs);

    let fullContent = '';
    let inputTokens = 0;
    let outputTokens = 0;

    try {
      for await (const chunk of provider.generateStream!({
        model: target.model,
        messages,
        temperature: 0.3,
        maxTokens: request.maxTokens ?? 4096,
        signal: controller.signal,
      })) {
        if (chunk.type === 'error') {
          yield { type: 'error', content: chunk.content };
          return;
        }
        if (chunk.type === 'done' && chunk.usage) {
          inputTokens = chunk.usage.inputTokens;
          outputTokens = chunk.usage.outputTokens;
        }
        if (chunk.type === 'content') {
          fullContent += (chunk.content || '');
        }
        yield chunk;
      }

      // Record usage after stream completes
      if (inputTokens > 0) {
        const estimatedCost = this.costCalculator.calculate(target.provider, target.model, inputTokens, outputTokens);
        this.usageLog.log({
          userId: request.userId, feature: request.feature,
          provider: target.provider, model: target.model, tier: request.tier,
          promptKey: request.promptKey, promptVersion: prompt.version,
          inputTokens, outputTokens, estimatedCost, latencyMs: 0, success: true,
        }).catch(() => {});
      }
    } catch (error: any) {
      yield { type: 'error', content: error.message };
    } finally {
      clearTimeout(timeoutId);
    }
  }

  private buildSystemPrompt(systemPrompt: string, schemaDesc: string): string {
    return `${systemPrompt}\n\n请严格按照以下 JSON Schema 输出，只输出 JSON，不要包含其他内容：\n${schemaDesc}`;
  }
}
