import { cuidSchema, scoreIdSchema, sessionIdSchema, spanIdSchema, traceIdSchema } from "@domain/shared"
import { z } from "zod"
import { ANNOTATION_SCORE_PARTIAL_SOURCE_IDS, SCORE_SOURCE_ID_MAX_LENGTH, SCORE_SOURCES } from "../constants.ts"

const scoreSourceIdSchema = z.string().min(1).max(SCORE_SOURCE_ID_MAX_LENGTH)

export const scoreSourceSchema = z.enum(SCORE_SOURCES)
export type ScoreSource = z.infer<typeof scoreSourceSchema>

export const annotationScorePartialSourceIdSchema = z.enum(ANNOTATION_SCORE_PARTIAL_SOURCE_IDS)
export type AnnotationScorePartialSourceId = z.infer<typeof annotationScorePartialSourceIdSchema>

export const annotationScoreSourceIdSchema = z.union([annotationScorePartialSourceIdSchema, cuidSchema])
export type AnnotationScoreSourceId = z.infer<typeof annotationScoreSourceIdSchema>

export const scoreValueSchema = z.number().min(0).max(1)

export const ANNOTATION_ANCHOR_TEXT_FORMATS = ["pretty-json"] as const
export type AnnotationAnchorTextFormat = (typeof ANNOTATION_ANCHOR_TEXT_FORMATS)[number]

type AnnotationAnchorInput = {
  readonly messageIndex?: number | undefined
  readonly partIndex?: number | undefined
  readonly startOffset?: number | undefined
  readonly endOffset?: number | undefined
  readonly textFormat?: AnnotationAnchorTextFormat | undefined
}

const annotationAnchorFields = {
  messageIndex: z.number().int().nonnegative().optional(), // optional message index in the canonical conversation; omit for conversation-level annotations
  partIndex: z.number().int().nonnegative().optional(), // optional raw GenAI `parts[]` index inside the target message
  startOffset: z.number().int().nonnegative().optional(), // optional start offset for substring annotations within a textual part
  endOffset: z.number().int().nonnegative().optional(), // optional end offset for substring annotations within a textual part
  textFormat: z.enum(ANNOTATION_ANCHOR_TEXT_FORMATS).optional(), // optional UI-side text transform applied before the offsets were captured (e.g. prettified JSON); resolvers must apply the same transform before slicing
} as const

function validateAnnotationAnchor(anchor: AnnotationAnchorInput, ctx: z.core.$RefinementCtx<unknown>) {
  const hasMessageIndex = anchor.messageIndex !== undefined
  const hasPartIndex = anchor.partIndex !== undefined
  const hasStartOffset = anchor.startOffset !== undefined
  const hasEndOffset = anchor.endOffset !== undefined
  const hasTextRange = hasStartOffset || hasEndOffset

  if (!hasMessageIndex && hasPartIndex) {
    ctx.addIssue({
      code: "custom",
      message: "partIndex requires messageIndex",
      path: ["partIndex"],
    })
  }

  if (hasStartOffset !== hasEndOffset) {
    ctx.addIssue({
      code: "custom",
      message: "startOffset and endOffset must be provided together",
      path: hasStartOffset ? ["endOffset"] : ["startOffset"],
    })
  }

  if (hasTextRange && !hasPartIndex) {
    ctx.addIssue({
      code: "custom",
      message: "startOffset and endOffset require partIndex",
      path: ["partIndex"],
    })
  }

  if (anchor.startOffset !== undefined && anchor.endOffset !== undefined && anchor.startOffset > anchor.endOffset) {
    ctx.addIssue({
      code: "custom",
      message: "endOffset must be greater than or equal to startOffset",
      path: ["endOffset"],
    })
  }
}

export const annotationAnchorSchema = z.object(annotationAnchorFields).superRefine(validateAnnotationAnchor)
export type AnnotationAnchor = z.infer<typeof annotationAnchorSchema>

export const baseScoreMetadataSchema = z.object({
  // common metadata shared by all score sources
})
export type BaseScoreMetadata = z.infer<typeof baseScoreMetadataSchema>

export const evaluationScoreMetadataSchema = baseScoreMetadataSchema.extend({
  evaluationHash: z.string(), // hash of the script so we know whether this score was generated by the current version
})
export type EvaluationScoreMetadata = z.infer<typeof evaluationScoreMetadataSchema>

