import { Injectable, Logger } from '@nestjs/common';
import { ConfigService } from '@nestjs/config';
import { generateObject, streamObject } from 'ai';
import type { Prisma } from '@prisma/client';
import type { AppConfig } from '../../config/configuration';
import { AiProviderService } from './core/ai-provider.service';
import { InvocationRecorderService } from './core/invocation-recorder.service';
import { PromptCacheService } from './core/prompt-cache.service';
import { SafetyGateRejectError, SafetyGateService } from './core/safety-gate.service';
import { computeInputHash } from './core/hash.util';
import type { AiCall, AiCallContext, AiCallResult } from './ai-call.interface';

/**
 * 流式事件 — orchestrator / controller 转换成 SSE 后吐给客户端
 */
export type StreamEvent<TOutput> =
  | { type: 'start'; invocationId: string; modelId: string; promptVersion: string }
  | { type: 'partial'; partial: Partial<TOutput> }
  | {
      type: 'done';
      output: TOutput;
      source: 'agent' | 'template_fallback';
      invocationId: string;
      promptTokens: number;
      completionTokens: number;
      costYuan: number;
      fallbackReason?: string;
    }
  | { type: 'error'; message: string; invocationId: string };

/**
 * AiCallRunner — harness 的执行器。
 *
 * 把 call 定义 + input + 上下文 跑通统一流程:
 *   1. 计算 inputHash
 *   2. 查缓存(可选 bustCache 跳过)
 *      命中 → 落 cached invocation → 返回 cached output
 *   3. 起 invocation(status=running)
 *   4. 调 generateObject(model + schema + system + prompt)
 *   5. safety gate 校验 → 失败走 fallback
 *   6. LLM 失败 → 走 fallback
 *   7. 落 invocation(succeeded / failed)+ 写缓存
 *   8. 返回 AiCallResult
 *
 * 任何 AiCall 都通过本 runner 跑,不允许业务代码自己 new generateObject。
 * 这是 audit / cost / eval / cache 一致性的唯一保证。
 */
@Injectable()
export class AiCallRunnerService {
  private readonly logger = new Logger(AiCallRunnerService.name);

  constructor(
    private readonly provider: AiProviderService,
    private readonly cache: PromptCacheService,
    private readonly recorder: InvocationRecorderService,
    private readonly safety: SafetyGateService,
    private readonly config: ConfigService<AppConfig, true>,
  ) {}

  async run<TInput, TOutput>(
    call: AiCall<TInput, TOutput>,
    input: TInput,
    ctx: AiCallContext,
  ): Promise<AiCallResult<TOutput>> {
    const inputHash = computeInputHash(call.callKey, call.promptVersion, input);
    const modelId = ctx.modelIdOverride ?? call.defaultModelId;
    const { model, provider } = this.provider.resolve(modelId);

    // ─── 1. 缓存查询 ───
    if (!ctx.bustCache) {
      const cached = this.cache.get<TOutput>(inputHash);
      if (cached !== undefined) {
        const invId = await this.recorder.start({
          ...this.baseStart(call, ctx, modelId, provider, inputHash, input, true),
        });
        await this.recorder.end(invId, { status: 'cached', output: cached as Prisma.InputJsonValue });
        this.logger.debug(`cache hit call=${call.callKey} hash=${inputHash.slice(0, 8)}`);
        return {
          output: cached,
          invocationId: invId,
          cacheHit: true,
          source: 'agent',
          promptTokens: 0,
          completionTokens: 0,
          costYuan: 0,
        };
      }
    }

    // ─── 2. 起 invocation ───
    const { system, prompt } = call.buildPrompt(input);
    const invocationId = await this.recorder.start({
      ...this.baseStart(call, ctx, modelId, provider, inputHash, input, false),
      promptTemplate: prompt.length > 8000 ? prompt.slice(0, 8000) + '…[truncated]' : prompt,
    });

    // ─── 3. 调 LLM ───
    try {
      // 类型说明:AI SDK v6 的 generateObject 是判别联合类型(基于 output mode),
      // TOutput 泛型在 runner 这层无法窄化,用 any 收口调用,runtime 由 Zod schema 校验。
      // 出参立刻 cast 回 TOutput,边界只在这一行。
      // eslint-disable-next-line @typescript-eslint/no-explicit-any
      const result = await generateObject<any>({
        model,
        schema: call.outputSchema,
        system,
        prompt,
      });

      const output = result.object as TOutput;

      // ─── 4. safety gate ───
      if (call.safetyRules && call.safetyRules.length > 0) {
        try {
          this.safety.check(output, call.safetyRules, {
            callKey: call.callKey,
            promptVersion: call.promptVersion,
            tenantId: ctx.tenantId,
          });
        } catch (err) {
          if (err instanceof SafetyGateRejectError) {
            return this.fallback(call, input, ctx, invocationId, `safety: ${err.message}`);
          }
          throw err;
        }
      }

      // ─── 5. 成功落账 ───
      const usage = result.usage ?? { inputTokens: 0, outputTokens: 0, totalTokens: 0 };
      const promptTokens = usage.inputTokens ?? 0;
      const completionTokens = usage.outputTokens ?? 0;
      const cachedInputTokens = (usage as { cachedInputTokens?: number }).cachedInputTokens ?? 0;
      const reasoningTokens = (usage as { reasoningTokens?: number }).reasoningTokens ?? 0;
      this.logger.log(
        `[usage] call=${call.callKey} usage=${JSON.stringify(usage)} ` +
          `providerMetadata=${JSON.stringify(result.providerMetadata ?? null)}`,
      );
      const costYuan = this.estimateCostYuan(modelId, promptTokens, completionTokens, cachedInputTokens);

      await this.recorder.end(invocationId, {
        output: output as Prisma.InputJsonValue,
        promptTokens,
        completionTokens,
        totalTokens: usage.totalTokens ?? promptTokens + completionTokens,
        cachedInputTokens,
        reasoningTokens,
        costYuan,
        status: 'succeeded',
      });

      this.cache.set(inputHash, output);

      return {
        output,
        invocationId,
        cacheHit: false,
        source: 'agent',
        promptTokens,
        completionTokens,
        costYuan,
      };
    } catch (err) {
      const errMsg = err instanceof Error ? err.message : String(err);
      this.logger.warn(`call=${call.callKey} LLM 失败: ${errMsg}`);
      return this.fallback(call, input, ctx, invocationId, `llm: ${errMsg}`);
    }
  }

