import { Injectable, Logger } from '@nestjs/common';
import { and, desc, eq, sql } from 'drizzle-orm';
import { NlQueryDto, PaginationQuery, InventoryQuerySchema, AsnListQuerySchema, OrderListQuerySchema, TaskQuerySchema, ChargeQuerySchema } from '@wms/shared';
import { z } from 'zod';
import { DbService } from '@/common/db/db.service';
import { aiConversations, aiMessages, warehouses, storers, skus } from '@/db/schema';
import { requireRequestContext } from '@/common/context/request-context';
import { NotFoundError, ValidationError } from '@/common/errors/domain-error';
import { assertWarehouseAccess, resolveStorerScope, todayIso, serialize } from '@/common/util';
import { ReportsService } from '@/modules/reports/reports.service';
import { InventoryService } from '@/modules/inventory/inventory.service';
import { AsnService } from '@/modules/inbound/asn.service';
import { OrdersService } from '@/modules/outbound/orders.service';
import { TasksService } from '@/modules/tasks/tasks.service';
import { InvoicesService } from '@/modules/billing/invoices.service';
import { LlmService } from './llm/llm.service';
import { LlmMessage, LlmTool } from './llm/llm.provider';

type ToolDef = {
  name: string;
  description: string;
  schema: z.ZodObject<z.ZodRawShape>;
  run: (input: Record<string, unknown>) => Promise<unknown>;
};

const MAX_TOOL_ROUNDS = 6;
const MAX_TOOL_RESULT_CHARS = 12_000;
const HISTORY_TURNS = 12;

/**
 * Natural-language copilot over operational data.
 *
 * The model never sees SQL and never writes anything: it can only call a
 * fixed set of read-only tools, each of which goes through the same domain
 * services (and therefore the same warehouse / storer scoping, RLS and
 * permission checks) as the REST API. Every question, tool call and answer
 * is persisted so supervisors can audit what the assistant was asked and
 * what data it looked at.
 */
@Injectable()
export class CopilotService {
  private readonly logger = new Logger(CopilotService.name);

  constructor(
    private readonly db: DbService,
    private readonly llm: LlmService,
    private readonly reports: ReportsService,
    private readonly inventory: InventoryService,
    private readonly asnService: AsnService,
    private readonly ordersService: OrdersService,
    private readonly tasksService: TasksService,
    private readonly invoices: InvoicesService,
  ) {}

