import { ScoreRepository, type TraceAnnotationCounts } from "@domain/scores"
import type { ProjectId, RepositoryError, SqlClient, TraceId } from "@domain/shared"
import type { Trace, TraceDetail } from "@domain/spans"
import { z } from "@hono/zod-openapi"
import { Effect } from "effect"
import { Paginated } from "../pagination.ts"

const nullableString = () => z.string().nullable()
const emptyToNull = (value: string): string | null => (value === "" ? null : value)

/**
 * Fields the trace repository accepts as `sortBy`. The API surface narrows the
 * domain's free-form string sort param to a stable allow-list so the SDK and
 * MCP tool inputs document exactly what's tunable.
 */
export const TRACE_SORT_FIELDS = ["startTime", "endTime", "durationNs", "tokensTotal", "costTotalMicrocents"] as const

/**
 * Opaque cursor over the wire — base64url JSON of `{ sortValue, traceId }`.
 * Keeps the public API surface a plain `string` while letting the ClickHouse
 * repo hand back its tuple cursor unchanged.
 */
export const encodeTraceCursor = (cursor: { sortValue: string; traceId: string }): string =>
  Buffer.from(JSON.stringify(cursor), "utf8").toString("base64url")

export const decodeTraceCursor = (raw: string): { sortValue: string; traceId: string } | null => {
  try {
    const json = Buffer.from(raw, "base64url").toString("utf8")
    const parsed = JSON.parse(json) as unknown
    if (
      parsed === null ||
      typeof parsed !== "object" ||
      typeof (parsed as { sortValue?: unknown }).sortValue !== "string" ||
      typeof (parsed as { traceId?: unknown }).traceId !== "string"
    ) {
      return null
    }
    return parsed as { sortValue: string; traceId: string }
  } catch {
    return null
  }
}

const traceFields = {
  organizationId: z.string().describe("Organization that owns this trace."),
  projectId: z.string().describe("Project this trace belongs to."),
  traceId: z.string().describe("32-character trace identifier."),
  spanCount: z.number().int().nonnegative().describe("Total number of spans in the trace."),
  errorCount: z.number().int().nonnegative().describe("Number of spans flagged with an error status."),
  startTime: z.string().describe("ISO-8601 timestamp of the trace's earliest span."),
  endTime: z.string().describe("ISO-8601 timestamp of the trace's latest span."),
  durationNs: z.number().describe("Wall-clock duration of the trace in nanoseconds."),
  timeToFirstTokenNs: z
    .number()
    .describe("Nanoseconds from the start of the first LLM span to its first emitted token. `0` if not measured."),
  tokensInput: z.number().describe("Total input tokens across LLM spans."),
  tokensOutput: z.number().describe("Total output tokens across LLM spans."),
  tokensCacheRead: z.number().describe("Total tokens served from the provider's prompt cache."),
  tokensCacheCreate: z.number().describe("Total tokens written to the provider's prompt cache."),
  tokensReasoning: z.number().describe("Total reasoning tokens reported by the model."),
  tokensTotal: z.number().describe("Sum of all token counters."),
  costInputMicrocents: z.number().describe("Cost of input tokens in microcents (1/1,000,000 USD)."),
  costOutputMicrocents: z.number().describe("Cost of output tokens in microcents (1/1,000,000 USD)."),
  costTotalMicrocents: z.number().describe("Total cost in microcents (1/1,000,000 USD)."),
  sessionId: nullableString().describe("Conversation/session identifier set by the SDK. `null` when absent."),
  userId: nullableString().describe("End-user identifier set by the SDK. `null` when absent."),
  simulationId: nullableString().describe(
    "CUID of the simulation that produced this trace. `null` when not a simulation.",
  ),
  tags: z.array(z.string()).describe("Free-form tags attached at ingest time."),
  models: z.array(z.string()).describe("Model identifiers seen across the trace's LLM spans."),
  providers: z.array(z.string()).describe("LLM-provider identifiers seen across the trace's spans."),
  serviceNames: z.array(z.string()).describe("OpenTelemetry `service.name` values seen in the trace."),
  rootSpanId: nullableString().describe(
    "Identifier of the trace's root span. `null` when no root span has been ingested.",
  ),
  rootSpanName: nullableString().describe(
    "`name` attribute of the root span. `null` when no root span has been ingested.",
  ),
  positiveAnnotationCount: z
    .number()
    .int()
    .nonnegative()
    .describe("Number of `passed = true` annotations attached to this trace."),
  negativeAnnotationCount: z
    .number()
    .int()
    .nonnegative()
    .describe("Number of `passed = false` annotations attached to this trace."),
} as const

const TraceSchema = z.object(traceFields).openapi("Trace")

export const PaginatedTracesSchema = Paginated(TraceSchema, "PaginatedTraces")

const GenAIMessageSchema = z
  .record(z.string(), z.unknown())
  .openapi("GenAIMessage")
  .describe("Message in OpenTelemetry GenAI format (`role` + content parts + optional tool calls).")

