import {
  AbstractChat,
  type ChatInit,
  type ChatState,
  type ChatStatus,
  type UIMessage,
} from 'ai';

type Listener = () => void;

class ObservableChatState<UI_MESSAGE extends UIMessage> implements ChatState<UI_MESSAGE> {
  private messageValue: UI_MESSAGE[];
  private statusValue: ChatStatus = 'ready';
  private errorValue: Error | undefined;
  private readonly listeners = new Set<Listener>();

  constructor(messages: UI_MESSAGE[] = []) {
    this.messageValue = messages;
  }

  get status(): ChatStatus {
    return this.statusValue;
  }

  set status(value: ChatStatus) {
    this.statusValue = value;
    this.notify();
  }

  get error(): Error | undefined {
    return this.errorValue;
  }

  set error(value: Error | undefined) {
    this.errorValue = value;
    this.notify();
  }

  get messages(): UI_MESSAGE[] {
    return this.messageValue;
  }

  set messages(value: UI_MESSAGE[]) {
    this.messageValue = [...value];
    this.notify();
  }

  pushMessage = (message: UI_MESSAGE): void => {
    this.messageValue = [...this.messageValue, message];
    this.notify();
  };

  popMessage = (): void => {
    this.messageValue = this.messageValue.slice(0, -1);
    this.notify();
  };

  replaceMessage = (index: number, message: UI_MESSAGE): void => {
    this.messageValue = [
      ...this.messageValue.slice(0, index),
      this.snapshot(message),
      ...this.messageValue.slice(index + 1),
    ];
    this.notify();
  };

  snapshot = <T>(value: T): T => structuredClone(value);

  subscribe(listener: Listener): () => void {
    this.listeners.add(listener);
    return () => this.listeners.delete(listener);
  }

  private notify(): void {
    this.listeners.forEach((listener) => listener());
  }
}

export class FrameworkAgnosticChat<UI_MESSAGE extends UIMessage> extends AbstractChat<UI_MESSAGE> {
  private readonly observableState: ObservableChatState<UI_MESSAGE>;

  constructor(init: ChatInit<UI_MESSAGE>) {
    const state = new ObservableChatState(init.messages);
    const { messages: _messages, ...options } = init;
    super({ ...options, state });
    this.observableState = state;
  }

  subscribe(listener: Listener): () => void {
    return this.observableState.subscribe(listener);
  }
}