import { Effect, JsonSchema, Schema, Stream } from "effect"
import { tryRequest } from "./media-model.js"
import {
  LLMClient,
  Service,
  canCompact,
  compactRequest,
  type CheckpointRequest,
  type CompactionRequest,
  type EndpointCompactOptions,
  type StreamOptions,
  type TriggerCompactOptions,
} from "./route/client.js"
import {
  InvalidProviderOutputError,
  InvalidRequestError,
  mergeGeneration,
  AIError,
  LLMEvent,
  LLMResponse,
  ResponseFold,
  Message,
  LanguageModel,
  SystemPart,
  type BaseRequest,
  type CachePolicy,
  type CompactionCheckpointResponse,
  type CompactionResponse,
  type Finish,
  ToolChoice,
  ToolEntry,
  Usage,
  type LanguageModelProviderOptions,
  type GenerationOptionsFields,
  type LLMRequest,
  type ToolEntryInput,
  withPartial,
} from "./schema/index.js"
import { tool, type ExecutableTool, type Tool, type ToolSchema } from "./tool.js"
import { runTools } from "./tool-runtime.js"
import { addUsage, nextRequest, pendingCalls } from "./loop.js"

/** Input accepted by `ai.request`, `ai.generate`, `ai.stream` and `ai.compact`. An `LLMRequest` is also a valid input. */
export type LLMRequestInput<Model extends LanguageModel = LanguageModel> = TextRequestFields<Model> & PromptOrMessages

/** Exactly one of `prompt`, a single user message, or `messages`, the full conversation. */
export type PromptOrMessages =
  | { readonly prompt: Message.UserContentInput; readonly messages?: never }
  | { readonly messages: ReadonlyArray<Message | Message.Input>; readonly prompt?: never }

interface TextRequestFields<Model extends LanguageModel>
  extends BaseRequest<Model, NoInfer<LanguageModelProviderOptions<Model>>>,
    GenerationOptionsFields {
  readonly id?: string
  readonly system?: string | SystemPart | ReadonlyArray<SystemPart>
  /** Sent in array order. Pass `tool(...)` values or wire entries such as hosted provider tools. */
  readonly tools?: ReadonlyArray<Tool | ToolEntryInput>
  readonly toolChoice?: ToolChoice.Input
  readonly cache?: CachePolicy
  /** Stable cache affinity for protocols that support provider-managed prompt caching. */
  readonly promptCacheKey?: string
  /** Emit a `raw` event with each provider frame before the events parsed from it. */
  readonly includeRaw?: boolean
}

/** An Effect Schema decodes and types `object`; a raw JSON Schema leaves it `unknown`. */
export type OutputSchema = ToolSchema<any> | JsonSchema.JsonSchema

export type OutputType<S extends OutputSchema> = S extends Schema.Top ? Schema.Schema.Type<S> : unknown

export type ObjectResponse<T> = LLMResponse & { readonly object: T }

/** Stream events for a request with `schema`: the `finish` event carries the decoded `object`. */
export type ObjectEvent<T> = Exclude<LLMEvent, Finish> | (Finish & { readonly object: T })

/** Loop events for a request with `schema`: the loop's `finish` event carries the decoded `object`. */
export type LoopObjectEvent<T> = Exclude<LLMEvent, Finish> | (Finish & { readonly object?: T })

/** Run one model call. With `schema`, the model must answer through a forced tool call decoded into `object`. */
export function generate<const Model extends LanguageModel, S extends OutputSchema>(
  input: LLMRequestInput<Model> & { readonly schema: S },
  options?: StreamOptions,
): Effect.Effect<ObjectResponse<OutputType<S>>, AIError, Service>
export function generate<const Model extends LanguageModel>(
  input: LLMRequestInput<Model>,
  options?: StreamOptions,
): Effect.Effect<LLMResponse, AIError, Service>
export function generate(input: Input, options?: StreamOptions) {
  return prepare(input).pipe(
    Effect.flatMap(({ request, output }) =>
      LLMClient.generate(request, options).pipe(
        Effect.flatMap((response) =>
          output === undefined
            ? Effect.succeed(response)
            : decodeObject(output, response).pipe(Effect.map((object) => Object.assign(response, { object }))),
        ),
      ),
    ),
  )
}

/**
 * Stream one model call. The `finish` event carries the assistant `message` and `usage`; with `schema`, also the
 * decoded `object`.
 */
