import { fromCallback, setup } from 'xstate';

import { createDecorator, IInstantiationService } from '#/_base/di/instantiation';
import { toDisposable, type IDisposable } from '#/_base/di/lifecycle';
import { ILogService } from '#/_base/log/log';
import {
  AgentActorService,
  type AgentActorContext,
  type AgentActorRestoreEvent,
} from '#/agent/actorService/agentActorService';
import { IAgentContextMemoryService } from '#/agent/contextMemory/contextMemory';
import { isCompactionSummaryMessage } from '#/agent/contextMemory/compactionHandoff';
import { ContextSpliced } from '#/agent/contextMemory/contextEvents';
import type { ContextMessage } from '#/agent/contextMemory/types';
import { IAgentLoopService, type BeforeStepContext } from '#/agent/loop/loop';
import { IAgentScopeContext } from '#/agent/scopeContext/scopeContext';
import { IEventBus } from '#/app/event/eventBus';
import { IEventDispatcher } from '#/state/eventDispatcher';

import { wrapSystemReminder } from './systemReminder';
import type {
  ContextInjectionContent,
  ContextInjectionContext,
  ContextInjectionMessage,
  ContextInjectionProvider,
  ContextInjectionResult,
  ReminderNotification,
  ReminderRegistration,
} from './types';

interface ReminderEntry {
  readonly provider: ContextInjectionProvider<unknown>;
  readonly variant: string;
}

const REMINDER_VARIANT_PRIORITY = new Map<string, number>([['date_change', -1]]);

interface ReminderActorContext {
  readonly entries: Set<ReminderEntry>;
  readonly runtime: AgentActorContext<null>;
}

interface ReminderRegisterEvent {
  readonly type: 'reminder.register';
  readonly entry: ReminderEntry;
}

interface ReminderUnregisterEvent {
  readonly type: 'reminder.unregister';
  readonly entry: ReminderEntry;
}

type ReminderActorEvent = AgentActorRestoreEvent | ReminderRegisterEvent | ReminderUnregisterEvent;

function actorContext(runtime: AgentActorContext<null>): ReminderActorContext {
  return runtime.getLogicState<ReminderActorContext>();
}

function appendReminder(
  runtime: AgentActorContext<null>,
  content: string,
  notification: ReminderNotification,
): void {
  runtime.get(IAgentContextMemoryService).append({
    role: 'user',
    content: [{ type: 'text', text: wrapSystemReminder(content) }],
    toolCalls: [],
    origin: {
      kind: 'injection',
      variant: notification.variant,
      ownerPromptId: notification.ownerPromptId,
    },
  });
}

function providerContext(
  runtime: AgentActorContext<null>,
  entry: ReminderEntry,
  isNewTurn: boolean,
): ContextInjectionContext<unknown> {
  const history = runtime.get(IAgentContextMemoryService).get();
  const injectedPositions = findInjections(history, entry.variant);
  const lastInjectedAt = injectedPositions.at(-1) ?? null;
  const lastInjection = lastInjectedAt === null ? undefined : history[lastInjectedAt];
  return {
    injectedPositions,
    lastInjectedAt,
    lastInjection,
    lastDisclosure:
      lastInjection?.origin?.kind === 'injection'
        ? lastInjection.origin.disclosure
        : undefined,
    isNewTurn,
  };
}

async function injectEntry(
  runtime: AgentActorContext<null>,
  entry: ReminderEntry,
  isNewTurn: boolean,
): Promise<void> {
  let content: Awaited<ReturnType<ContextInjectionProvider>>;
  try {
    content = await entry.provider(providerContext(runtime, entry, isNewTurn));
  } catch (error) {
    runtime.get(ILogService).error('context provider failed; skipping it', {
      name: entry.variant,
      error,
    });
    return;
  }
  if (!actorContext(runtime).entries.has(entry)) return;
  appendResult(runtime, entry, content);
}

function appendResult(
  runtime: AgentActorContext<null>,
  entry: ReminderEntry,
  content: ContextInjectionContent | ContextInjectionResult<unknown> | undefined,
): void {
  if (content === undefined) return;
  const result: ContextInjectionResult<unknown> = isInjectionResult(content)
    ? content
    : { content };
  const origin = {
    kind: 'injection' as const,
    variant: entry.variant,
    disclosure: result.disclosure,
  };
  const resolved = result.content;
  if (typeof resolved === 'string') {
    if (resolved.trim().length === 0) return;
    runtime.get(IAgentContextMemoryService).append({
      role: 'user',
      content: [{ type: 'text', text: wrapSystemReminder(resolved) }],
      toolCalls: [],
      origin,
    });
    return;
  }
  if (isRawInjectionMessage(resolved)) {
    const message = resolved.message;
    if (message.content.length === 0 && (message.tools === undefined || message.tools.length === 0)) {
      return;
    }
    runtime.get(IAgentContextMemoryService).append({
      role: message.role,
      content: [...message.content],
      toolCalls: [],
      tools: message.tools,
      origin,
    });
    return;
  }
  if (resolved.length === 0) return;
  runtime.get(IAgentContextMemoryService).append({
    role: 'user',
    content: [...resolved],
    toolCalls: [],
    origin,
  });
}