export const annotationScoreMetadataSchema = baseScoreMetadataSchema
  .extend({
    rawFeedback: z.string(), // original feedback text before enrichment; human-authored for human drafts/published annotations, model-authored for system-created drafts
    ...annotationAnchorFields,
  })
  .superRefine(validateAnnotationAnchor)
export type AnnotationScoreMetadata = z.infer<typeof annotationScoreMetadataSchema>

export const customScoreMetadataSchema = z.record(z.string(), z.unknown()) // whatever the user wants to store in custom score metadata
export type CustomScoreMetadata = z.infer<typeof customScoreMetadataSchema>

export const scoreMetadataSchemas = {
  evaluation: evaluationScoreMetadataSchema,
  annotation: annotationScoreMetadataSchema,
  custom: customScoreMetadataSchema,
} as const

export const baseScoreSchema = z.object({
  id: scoreIdSchema, // CUID score identifier
  organizationId: cuidSchema, // owning organization
  projectId: cuidSchema, // owning project
  sessionId: sessionIdSchema.nullable(), // optional session id inherited from instrumentation
  traceId: traceIdSchema.nullable(), // optional trace id inherited from instrumentation
  spanId: spanIdSchema.nullable(), // optional span id inherited from instrumentation
  simulationId: cuidSchema.nullable(), // optional simulation CUID link
  issueId: cuidSchema.nullable(), // optional issue CUID assignment
  value: scoreValueSchema, // normalized [0, 1] score value
  passed: z.boolean(), // true if passed, false if failed or errored
  feedback: z.string(), // clusterable feedback text used by issues
  error: z.string().min(1).nullable(), // canonical error text when the score generation truly errored
  errored: z.boolean(), // maintained in application/domain code on create or update
  duration: z.number().int().nonnegative(), // duration of score generation in nanoseconds
  tokens: z.number().int().nonnegative(), // total LLM token usage for this score generation
  cost: z.number().int().nonnegative(), // total LLM cost in microcents
  draftedAt: z.date().nullable(), // set while the score is still editable or awaiting human confirmation
  annotatorId: cuidSchema.nullable(), // user who created this score (nullable for system-generated scores)
  createdAt: z.date(), // score creation time
  updatedAt: z.date(), // score update time
})

function validateScoreLifecycle(
  score: {
    readonly error: string | null
    readonly errored: boolean
    readonly passed: boolean
  },
  ctx: z.core.$RefinementCtx<unknown>,
) {
  const hasError = score.error !== null

  if (score.errored !== hasError) {
    ctx.addIssue({
      code: "custom",
      message: "errored must match whether error is present",
      path: ["errored"],
      input: score.errored,
    })
  }

  if (score.passed && hasError) {
    ctx.addIssue({
      code: "custom",
      message: "passed scores cannot include an error",
      path: ["passed"],
      input: score.passed,
    })
  }
}

export const evaluationScoreSchema = baseScoreSchema.extend({
  source: z.literal("evaluation"),
  sourceId: cuidSchema, // evaluation cuid that produced this score
  metadata: evaluationScoreMetadataSchema,
})
export type EvaluationScore = z.infer<typeof evaluationScoreSchema>

export const annotationScoreSchema = baseScoreSchema.extend({
  source: z.literal("annotation"),
  sourceId: annotationScoreSourceIdSchema, // sentinel `"UI"` / `"API"` / `"SYSTEM"` for drafts and automation, or annotation-queue cuid for queue-authored rows
  metadata: annotationScoreMetadataSchema,
})
export type AnnotationScore = z.infer<typeof annotationScoreSchema>

export const customScoreSchema = baseScoreSchema.extend({
  source: z.literal("custom"),
  sourceId: scoreSourceIdSchema, // user-supplied tag for custom-source scores
  metadata: customScoreMetadataSchema,
})
export type CustomScore = z.infer<typeof customScoreSchema>

export const scoreSchema = z
  .discriminatedUnion("source", [evaluationScoreSchema, annotationScoreSchema, customScoreSchema])
  .superRefine(validateScoreLifecycle)
export type Score = z.infer<typeof scoreSchema>

export type ScoreMetadata = Score["metadata"]
