import { Inject, Injectable, Logger } from '@nestjs/common';
import { ZodSchema, z } from 'zod';
import { FeatureKey } from '@wms/shared';
import { DbService } from '@/common/db/db.service';
import { requireRequestContext } from '@/common/context/request-context';
import { DomainError, ErrorCodes } from '@/common/errors/domain-error';
import { PlatformService } from '@/modules/platform/platform.service';
import { todayIso, monthKey } from '@/common/util';
import { LLM_PROVIDER, LlmProvider, LlmRequest, LlmResponse } from './llm.provider';

export interface LlmCallMeta {
  /** Feature gate for this call; the tenant must have it enabled. */
  feature: FeatureKey;
  /** Short label for logs / audit (e.g. 'document_intake.asn'). */
  purpose: string;
}

export interface StructuredResult<T> {
  data: T;
  raw: string;
  usage: { inputTokens: number; outputTokens: number };
  model: string;
  attempts: number;
}

/**
 * Single entry point for every LLM call in the system.
 *
 * Guarantees:
 *  - feature gating per tenant (platform admins bypass),
 *  - monthly token budget from the subscription plan (aiTokensPerMonth; 0 = unlimited),
 *  - usage metering (AI_TOKENS) for billing and abuse control,
 *  - structured outputs validated with zod, with one corrective retry.
 */
@Injectable()
export class LlmService {
  private readonly logger = new Logger(LlmService.name);

  constructor(
    @Inject(LLM_PROVIDER) private readonly provider: LlmProvider,
    private readonly db: DbService,
    private readonly platform: PlatformService,
  ) {}

  get enabled(): boolean {
    return this.provider.name !== 'disabled';
  }

  get model(): string {
    return this.provider.model;
  }

  status() {
    return { enabled: this.enabled, provider: this.provider.name, model: this.provider.model };
  }

  /** Throws AI_DISABLED / FEATURE_DISABLED / PLAN_LIMIT before any tokens are spent. */
  async assertAllowed(feature: FeatureKey): Promise<void> {
    if (!this.enabled) throw new DomainError(ErrorCodes.AI_DISABLED, 'AI provider is not configured on this server.', 503);
    const ctx = requireRequestContext();
    if (ctx.actorType !== 'PLATFORM_ADMIN' && !ctx.features.includes(feature)) {
      throw new DomainError(ErrorCodes.FEATURE_DISABLED, `Feature ${feature} is not enabled for this tenant.`, 403);
    }
    if (ctx.tenantId) await this.assertBudget(ctx.tenantId);
  }

  async complete(req: LlmRequest, meta: LlmCallMeta): Promise<LlmResponse> {
    await this.assertAllowed(meta.feature);
    const started = Date.now();
    const res = await this.provider.complete(req);
    const tokens = res.usage.inputTokens + res.usage.outputTokens;
    this.logger.log(`${meta.purpose}: ${res.model} in=${res.usage.inputTokens} out=${res.usage.outputTokens} stop=${res.stopReason} ${Date.now() - started}ms`);
    await this.meter(tokens);
    return res;
  }

  /**
   * Ask for a JSON object matching `schema`. On a parse/validation failure the
   * model gets one corrective turn with the validation errors; after that the
   * caller receives a clear error (never a half-valid object).
   */
  async structured<T>(req: Omit<LlmRequest, 'json'>, schema: ZodSchema<T>, meta: LlmCallMeta): Promise<StructuredResult<T>> {
    let messages = [...req.messages];
    let usage = { inputTokens: 0, outputTokens: 0 };
    let lastError = '';
    let model = this.provider.model;
    for (let attempt = 1; attempt <= 2; attempt++) {
      const res = await this.complete({ ...req, messages, json: true }, meta);
      usage = { inputTokens: usage.inputTokens + res.usage.inputTokens, outputTokens: usage.outputTokens + res.usage.outputTokens };
      model = res.model;
      const parsed = parseJsonObject(res.text);
      if (parsed.ok) {
        const v = schema.safeParse(parsed.value);
        if (v.success) return { data: v.data, raw: res.text, usage, model, attempts: attempt };
        lastError = formatZodError(v.error);
      } else {
        lastError = parsed.error;
      }
      messages = [
        ...messages,
        { role: 'assistant', content: res.text || '(empty)' },
        { role: 'user', content: `Your previous answer was not valid. Problems:\n${lastError}\nReturn the corrected JSON object only.` },
      ];
    }
    throw new DomainError('AI_INVALID_OUTPUT', `The AI model did not return a valid structured answer: ${lastError}`, 502);
  }

  // ---- budget & metering ----

  private async assertBudget(tenantId: string) {
    const today = todayIso();
    const month = monthKey(today); // YYYY-MM
    await this.db.runBypass(async (tx) => {
      const limits = await this.platform.planLimits(tenantId, tx);
      const budget = Number(limits?.aiTokensPerMonth ?? 0);
      if (!budget) return; // 0 or missing = unlimited
      const used = await this.platform.usageTotal(tenantId, 'AI_TOKENS', `${month}-01`, today, tx);
      if (used >= budget) {
        throw new DomainError(ErrorCodes.PLAN_LIMIT, `Monthly AI token budget exhausted (${used.toLocaleString()} of ${budget.toLocaleString()}). Upgrade the plan or wait for the next period.`, 429);
      }
    });
  }

  private async meter(tokens: number) {
    const ctx = requireRequestContext();
    if (!ctx.tenantId || tokens <= 0) return;
    try {
      await this.db.runBypass((tx) => this.platform.recordUsage(ctx.tenantId!, 'AI_TOKENS', tokens, todayIso(), tx));
    } catch (e) {
      this.logger.error(`Failed to meter AI usage: ${(e as Error).message}`);
    }
  }
}

function parseJsonObject(text: string): { ok: true; value: unknown } | { ok: false; error: string } {
  const trimmed = text.trim().replace(/^```(?:json)?\s*/i, '').replace(/\s*```$/, '');
  const start = trimmed.indexOf('{');
  const end = trimmed.lastIndexOf('}');
  if (start < 0 || end <= start) return { ok: false, error: 'No JSON object found in the response.' };
  try {
    return { ok: true, value: JSON.parse(trimmed.slice(start, end + 1)) };
  } catch (e) {
    return { ok: false, error: `JSON parse error: ${(e as Error).message}` };
  }
}

function formatZodError(err: z.ZodError): string {
  return err.issues.slice(0, 10).map((i) => `- ${i.path.join('.') || '(root)'}: ${i.message}`).join('\n');
}