async function inject(runtime: AgentActorContext<null>, isNewTurn: boolean): Promise<void> {
  const entries = [...actorContext(runtime).entries].sort(
    (left, right) =>
      (REMINDER_VARIANT_PRIORITY.get(left.variant) ?? 0) -
      (REMINDER_VARIANT_PRIORITY.get(right.variant) ?? 0),
  );
  for (const entry of entries) await injectEntry(runtime, entry, isNewTurn);
}

const reminderEffects = fromCallback(({ input }: { input: { readonly runtime: AgentActorContext<null> } }) => {
  let compactionRearmPending = false;
  const loop = input.runtime.get(IAgentLoopService);
  const takeCompactionRearm = (): boolean => {
    const pending = compactionRearmPending;
    compactionRearmPending = false;
    return pending;
  };
  const reconcileAroundStep = async (
    context: BeforeStepContext,
    next: (context?: BeforeStepContext) => Promise<void>,
  ): Promise<void> => {
    const rearmed = takeCompactionRearm();
    await inject(input.runtime, context.firstStepOfTurn || rearmed);
    await next();
    if (takeCompactionRearm()) await inject(input.runtime, true);
  };
  let hook: IDisposable;
  try {
    hook = loop.hooks.onWillBeginStep.register('context-injector', reconcileAroundStep, {
      before: 'full-compaction',
    });
  } catch {
    hook = loop.hooks.onWillBeginStep.register('context-injector', reconcileAroundStep);
  }
  const splice = input.runtime.get(IEventBus).subscribe(ContextSpliced, (event) => {
    if (isCompactionSplice(event)) compactionRearmPending = true;
  });
  return () => {
    splice.dispose();
    hook.dispose();
    actorContext(input.runtime).entries.clear();
  };
});

const reminderActorLogic = setup({
  types: {} as {
    context: ReminderActorContext;
    input: AgentActorContext<null>;
    events: ReminderActorEvent;
  },
  actors: { reminderEffects },
}).createMachine({
  context: ({ input }) => ({ entries: new Set(), runtime: input }),
  initial: 'beforeRestore',
  states: {
    beforeRestore: {
      on: { 'runtime.restore': 'active' },
    },
    active: {
      invoke: {
        src: 'reminderEffects',
        input: ({ context }) => ({ runtime: context.runtime }),
      },
    },
  },
  on: {
    'reminder.register': {
      actions: ({ context, event }) => { context.entries.add(event.entry); },
    },
    'reminder.unregister': {
      actions: ({ context, event }) => { context.entries.delete(event.entry); },
    },
  },
});

export interface IAgentReminderService {
  readonly _serviceBrand: undefined;
  register<D = unknown>(variant: string, provider: ContextInjectionProvider<D>): ReminderRegistration;
  notify(content: string, notification: ReminderNotification): void;
  reconcileWhenIdle(variant: string): Promise<void>;
}

export const IAgentReminderService = createDecorator<IAgentReminderService>('agentReminderService');

export class AgentReminderService extends AgentActorService<null> implements IAgentReminderService {
  declare readonly _serviceBrand: undefined;

  private readonly actor: AgentActorContext<null>;
  private disposed = false;

  constructor(
    @IEventDispatcher dispatcher: IEventDispatcher,
    @IAgentScopeContext scopeContext: IAgentScopeContext,
    @IInstantiationService instantiation: IInstantiationService,
  ) {
    super(dispatcher, scopeContext, instantiation);
    this.actor = this.attachActor(reminderActorLogic, { id: 'reminder' });
    this._register(toDisposable(() => { this.disposed = true; }));
  }

  register<D = unknown>(variant: string, provider: ContextInjectionProvider<D>): ReminderRegistration {
    const entry: ReminderEntry = {
      provider: provider as ContextInjectionProvider<unknown>,
      variant,
    };
    this.actor.send({ type: 'reminder.register', entry });
    return toDisposable(() => {
      if (this.disposed) return;
      try {
        this.actor.send({ type: 'reminder.unregister', entry });
      } catch {}
    });
  }

  notify(content: string, notification: ReminderNotification): void {
    appendReminder(this.actor, content, notification);
  }

  async reconcileWhenIdle(variant: string): Promise<void> {
    const loop = this.actor.get(IAgentLoopService);
    const quiescence = loop.tryAcquireQuiescence();
    if (quiescence === undefined) return;
    try {
      for (const entry of actorContext(this.actor).entries) {
        if (entry.variant === variant) await injectEntry(this.actor, entry, false);
      }
    } finally {
      quiescence.dispose();
    }
  }
}

function isCompactionSplice(splice: {
  readonly deleteCount: number;
  readonly messages: readonly ContextMessage[];
}): boolean {
  return splice.deleteCount > 0 && splice.messages.some(isCompactionSummaryMessage);
}

function isRawInjectionMessage(
  content: Exclude<ContextInjectionContent, string>,
): content is { readonly message: ContextInjectionMessage } {
  return !Array.isArray(content);
}

function isInjectionResult(
  content: ContextInjectionContent | ContextInjectionResult<unknown>,
): content is ContextInjectionResult<unknown> {
  return typeof content === 'object' && content !== null && !Array.isArray(content) && 'content' in content;
}

function findInjections(history: readonly ContextMessage[], variant: string): number[] {
  const positions: number[] = [];
  history.forEach((message, index) => {
    if (message.origin?.kind === 'injection' && message.origin.variant === variant) positions.push(index);
  });
  return positions;
}