import { z } from 'zod';

export const KnowledgeBaseStatusSchema = z.enum(['draft', 'active', 'archived']);
export const KnowledgeSourceTypeSchema = z.enum(['url', 'file', 'manual', 'dataset', 'api']);
export const KnowledgeSourceStatusSchema = z.enum(['pending', 'indexed', 'failed', 'archived']);
export const KnowledgeChunkStatusSchema = z.enum(['pending', 'ready', 'archived']);

export const KnowledgeBaseSchema = z.object({
  id: z.string().min(1),
  productId: z.string().min(1),
  orgId: z.string().optional(),
  workspaceId: z.string().optional(),
  name: z.string().min(1),
  description: z.string().optional(),
  status: KnowledgeBaseStatusSchema,
  embeddingModel: z.string().optional(),
  retrievalStrategy: z.string().default('hybrid'),
  tags: z.array(z.string()).default([]),
  metadata: z.record(z.unknown()).optional(),
  createdBy: z.string().min(1),
  createdAt: z.string(),
  updatedAt: z.string(),
});

export type KnowledgeBaseDoc = z.infer<typeof KnowledgeBaseSchema> & {
  _ts?: number;
  _etag?: string;
};

export const KnowledgeSourceSchema = z.object({
  id: z.string().min(1),
  knowledgeBaseId: z.string().min(1),
  productId: z.string().min(1),
  type: KnowledgeSourceTypeSchema,
  status: KnowledgeSourceStatusSchema,
  title: z.string().min(1),
  uri: z.string().min(1).optional(),
  extractionJobId: z.string().optional(),
  contentHash: z.string().optional(),
  tags: z.array(z.string()).default([]),
  metadata: z.record(z.unknown()).optional(),
  createdAt: z.string(),
  updatedAt: z.string(),
});

export type KnowledgeSourceDoc = z.infer<typeof KnowledgeSourceSchema> & {
  _ts?: number;
  _etag?: string;
};

export const KnowledgeChunkSchema = z.object({
  id: z.string().min(1),
  knowledgeBaseId: z.string().min(1),
  sourceId: z.string().min(1),
  productId: z.string().min(1),
  ordinal: z.number().int().min(0),
  contentText: z.string().min(1),
  tokenCount: z.number().int().min(0).default(0),
  embeddingStatus: KnowledgeChunkStatusSchema,
  citations: z.array(z.string()).default([]),
  tags: z.array(z.string()).default([]),
  metadata: z.record(z.unknown()).optional(),
  createdAt: z.string(),
});

export type KnowledgeChunkDoc = z.infer<typeof KnowledgeChunkSchema> & {
  _ts?: number;
  _etag?: string;
};

export const CreateKnowledgeBaseSchema = z.object({
  orgId: z.string().optional(),
  workspaceId: z.string().optional(),
  name: z.string().min(1),
  description: z.string().optional(),
  embeddingModel: z.string().optional(),
  retrievalStrategy: z.string().default('hybrid'),
  tags: z.array(z.string()).default([]),
  metadata: z.record(z.unknown()).optional(),
});

export const UpdateKnowledgeBaseSchema = z.object({
  name: z.string().min(1).optional(),
  description: z.string().optional(),
  status: KnowledgeBaseStatusSchema.optional(),
  embeddingModel: z.string().optional(),
  retrievalStrategy: z.string().optional(),
  tags: z.array(z.string()).optional(),
  metadata: z.record(z.unknown()).optional(),
});

export const CreateKnowledgeSourceSchema = z.object({
  type: KnowledgeSourceTypeSchema,
  title: z.string().min(1),
  uri: z.string().min(1).optional(),
  extractionJobId: z.string().optional(),
  contentHash: z.string().optional(),
  tags: z.array(z.string()).default([]),
  metadata: z.record(z.unknown()).optional(),
});

export const UpsertKnowledgeChunksSchema = z.object({
  sourceStatus: KnowledgeSourceStatusSchema.optional(),
  chunks: z
    .array(
      z.object({
        ordinal: z.number().int().min(0),
        contentText: z.string().min(1),
        tokenCount: z.number().int().min(0).default(0),
        embeddingStatus: KnowledgeChunkStatusSchema.default('pending'),
        citations: z.array(z.string()).default([]),
        tags: z.array(z.string()).default([]),
        metadata: z.record(z.unknown()).optional(),
      })
    )
    .min(1),
});

export const ListKnowledgeBasesQuerySchema = z.object({
  status: KnowledgeBaseStatusSchema.optional(),
  limit: z.coerce.number().min(1).max(100).default(20),
});

export const RetrievalQuerySchema = z.object({
  query: z.string().min(1),
  limit: z.coerce.number().min(1).max(50).default(10),
  sourceId: z.string().optional(),
});

export type ListKnowledgeBasesQuery = z.infer<typeof ListKnowledgeBasesQuerySchema>;