// `GenAISystem` from `rosetta-ai` is an *array* of part objects (`text`, `tool_definition`, …),
// not a single record. The previous record-shaped schema rejected real responses inside MCP
// clients with `expected: "record", received: "array"`. The array shape now matches the runtime
// value emitted by `toTraceDetailResponse` / `toSpanDetailResponse`.
const GenAISystemSchema = z
  .array(z.record(z.string(), z.unknown()))
  .openapi("GenAISystem")
  .describe("System instructions in OpenTelemetry GenAI format — an array of part objects.")

export const TraceDetailSchema = z
  .object({
    ...traceFields,
    metadata: z.record(z.string(), z.string()).describe("Free-form metadata attached at ingest time."),
    systemInstructions: GenAISystemSchema,
    inputMessages: z
      .array(GenAIMessageSchema)
      .describe("Input messages sent into the first LLM span of the trace, in OpenTelemetry GenAI format."),
    outputMessages: z
      .array(GenAIMessageSchema)
      .describe("Output messages from the last LLM span of the trace, in OpenTelemetry GenAI format."),
    allMessages: z
      .array(GenAIMessageSchema)
      .describe(
        "Full conversation view for the trace's final turn — the last span's input messages followed by its output messages, in OpenTelemetry GenAI format.",
      ),
  })
  .openapi("TraceDetail")

// Indicators data — positive/negative annotation counts shown in the trace table's
// "Indicators" column — comes from a separate Postgres read (the `scores` table), not
// from the ClickHouse-backed trace row. Shared across every trace-returning endpoint so
// the inline indicators on each `TraceSchema` row are fetched in exactly one place:
// handlers call `fetchTraceIndicators` after listing traces, then pass the resulting
// `Map` to `toTraceResponse` / `toTraceDetailResponse`. Missing entries default to `0/0`
// inside the mapper so callers never see `undefined` for these fields.
type TraceIndicators = ReadonlyMap<string, TraceAnnotationCounts>

export const fetchTraceIndicators = (input: {
  readonly projectId: ProjectId
  readonly traceIds: readonly TraceId[]
}): Effect.Effect<TraceIndicators, RepositoryError, ScoreRepository | SqlClient> =>
  Effect.gen(function* () {
    if (input.traceIds.length === 0) return new Map<string, TraceAnnotationCounts>()
    const scoreRepo = yield* ScoreRepository
    const counts = yield* scoreRepo.countAnnotationsByTraceIds(input)
    return new Map(counts.map((entry) => [entry.traceId as string, entry] as const))
  })

// ClickHouse doesn't support NULL in our trace columns, so the domain entity
// uses `""` as the absent sentinel for SDK-optional ids. Normalise here so
// the public API surface exposes a proper `string | null` instead of leaking
// the storage-layer encoding.
export const toTraceResponse = (trace: Trace, indicators?: TraceIndicators | TraceAnnotationCounts) => {
  const traceIndicators =
    indicators instanceof Map ? indicators.get(trace.traceId as string) : (indicators ?? undefined)
  return {
    organizationId: trace.organizationId as string,
    projectId: trace.projectId as string,
    traceId: trace.traceId as string,
    spanCount: trace.spanCount,
    errorCount: trace.errorCount,
    startTime: trace.startTime.toISOString(),
    endTime: trace.endTime.toISOString(),
    durationNs: trace.durationNs,
    timeToFirstTokenNs: trace.timeToFirstTokenNs,
    tokensInput: trace.tokensInput,
    tokensOutput: trace.tokensOutput,
    tokensCacheRead: trace.tokensCacheRead,
    tokensCacheCreate: trace.tokensCacheCreate,
    tokensReasoning: trace.tokensReasoning,
    tokensTotal: trace.tokensTotal,
    costInputMicrocents: trace.costInputMicrocents,
    costOutputMicrocents: trace.costOutputMicrocents,
    costTotalMicrocents: trace.costTotalMicrocents,
    sessionId: emptyToNull(trace.sessionId as string),
    userId: emptyToNull(trace.userId as string),
    simulationId: emptyToNull(trace.simulationId as string),
    tags: [...trace.tags],
    models: [...trace.models],
    providers: [...trace.providers],
    serviceNames: [...trace.serviceNames],
    rootSpanId: emptyToNull(trace.rootSpanId as string),
    rootSpanName: emptyToNull(trace.rootSpanName),
    positiveAnnotationCount: traceIndicators?.positiveCount ?? 0,
    negativeAnnotationCount: traceIndicators?.negativeCount ?? 0,
  }
}

// rosetta-ai's `GenAIMessage` / `GenAISystem` types declare more specific
// shapes than `Record<string, unknown>` and TS rejects the direct cast as
// non-overlapping. Going through `unknown` is the documented escape hatch —
// the payload IS a JSON object at runtime, the cast is purely a type-level
// bridge for the response schema.
export const toTraceDetailResponse = (trace: TraceDetail, indicators?: TraceIndicators | TraceAnnotationCounts) => ({
  ...toTraceResponse(trace, indicators),
  metadata: { ...trace.metadata },
  systemInstructions: trace.systemInstructions.map((part) => part as unknown as Record<string, unknown>),
  inputMessages: trace.inputMessages.map((m) => m as unknown as Record<string, unknown>),
  outputMessages: trace.outputMessages.map((m) => m as unknown as Record<string, unknown>),
  allMessages: trace.allMessages.map((m) => m as unknown as Record<string, unknown>),
})