export function stream<const Model extends LanguageModel, S extends OutputSchema>(
  input: LLMRequestInput<Model> & { readonly schema: S },
  options?: StreamOptions,
): Stream.Stream<ObjectEvent<OutputType<S>>, AIError, Service>
export function stream<const Model extends LanguageModel>(
  input: LLMRequestInput<Model>,
  options?: StreamOptions,
): Stream.Stream<LLMEvent, AIError, Service>
export function stream(input: Input, options?: StreamOptions) {
  return Stream.unwrap(
    prepare(input).pipe(
      Effect.map(({ request, output }) =>
        stepStream(request, undefined, options, (step) =>
          Stream.fromEffect(finishStep(step, step.response.usage, output)),
        ),
      ),
    ),
  )
}

export type LLMLoopInput<Model extends LanguageModel = LanguageModel> = Omit<TextRequestFields<Model>, "tools"> &
  PromptOrMessages & {
    readonly tools: ReadonlyArray<ExecutableTool>
    /** Called after each step that has tool calls; returning `true` ends the loop without running them. */
    readonly until?: (state: { readonly steps: ReadonlyArray<LLMResponse> }) => boolean
  }

/**
 * Stream model calls, running the model's tool calls between them, until a step has no client-executed tool calls,
 * hits the output token limit, or `until` returns `true`. Tool calls whose input is not valid JSON get an error result
 * so the model can retry. A forced `toolChoice` applies to the first step only. Each step's `step-start` and `step-finish` carry its step index and usage; one `finish` ends
 * the loop with the final step's reason and `message` and the usage summed across steps. With `schema`, the model ends
 * by calling the hidden answer tool, decoded into `object` on the `finish`.
 */
export function loop<const Model extends LanguageModel, S extends OutputSchema>(
  input: LLMLoopInput<Model> & { readonly schema: S },
  options?: StreamOptions,
): Stream.Stream<LoopObjectEvent<OutputType<S>>, AIError, Service>
export function loop<const Model extends LanguageModel>(
  input: LLMLoopInput<Model>,
  options?: StreamOptions,
): Stream.Stream<LLMEvent, AIError, Service>
export function loop(input: LLMLoopInput & { readonly schema?: OutputSchema }, options?: StreamOptions) {
  const { tools, until, ...rest } = input
  return Stream.unwrap(
    prepare({ ...rest, tools }).pipe(
      Effect.map(({ request, output }) => driveLoop({ tools, until, output, options }, request, [], undefined)),
    ),
  )
}

export { collect, type Collected, type CollectResult } from "./loop.js"

interface Step {
  readonly response: LLMResponse
  readonly finish: Finish
}

/**
 * Stream one model call. Its `finish` is held back: `after` gets the completed step and streams what follows it.
 * With `index`, the step's `step-start` and `step-finish` carry it.
 */
const stepStream = (
  request: LLMRequest,
  index: number | undefined,
  options: StreamOptions | undefined,
  after: (step: Step) => Stream.Stream<LLMEvent, AIError, Service>,
): Stream.Stream<LLMEvent, AIError, Service> =>
  Stream.suspend(() => {
    let state = ResponseFold.empty()
    let finish: Finish | undefined
    return LLMClient.stream(request, options).pipe(
      Stream.map((event) =>
        index !== undefined && (event.type === "step-start" || event.type === "step-finish")
          ? { ...event, index }
          : event,
      ),
      Stream.tap((event) =>
        Effect.sync(() => {
          state = ResponseFold.reduce(state, event)
          if (event.type === "finish") finish = event
        }),
      ),
      Stream.filter((event) => event.type !== "finish"),
      Stream.mapError((error) => withPartial(error, ResponseFold.partial(state))),
      Stream.concat(
        Stream.suspend(() => {
          const response = ResponseFold.complete(state)
          return response && finish ? after({ response, finish }) : Stream.empty
        }),
      ),
    )
  })

/** The `finish` that ends a stream or loop: the final step's `message`, the given `usage`, and the decoded `object`. */
const finishStep = (
  step: Step,
  usage: Usage | undefined,
  output: Tool | undefined,
): Effect.Effect<LLMEvent, AIError> => {
  const event: Finish = { ...step.finish, usage, message: step.response.message }
  if (output === undefined) return Effect.succeed(event)
  return decodeObject(output, step.response).pipe(Effect.map((object) => ({ ...event, object })))
}

interface LoopContext {
  readonly tools: ReadonlyArray<ExecutableTool>
  readonly until: LLMLoopInput["until"]
  readonly output: Tool | undefined
  readonly options: StreamOptions | undefined
}