  async ask(dto: NlQueryDto) {
    await this.llm.assertAllowed('ai.nl_query');
    const ctx = requireRequestContext();
    if (dto.warehouseId) assertWarehouseAccess(dto.warehouseId);

    const conversation = await this.db.run(async (tx) => {
      if (dto.conversationId) {
        const [c] = await tx.select().from(aiConversations).where(and(eq(aiConversations.id, dto.conversationId), eq(aiConversations.userId, ctx.userId!))).limit(1);
        if (!c) throw new NotFoundError('Conversation', dto.conversationId);
        return c;
      }
      const [c] = await tx.insert(aiConversations).values({ tenantId: ctx.tenantId!, userId: ctx.userId!, warehouseId: dto.warehouseId ?? null, title: dto.question.slice(0, 80) }).returning();
      return c;
    });

    const history = await this.db.run((tx) =>
      tx.select({ role: aiMessages.role, content: aiMessages.content }).from(aiMessages).where(eq(aiMessages.conversationId, conversation.id)).orderBy(desc(aiMessages.createdAt)).limit(HISTORY_TURNS),
    );
    const context = await this.db.run(async (tx) => {
      const whs = await tx.select({ id: warehouses.id, code: warehouses.code, name: warehouses.name }).from(warehouses).where(eq(warehouses.status, 'ACTIVE'));
      const visible = ctx.warehouseIds === 'ALL' ? whs : whs.filter((w) => (ctx.warehouseIds as string[]).includes(w.id));
      const sts = await tx.select({ id: storers.id, code: storers.code, name: storers.name }).from(storers).where(and(eq(storers.status, 'ACTIVE'), ctx.storerId ? eq(storers.id, ctx.storerId) : undefined)).limit(50);
      return { warehouses: visible, storers: sts };
    });

    const tools = this.tools(context);
    const system = [
      'You are the operations copilot inside a warehouse management system (WMS). Answer questions about inbound, inventory, outbound, tasks, billing and KPIs using ONLY the tools provided. Never guess numbers; if a tool cannot answer, say so.',
      `Today is ${todayIso()}. The user is ${ctx.actorName ?? 'a warehouse user'}.`,
      `Warehouses the user can see: ${context.warehouses.map((w) => `${w.code} (${w.name}, id ${w.id})`).join('; ') || 'none'}.${dto.warehouseId ? ` Default warehouse id: ${dto.warehouseId}.` : ''}`,
      `Storers (customers of the warehouse): ${context.storers.map((s) => `${s.code} (${s.name}, id ${s.id})`).join('; ') || 'none'}.`,
      'Guidelines: be concise and factual, quantities with units, reference document numbers (ASN/order/task numbers) so the user can open them. Use plain text, short paragraphs or simple lists. If the question needs a warehouse and none is given and several are visible, ask which one. You cannot change any data; if the user asks to create or modify something, explain which screen or action does it.',
    ].join('\n');

    const messages: LlmMessage[] = [
      ...history.reverse().filter((m) => m.role === 'user' || m.role === 'assistant').map((m) => ({ role: m.role as 'user' | 'assistant', content: m.content })),
      { role: 'user', content: dto.question },
    ];

    const usage = { inputTokens: 0, outputTokens: 0 };
    const toolTrace: { name: string; input: unknown; ok: boolean; ms: number; preview: string }[] = [];
    let answer = '';
    let model = this.llm.model;

    for (let round = 0; round <= MAX_TOOL_ROUNDS; round++) {
      const res = await this.llm.complete({ system, messages, tools: tools.map(toLlmTool), maxTokens: 1500, temperature: 0 }, { feature: 'ai.nl_query', purpose: 'copilot' });
      usage.inputTokens += res.usage.inputTokens;
      usage.outputTokens += res.usage.outputTokens;
      model = res.model;

      if (res.stopReason !== 'tool_use' || !res.toolCalls.length) {
        answer = res.text || 'I could not produce an answer from the available data.';
        break;
      }
      if (round === MAX_TOOL_ROUNDS) {
        answer = res.text || 'I looked at several data sources but could not finish within the allowed number of steps. Please narrow the question.';
        break;
      }

      // Replay the assistant turn (text + tool_use blocks) then feed tool results.
      messages.push({ role: 'assistant', content: [...(res.text ? [{ type: 'text' as const, text: res.text }] : []), ...res.toolCalls.map((c) => ({ type: 'tool_use' as const, id: c.id, name: c.name, input: c.input }))] });
      const results: LlmMessage['content'] = [];
      for (const call of res.toolCalls) {
        const started = Date.now();
        const tool = tools.find((t) => t.name === call.name);
        let content: string;
        let ok = true;
        if (!tool) {
          ok = false;
          content = `Unknown tool ${call.name}`;
        } else {
          try {
            const input = tool.schema.parse(call.input);
            const out = await tool.run(input);
            content = truncate(JSON.stringify(serialize(out)), MAX_TOOL_RESULT_CHARS);
          } catch (e) {
            ok = false;
            content = `Error: ${(e as Error).message}`;
          }
        }
        toolTrace.push({ name: call.name, input: call.input, ok, ms: Date.now() - started, preview: content.slice(0, 200) });
        (results as Extract<LlmMessage['content'], unknown[]>).push({ type: 'tool_result', toolUseId: call.id, content, isError: !ok });
      }
      messages.push({ role: 'user', content: results });
    }

    await this.db.run(async (tx) => {
      await tx.insert(aiMessages).values([
        { tenantId: ctx.tenantId!, conversationId: conversation.id, role: 'user', content: dto.question },
        { tenantId: ctx.tenantId!, conversationId: conversation.id, role: 'assistant', content: answer, toolCalls: toolTrace, promptTokens: usage.inputTokens, completionTokens: usage.outputTokens },
      ]);
      await tx.update(aiConversations).set({ updatedAt: new Date() }).where(eq(aiConversations.id, conversation.id));
    });

    return { conversationId: conversation.id, answer, toolsUsed: toolTrace.map((t) => ({ name: t.name, ok: t.ok, ms: t.ms })), usage, model };
  }

  listConversations() {
    const ctx = requireRequestContext();
    return this.db.run((tx) => tx.select().from(aiConversations).where(eq(aiConversations.userId, ctx.userId!)).orderBy(desc(aiConversations.updatedAt)).limit(50));
  }

  conversation(id: string) {
    const ctx = requireRequestContext();
    return this.db.run(async (tx) => {
      const [c] = await tx.select().from(aiConversations).where(and(eq(aiConversations.id, id), eq(aiConversations.userId, ctx.userId!))).limit(1);
      if (!c) throw new NotFoundError('Conversation', id);
      const msgs = await tx.select().from(aiMessages).where(eq(aiMessages.conversationId, id)).orderBy(aiMessages.createdAt);
      return { ...c, messages: msgs };
    });
  }

  // ---- read-only tool catalogue ----

