import { Clock, Duration, Effect, Schedule, Schema, Stream } from "effect"
import { AIError, TimeoutError } from "./schema/errors.js"

export const Status = Schema.Literals(["queued", "running", "completed", "failed", "cancelled", "expired"])
export type Status = Schema.Schema.Type<typeof Status>

/** Provider-neutral view of one generation observation. */
export interface Snapshot {
  readonly id: string
  readonly status: Status
  /** Normalized 0..1 when the provider reports progress. */
  readonly progress?: number
  readonly position?: number
}

/**
 * Route-owned generation operations for one generation. The media route decodes its serializable token once (from the
 * submission response or a `resume` input) and closes over it, so `Generation` never sees the token's shape.
 */
export interface Route<Response> {
  readonly status: Effect.Effect<Snapshot, AIError>
  readonly result: Effect.Effect<Response, AIError>
  readonly cancel?: Effect.Effect<void, AIError>
}

export interface Poll {
  readonly interval?: Duration.Input
  readonly timeout?: Duration.Input
}

export interface AwaitOptions {
  readonly poll?: Poll
}

const DEFAULT_POLL_INTERVAL = Duration.seconds(5)
const DEFAULT_POLL_TIMEOUT = Duration.minutes(10)

export const QueuedEvent = Schema.Struct({
  type: Schema.tag("generation-queued"),
  id: Schema.String,
  position: Schema.optional(Schema.Number),
}).annotate({ identifier: "Generation.Event.Queued" })

export const ProgressEvent = Schema.Struct({
  type: Schema.tag("generation-progress"),
  id: Schema.String,
  progress: Schema.optional(Schema.Number),
}).annotate({ identifier: "Generation.Event.Progress" })

export type Observation = Schema.Schema.Type<typeof QueuedEvent> | Schema.Schema.Type<typeof ProgressEvent>

export type Event = Observation | { readonly type: "generation-finished"; readonly id: string; readonly status: Status }

const TERMINAL: ReadonlySet<Status> = new Set(["completed", "failed", "cancelled", "expired"])

export const isTerminal = (status: Status) => TERMINAL.has(status)

export class Generation<Response> {
  readonly id: string
  readonly status: Status
  readonly progress?: number
  readonly position?: number

  constructor(
    readonly route: Route<Response>,
    /** Route-owned serializable JSON; pass it to the modality's `resume` from another process. */
    readonly token: unknown,
    snapshot: Snapshot,
  ) {
    this.id = snapshot.id
    this.status = snapshot.status
    this.progress = snapshot.progress
    this.position = snapshot.position
  }

  get snapshot(): Snapshot {
    return {
      id: this.id,
      status: this.status,
      progress: this.progress,
      position: this.position,
    }
  }

  get terminal() {
    return isTerminal(this.status)
  }

  refresh(): Effect.Effect<Generation<Response>, AIError> {
    return this.route.status.pipe(Effect.map((snapshot) => new Generation(this.route, this.token, snapshot)))
  }

  /** Fetch the result without polling; non-completed terminal generations fail with the provider's terminal body. */
  result(): Effect.Effect<Response, AIError> {
    return this.route.result
  }

  /** Poll until the generation reaches a terminal status, then fetch the result. Fails with a `Timeout` reason on deadline. */
  await(options?: AwaitOptions): Effect.Effect<Response, AIError> {
    const timeout = Duration.fromInputUnsafe(options?.poll?.timeout ?? DEFAULT_POLL_TIMEOUT)
    const settled = this.terminal ? Effect.succeed(this) : this.poll(options?.poll)
    return settled.pipe(
      // Non-completed terminal states also go through `result` so the route can surface its provider failure body.
      Effect.flatMap((generation) => generation.result()),
      Effect.timeoutOrElse({ duration: timeout, orElse: () => timeoutError(this.id, timeout) }),
    )
  }

  cancel(): Effect.Effect<void, AIError> {
    return this.route.cancel ?? Effect.void
  }