const driveLoop = (
  context: LoopContext,
  request: LLMRequest,
  steps: ReadonlyArray<LLMResponse>,
  usage: Usage | undefined,
): Stream.Stream<LLMEvent, AIError, Service> =>
  stepStream(request, steps.length, context.options, (step) => {
    const total = addUsage(usage, step.response.usage)
    const calls = pendingCalls(step.response, context.output && OBJECT_TOOL)
    const done = [...steps, step.response]
    if (calls.length === 0 || context.until?.({ steps: done }))
      return Stream.fromEffect(finishStep(step, total, calls.length === 0 ? context.output : undefined))
    return Stream.unwrap(
      runTools(context.tools, calls).pipe(
        Effect.map((results) =>
          Stream.fromIterable(results.flatMap((result) => result.events)).pipe(
            Stream.concat(
              driveLoop(
                context,
                nextRequest(request, step.response, calls, results, context.output !== undefined),
                done,
                total,
              ),
            ),
          ),
        ),
      ),
    )
  })

const OBJECT_TOOL = "generate_object"

type Input = LLMRequestInput & { readonly schema?: OutputSchema }

// `schema` adds a hidden tool whose input is the answer. With other tools the model only has to call some tool.
const prepare = (input: Input) => {
  const { schema, ...rest } = input
  return tryRequest(() => {
    const description = "Return the structured result by calling this tool."
    const output: Tool | undefined =
      schema === undefined ? undefined : tool({ name: OBJECT_TOOL, description, input: schema })
    const tools = rest.tools ?? []
    const built = request(
      output === undefined
        ? rest
        : {
            ...rest,
            tools: [...tools, output],
            toolChoice: tools.length === 0 ? ToolChoice.named(OBJECT_TOOL) : new ToolChoice({ type: "required" }),
          },
    )
    const duplicate = built.tools.find(
      (entry, index) => built.tools.findIndex((other) => other.name === entry.name) !== index,
    )
    if (duplicate) throw new Error(`Duplicate tool name: ${duplicate.name}`)
    return { request: built, output }
  })
}

const decodeObject = (output: Tool, response: LLMResponse) => {
  const call = response.toolCalls.find((event) => event.name === OBJECT_TOOL)
  const invalid = (message: string, cause?: unknown) =>
    new AIError({ reason: new InvalidProviderOutputError({ message, cause }), partial: response })
  // A cut-off call arrives with truncated or empty input, so its decode error would hide the real cause.
  if (response.finishReason.normalized === "length")
    return Effect.fail(invalid("output hit the token limit before the forced tool call completed"))
  if (call)
    return output
      ._decode(call.input)
      .pipe(Effect.mapError((error) => invalid(`schema decode failed: ${error.message}`, error)))
  return Effect.fail(invalid(`model did not call the forced \`${OBJECT_TOOL}\` tool`))
}

/** Normalize text request input into a plain `LLMRequest`, with `prompt` turned into the one user message. */
export const request = <const Model extends LanguageModel>(input: LLMRequestInput<Model>): LLMRequest<Model> => {
  if ((input.prompt === undefined) === (input.messages === undefined))
    throw new AIError({
      reason: new InvalidRequestError({ message: "Pass exactly one of `prompt` or `messages`, not both or neither" }),
    })
  return {
    id: input.id,
    model: input.model,
    system: SystemPart.content(input.system),
    messages: input.prompt === undefined ? (input.messages ?? []).map(Message.make) : [Message.user(input.prompt)],
    tools: input.tools?.map((item) => ("_definition" in item ? item._definition : ToolEntry.make(item))) ?? [],
    toolChoice: input.toolChoice ? ToolChoice.make(input.toolChoice) : undefined,
    ...mergeGeneration(input),
    providerOptions: input.providerOptions,
    headers: input.headers,
    fetch: input.fetch,
    timeout: input.timeout,
    http: input.http,
    cache: input.cache,
    promptCacheKey: input.promptCacheKey,
    includeRaw: input.includeRaw,
  }
}

/** Compact a text request: a replacement window, or a checkpoint with `mechanism: "trigger"`. */
export function compact(
  input: CheckpointRequest,
  options: TriggerCompactOptions,
): Effect.Effect<CompactionCheckpointResponse, AIError, Service>
export function compact(
  input: CompactionRequest,
  options?: EndpointCompactOptions,
): Effect.Effect<CompactionResponse, AIError, Service>
export function compact(
  input: CheckpointRequest | CompactionRequest,
  options?: EndpointCompactOptions | TriggerCompactOptions,
) {
  return tryRequest(() => request<LanguageModel>(input)).pipe(
    Effect.flatMap((normalized) => compactRequest(normalized, options)),
  )
}

export { canCompact }

export * as LLM from "./llm.js"