  private tools(context: { warehouses: { id: string; code: string }[]; storers: { id: string; code: string }[] }): ToolDef[] {
    const whId = (code?: string, id?: string) => {
      if (id) return id;
      if (code) {
        const w = context.warehouses.find((x) => x.code.toLowerCase() === code.toLowerCase());
        if (!w) throw new ValidationError(`Unknown warehouse code ${code}`);
        return w.id;
      }
      if (context.warehouses.length === 1) return context.warehouses[0].id;
      throw new ValidationError('warehouseCode is required (several warehouses visible)');
    };
    const stId = (code?: string) => {
      if (!code) return undefined;
      const s = context.storers.find((x) => x.code.toLowerCase() === code.toLowerCase());
      if (!s) throw new ValidationError(`Unknown storer code ${code}`);
      return resolveStorerScope(s.id);
    };
    const page = (n?: number) => PaginationQuery.parse({ page: 1, pageSize: Math.min(n ?? 25, 100) });

    return [
      {
        name: 'get_dashboard',
        description: 'Warehouse control-tower summary: inbound/outbound by status, open tasks, stock by status, utilisation, expiring stock and 30-day KPIs (dock-to-stock, order cycle, accuracy, on-time).',
        schema: z.object({ warehouseCode: z.string().optional() }),
        run: (i) => this.reports.dashboard(whId(i.warehouseCode as string)),
      },
      {
        name: 'inventory_lookup',
        description: 'Stock balances. Filter by SKU code, location code, LPN number, lot, status (AVAILABLE, DAMAGED, QUARANTINE, HOLD, EXPIRED). groupBy SKU for totals per item, LOCATION, LPN, LOT or BALANCE for full detail.',
        schema: z.object({ warehouseCode: z.string().optional(), storerCode: z.string().optional(), skuCode: z.string().optional(), locationCode: z.string().optional(), lpnNumber: z.string().optional(), lotNumber: z.string().optional(), status: z.string().optional(), expiryBefore: z.string().optional(), groupBy: z.enum(['BALANCE', 'SKU', 'LOCATION', 'LPN', 'LOT']).default('SKU'), limit: z.number().int().optional() }),
        run: (i) => this.inventory.query({ ...InventoryQuerySchema.parse({ warehouseId: i.warehouseCode || context.warehouses.length === 1 ? whId(i.warehouseCode as string | undefined) : undefined, storerId: stId(i.storerCode as string), skuCode: i.skuCode, locationCode: i.locationCode, lpnNumber: i.lpnNumber, lotNumber: i.lotNumber, status: i.status, expiryBefore: i.expiryBefore, groupBy: i.groupBy }), ...page(i.limit as number) }),
      },
      {
        name: 'find_asns',
        description: 'Inbound deliveries (ASNs). Filter by status (EXPECTED, SCHEDULED, ARRIVED, RECEIVING, RECEIVED, CLOSED, CANCELLED), external reference (PO number) or expected date range.',
        schema: z.object({ warehouseCode: z.string().optional(), storerCode: z.string().optional(), status: z.string().optional(), externalRef: z.string().optional(), expectedFrom: z.string().optional(), expectedTo: z.string().optional(), limit: z.number().int().optional() }),
        run: (i) => this.asnService.list({ ...AsnListQuerySchema.parse({ warehouseId: i.warehouseCode ? whId(i.warehouseCode as string) : undefined, storerId: stId(i.storerCode as string), status: i.status, externalRef: i.externalRef, expectedFrom: i.expectedFrom, expectedTo: i.expectedTo }), ...page(i.limit as number) }),
      },
      {
        name: 'get_asn',
        description: 'Full detail of one ASN by id, including lines, receipts and discrepancies.',
        schema: z.object({ asnId: z.string().uuid() }),
        run: (i) => this.asnService.get(i.asnId as string),
      },
      {
        name: 'find_orders',
        description: 'Outbound orders. Filter by status (OPEN, ON_HOLD, ALLOCATED, RELEASED, PICKING, PICKED, PACKED, STAGED, LOADED, SHIPPED, SHORT_CLOSED, CANCELLED), external reference or free text (customer / order number).',
        schema: z.object({ warehouseCode: z.string().optional(), storerCode: z.string().optional(), status: z.string().optional(), externalRef: z.string().optional(), q: z.string().optional(), limit: z.number().int().optional() }),
        run: (i) => this.ordersService.list({ ...OrderListQuerySchema.parse({ warehouseId: i.warehouseCode ? whId(i.warehouseCode as string) : undefined, storerId: stId(i.storerCode as string), status: i.status, externalRef: i.externalRef }), ...page(i.limit as number), q: i.q as string | undefined }),
      },
      {
        name: 'get_order',
        description: 'Full detail of one outbound order by id, including lines, allocations and shipment.',
        schema: z.object({ orderId: z.string().uuid() }),
        run: (i) => this.ordersService.get(i.orderId as string),
      },
      {
        name: 'task_queue',
        description: 'Open work queue summary per task type and status for a warehouse, and optionally the task list filtered by type/status.',
        schema: z.object({ warehouseCode: z.string().optional(), type: z.string().optional(), status: z.string().optional(), limit: z.number().int().optional() }),
        run: async (i) => {
          const warehouseId = whId(i.warehouseCode as string);
          const summary = await this.tasksService.summary(warehouseId);
          const list = i.type || i.status ? await this.tasksService.list({ ...TaskQuerySchema.parse({ warehouseId, type: i.type, status: i.status }), ...page(i.limit as number) }) : undefined;
          return { summary, list };
        },
      },
      {
        name: 'expiry_report',
        description: 'Stock expiring within N days (default 90), with lot, location and quantity.',
        schema: z.object({ warehouseCode: z.string().optional(), storerCode: z.string().optional(), days: z.number().int().min(0).max(730).default(90) }),
        run: (i) => this.reports.expiry(whId(i.warehouseCode as string), i.days as number, stId(i.storerCode as string)),
      },
      {
        name: 'stock_aging',
        description: 'Stock aging buckets (days since receipt) per SKU.',
        schema: z.object({ warehouseCode: z.string().optional(), storerCode: z.string().optional() }),
        run: (i) => this.reports.stockAging(whId(i.warehouseCode as string), stId(i.storerCode as string)),
      },
      {
        name: 'productivity',
        description: 'Completed tasks per user and type per day, with average minutes per task, for a date range (YYYY-MM-DD).',
        schema: z.object({ warehouseCode: z.string().optional(), from: z.string(), to: z.string() }),
        run: (i) => this.reports.productivity(whId(i.warehouseCode as string), i.from as string, i.to as string),
      },
      {
        name: 'billing_charges',
        description: 'Captured 3PL charges (handling, storage, VAS) for a storer and period, with status PENDING/INVOICED/WAIVED and tariff group.',
        schema: z.object({ storerCode: z.string(), warehouseCode: z.string().optional(), from: z.string().optional(), to: z.string().optional(), status: z.string().optional(), limit: z.number().int().optional() }),
        run: (i) => this.invoices.listCharges({ ...ChargeQuerySchema.parse({ storerId: stId(i.storerCode as string), warehouseId: i.warehouseCode ? whId(i.warehouseCode as string) : undefined, from: i.from, to: i.to, status: i.status }), ...page(i.limit as number) }),
      },
      {
        name: 'sku_lookup',
        description: 'Find SKUs (items) by code or description text for a storer.',
        schema: z.object({ storerCode: z.string().optional(), text: z.string().min(1), limit: z.number().int().optional() }),
        run: (i) =>
          this.db.run((tx) =>
            tx
              .select({ id: skus.id, code: skus.code, description: skus.description, storageCondition: skus.storageCondition, rotationRule: skus.rotationRule, lotControlled: skus.lotControlled, expiryControlled: skus.expiryControlled, storerId: skus.storerId })
              .from(skus)
              .where(and(stId(i.storerCode as string) ? eq(skus.storerId, stId(i.storerCode as string)!) : undefined, sql`(${skus.code} ilike ${'%' + (i.text as string) + '%'} or ${skus.description} ilike ${'%' + (i.text as string) + '%'})`))
              .limit(Math.min((i.limit as number) ?? 20, 50)),
          ),
      },
    ];
  }
}