  /**
   * 流式版本 — 返回 AsyncGenerator,逐步 yield partial 对象。
   *
   * 设计取舍:
   *   - 流式时**始终绕过缓存**(命中缓存的话流式没意义,客户端直接拿 done 即可,
   *     由 orchestrator/controller 决定是否走流式)
   *   - safety gate / fallback / 落库 都在 stream 收尾时统一做
   *   - generateObject 的 partial 是渐进式构造 JSON,字段顺序按 schema 定义
   */
  async *stream<TInput, TOutput>(
    call: AiCall<TInput, TOutput>,
    input: TInput,
    ctx: AiCallContext,
  ): AsyncGenerator<StreamEvent<TOutput>, void, void> {
    const inputHash = computeInputHash(call.callKey, call.promptVersion, input);
    const modelId = ctx.modelIdOverride ?? call.defaultModelId;
    const { model, provider } = this.provider.resolve(modelId);
    const { system, prompt } = call.buildPrompt(input);

    const invocationId = await this.recorder.start({
      ...this.baseStart(call, ctx, modelId, provider, inputHash, input, false),
      promptTemplate: prompt.length > 8000 ? prompt.slice(0, 8000) + '…[truncated]' : prompt,
    });

    yield { type: 'start', invocationId, modelId, promptVersion: call.promptVersion };

    try {
      // eslint-disable-next-line @typescript-eslint/no-explicit-any
      const stream = streamObject<any>({
        model,
        schema: call.outputSchema,
        system,
        prompt,
      });

      let lastPartial: Partial<TOutput> = {};
      for await (const partial of stream.partialObjectStream) {
        lastPartial = partial as Partial<TOutput>;
        yield { type: 'partial', partial: lastPartial };
      }

      // 等最终对象 + 计量(stream.object / stream.usage 都是 Promise)
      const finalObject = (await stream.object) as TOutput;
      const usage = await stream.usage;
      const promptTokens = usage?.inputTokens ?? 0;
      const completionTokens = usage?.outputTokens ?? 0;
      const cachedInputTokens = (usage as { cachedInputTokens?: number } | undefined)?.cachedInputTokens ?? 0;
      const reasoningTokens = (usage as { reasoningTokens?: number } | undefined)?.reasoningTokens ?? 0;
      // 流式时 providerMetadata 在 stream.providerMetadata(Promise)
      const providerMetadata = await (stream as { providerMetadata?: Promise<unknown> }).providerMetadata;
      this.logger.log(
        `[usage stream] call=${call.callKey} usage=${JSON.stringify(usage)} ` +
          `providerMetadata=${JSON.stringify(providerMetadata ?? null)}`,
      );
      const costYuan = this.estimateCostYuan(modelId, promptTokens, completionTokens, cachedInputTokens);

      // safety gate
      if (call.safetyRules && call.safetyRules.length > 0) {
        try {
          this.safety.check(finalObject, call.safetyRules, {
            callKey: call.callKey,
            promptVersion: call.promptVersion,
            tenantId: ctx.tenantId,
          });
        } catch (err) {
          if (err instanceof SafetyGateRejectError && call.fallback) {
            const fb = call.fallback(input, `safety: ${err.message}`);
            await this.recorder.end(invocationId, {
              output: fb as Prisma.InputJsonValue,
              status: 'failed',
              errorMessage: `safety rejected, fallback used: ${err.message}`,
              promptTokens,
              completionTokens,
              totalTokens: promptTokens + completionTokens,
              cachedInputTokens,
              reasoningTokens,
              costYuan,
            });
            yield {
              type: 'done',
              output: fb,
              source: 'template_fallback',
              invocationId,
              promptTokens,
              completionTokens,
              costYuan,
              fallbackReason: `safety: ${err.message}`,
            };
            return;
          }
          throw err;
        }
      }

      await this.recorder.end(invocationId, {
        output: finalObject as Prisma.InputJsonValue,
        promptTokens,
        completionTokens,
        totalTokens: promptTokens + completionTokens,
        cachedInputTokens,
        reasoningTokens,
        costYuan,
        status: 'succeeded',
      });
      this.cache.set(inputHash, finalObject);

      yield {
        type: 'done',
        output: finalObject,
        source: 'agent',
        invocationId,
        promptTokens,
        completionTokens,
        costYuan,
      };
    } catch (err) {
      const errMsg = err instanceof Error ? err.message : String(err);
      this.logger.warn(`stream call=${call.callKey} 失败: ${errMsg}`);

      if (call.fallback) {
        const fb = call.fallback(input, `llm: ${errMsg}`);
        await this.recorder.end(invocationId, {
          output: fb as Prisma.InputJsonValue,
          status: 'failed',
          errorMessage: `fallback used: ${errMsg}`,
        });
        yield {
          type: 'done',
          output: fb,
          source: 'template_fallback',
          invocationId,
          promptTokens: 0,
          completionTokens: 0,
          costYuan: 0,
          fallbackReason: errMsg,
        };
        return;
      }

      await this.recorder.end(invocationId, { status: 'failed', errorMessage: errMsg });
      yield { type: 'error', message: errMsg, invocationId };
    }
  }