  /**
   * Status observations as a stream, ending after the first terminal observation. Each poll and each sleep between polls
   * is bounded by the time remaining until `poll.timeout`, so a hung status request or a long interval fails the stream at
   * the deadline instead of stalling it. (`Stream.interruptWhen` would express this directly but deadlocks under
   * `TestClock` when the source completes while the timer sleeps.)
   */
  events(options?: AwaitOptions): Stream.Stream<Event, AIError> {
    if (this.terminal) return Stream.make(this.event())
    const timeout = Duration.fromInputUnsafe(options?.poll?.timeout ?? DEFAULT_POLL_TIMEOUT)
    return Stream.unwrap(
      Clock.currentTimeMillis.pipe(
        Effect.map((start) => {
          const deadline = start + Duration.toMillis(timeout)
          const refresh = within(this.refresh(), this.id, timeout, deadline)
          const schedule = this.schedule(options?.poll).pipe(
            Schedule.modifyDelay((meta) =>
              Effect.succeed(Duration.min(meta.duration, Duration.millis(Math.max(0, deadline - meta.now)))),
            ),
          )
          return Stream.fromEffectSchedule(refresh, schedule).pipe(
            Stream.takeUntil((generation) => generation.terminal),
            Stream.map((generation) => generation.event()),
          )
        }),
      ),
    )
  }

  private event(): Event {
    if (this.terminal) return { type: "generation-finished", id: this.id, status: this.status }
    if (this.status === "queued") return { type: "generation-queued", id: this.id, position: this.position }
    return { type: "generation-progress", id: this.id, progress: this.progress }
  }

  private poll(poll: Poll | undefined) {
    return this.refresh().pipe(
      Effect.repeat({ schedule: this.schedule(poll), until: (generation) => generation.terminal }),
    )
  }

  private schedule(poll: Poll | undefined) {
    return Schedule.spaced(poll?.interval ?? DEFAULT_POLL_INTERVAL)
  }
}

/** `events` followed by the expanded result, with the result fetch bounded by the same `poll.timeout` deadline. */
export const resultEvents = <Response, A>(
  generation: Generation<Response>,
  expand: (response: Response) => ReadonlyArray<A>,
  options?: AwaitOptions,
): Stream.Stream<Observation | A, AIError> => {
  const timeout = Duration.fromInputUnsafe(options?.poll?.timeout ?? DEFAULT_POLL_TIMEOUT)
  return Stream.unwrap(
    Clock.currentTimeMillis.pipe(
      Effect.map((start) =>
        generation.events(options).pipe(
          Stream.filter((event): event is Observation => event.type !== "generation-finished"),
          Stream.concat(
            Stream.fromIterableEffect(
              within(generation.result(), generation.id, timeout, start + Duration.toMillis(timeout)).pipe(
                Effect.map(expand),
              ),
            ),
          ),
        ),
      ),
    ),
  )
}

/**
 * Run `effect` within the time left until `deadline`. Fails before starting once the deadline has passed: a fast
 * request could otherwise win the zero-budget race and schedule another zero-delay poll.
 */
const within = <A>(effect: Effect.Effect<A, AIError>, id: string, timeout: Duration.Duration, deadline: number) =>
  Clock.currentTimeMillis.pipe(
    Effect.flatMap((now) =>
      now >= deadline
        ? Effect.fail(timeoutError(id, timeout))
        : effect.pipe(
            Effect.timeoutOrElse({
              duration: Duration.millis(deadline - now),
              orElse: () => Effect.fail(timeoutError(id, timeout)),
            }),
          ),
    ),
  )

const timeoutError = (id: string, timeout: Duration.Duration) =>
  new AIError({
    reason: new TimeoutError({
      message: `Generation ${id} did not finish within ${Duration.format(timeout)}`,
      timeoutMs: Math.min(Number.MAX_SAFE_INTEGER, Math.max(0, Math.ceil(Duration.toMillis(timeout)))),
    }),
  })