function toLlmTool(t: ToolDef): LlmTool {
  return { name: t.name, description: t.description, inputSchema: zodToJsonSchema(t.schema) };
}

/** Small zod -> JSON schema converter for flat object tools (strings, numbers, enums, booleans). */
function zodToJsonSchema(schema: z.ZodObject<z.ZodRawShape>): Record<string, unknown> {
  const properties: Record<string, unknown> = {};
  const required: string[] = [];
  for (const [key, raw] of Object.entries(schema.shape)) {
    let t: z.ZodTypeAny = raw as z.ZodTypeAny;
    let optional = false;
    let def: unknown;
    for (;;) {
      if (t instanceof z.ZodOptional) { optional = true; t = t.unwrap(); continue; }
      if (t instanceof z.ZodDefault) { optional = true; def = t._def.defaultValue(); t = t._def.innerType; continue; }
      if (t instanceof z.ZodNullable) { optional = true; t = t.unwrap(); continue; }
      break;
    }
    let prop: Record<string, unknown>;
    if (t instanceof z.ZodEnum) prop = { type: 'string', enum: t.options };
    else if (t instanceof z.ZodNumber) prop = { type: t.isInt ? 'integer' : 'number' };
    else if (t instanceof z.ZodBoolean) prop = { type: 'boolean' };
    else prop = { type: 'string' };
    if (def !== undefined) prop.default = def;
    properties[key] = prop;
    if (!optional) required.push(key);
  }
  return { type: 'object', properties, required, additionalProperties: false };
}

function truncate(s: string, n: number): string {
  return s.length <= n ? s : `${s.slice(0, n)}... [truncated ${s.length - n} chars; narrow the query]`;
}