  /**
   * fallback 走模板兜底。
   *   - 有 fallback → 用,标 source=template_fallback,invocation status=succeeded(因为兜底了)
   *   - 没 fallback → 落 failed,抛错给上层(上层决定是否硬失败 / 显占位)
   */
  private async fallback<TInput, TOutput>(
    call: AiCall<TInput, TOutput>,
    input: TInput,
    _ctx: AiCallContext,
    invocationId: string,
    reason: string,
  ): Promise<AiCallResult<TOutput>> {
    if (!call.fallback) {
      await this.recorder.end(invocationId, { status: 'failed', errorMessage: reason });
      throw new Error(`call=${call.callKey} 失败且无 fallback: ${reason}`);
    }
    const output = call.fallback(input, reason);
    await this.recorder.end(invocationId, {
      output: output as Prisma.InputJsonValue,
      status: 'failed',
      errorMessage: `fallback used: ${reason}`,
    });
    return {
      output,
      invocationId,
      cacheHit: false,
      source: 'template_fallback',
      fallbackReason: reason,
      promptTokens: 0,
      completionTokens: 0,
      costYuan: 0,
    };
  }

  private baseStart<TInput, TOutput>(
    call: AiCall<TInput, TOutput>,
    ctx: AiCallContext,
    modelId: string,
    provider: string,
    inputHash: string,
    input: TInput,
    cacheHit: boolean,
  ) {
    return {
      hostId: ctx.hostId,
      tenantId: ctx.tenantId,
      kind: call.kind,
      callKey: call.callKey,
      promptVersion: call.promptVersion,
      modelProvider: provider,
      modelName: modelId,
      workflowRunId: ctx.workflowRunId,
      parentInvocationId: ctx.parentInvocationId,
      inputHash,
      cacheHit,
      evalMode: ctx.evalMode ?? 'production',
      goldenSetCaseId: ctx.goldenSetCaseId,
      inputSnapshot: input as unknown as Prisma.InputJsonValue,
      linkedPatientId: ctx.linkedPatientId,
      linkedPersonaId: ctx.linkedPersonaId,
      linkedPlanId: ctx.linkedPlanId,
    };
  }

  /**
   * 成本估算(人民币元,直接落 Decimal(12,6) 列)。
   *
   * 价格表从 config(env AI_PRICE_TABLE_JSON)读 — 调价改 env 重启即可,代码不动;
   * 历史 invocation 行的 costYuan 已快照入库,不会被新价污染。
   *
   * cachedInputTokens 由 SDK 从 usage.cachedInputTokens 取(DeepSeek prompt cache 命中部分);
   * 为 0 时退化到全 cache miss 上限估算。
   *
   * 同 prompt + 同输入连续调用时 cache hit 比例上升,实际成本明显低于纯 miss 价。
   */
  private estimateCostYuan(
    modelId: string,
    promptTokens: number,
    completionTokens: number,
    cachedInputTokens: number = 0,
  ): number {
    const priceTable = this.config.get('ai', { infer: true }).priceTable;
    const p = priceTable[modelId] ?? priceTable['deepseek-v4-pro'] ?? { inHit: 0.5, inMiss: 3.6, out: 25 };
    // 防御:cached > prompt 不该发生,clamp
    const hit = Math.min(cachedInputTokens, promptTokens);
    const miss = Math.max(0, promptTokens - hit);
    const yuan = (hit * p.inHit + miss * p.inMiss + completionTokens * p.out) / 1_000_000;
    return Math.max(0, yuan);
  }
}
