import type {
  ChatAssistantMessage,
  ChatMessagePart,
  StreamOperation,
  TransientPart,
} from "../types";

export interface StreamReducerState {
  message: ChatAssistantMessage | null;
  currentStepId?: string;
}

export function streamReducer(
  state: StreamReducerState,
  operation: StreamOperation,
): StreamReducerState {
  const { message, currentStepId } = state;

  switch (operation.type) {
    case "messageStart": {
      const id = operation.id;
      return {
        message: {
          id,
          sender: "bot",
          type: "assistant",
          parts: [],
          transientParts: [],
        },
      };
    }

    case "messageFinish": {
      if (!message) return state;
      return {
        ...state,
        message: {
          ...message,
          transientParts: [],
        },
      };
    }

    case "messageError": {
      if (!message) {
        return {
          message: {
            id: `error-msg`,
            sender: "bot",
            type: "assistant",
            parts: [
              {
                type: "text",
                id: `error-msg-text`,
                text: operation.message,
                status: "error",
              },
            ],
            transientParts: [],
          },
        };
      }
      const errorPart: ChatMessagePart = {
        type: "text",
        id: `${message.id}-error`,
        text: operation.message,
        status: "error",
      };
      return {
        ...state,
        message: {
          ...message,
          parts: [...message.parts, errorPart],
          transientParts: [],
        },
      };
    }

    case "stepStart": {
      if (!message) return state;
      const stepPart: ChatMessagePart = {
        type: "step",
        id: operation.id,
        label: operation.label,
        status: "running",
      };
      return {
        message: {
          ...message,
          parts: upsertPart(message.parts, stepPart),
        },
        currentStepId: operation.id,
      };
    }

    case "stepEnd": {
      if (!message) return state;
      const parts = message.parts.map((p) => {
        if (p.type === "step" && p.id === operation.id) {
          return {
            ...p,
            status: operation.status ?? "completed",
          };
        }
        return p;
      });

      const nextStepId =
        operation.id === currentStepId ? undefined : currentStepId;

      return {
        message: { ...message, parts },
        currentStepId: nextStepId,
      };
    }

    case "textStart": {
      if (!message) return state;
      const textPart: ChatMessagePart = {
        type: "text",
        id: operation.id,
        text: "",
        status: "streaming",
        stepId: currentStepId,
      };
      return {
        ...state,
        message: {
          ...message,
          parts: upsertPart(message.parts, textPart),
        },
      };
    }

    case "textDelta": {
      if (!message) return state;
      const parts = message.parts.map((p) => {
        if (p.type === "text" && p.id === operation.id) {
          return {
            ...p,
            text: p.text + operation.delta,
          };
        }
        return p;
      });
      return {
        ...state,
        message: { ...message, parts },
      };
    }

    case "textEnd": {
      if (!message) return state;
      const parts = message.parts.map((p) => {
        if (p.type === "text" && p.id === operation.id) {
          return {
            ...p,
            status: "done" as const,
          };
        }
        return p;
      });
      return {
        ...state,
        message: { ...message, parts },
      };
    }

    case "statusUpsert": {
      if (!message) return state;
      const transientPart: TransientPart = {
        type: "data-status",
        id: operation.id,
        message: operation.message,
        status: operation.status,
        stepId: currentStepId,
      };
      return {
        ...state,
        message: {
          ...message,
          transientParts: upsertTransientPart(
            message.transientParts ?? [],
            transientPart,
          ),
        },
      };
    }

    case "productsUpsert": {
      if (!message) return state;
      const dataPart: ChatMessagePart = {
        type: "data-products",
        id: operation.id,
        products: operation.products,
        stepId: currentStepId,
      };
      return {
        ...state,
        message: {
          ...message,
          parts: upsertPart(message.parts, dataPart),
        },
      };
    }

    case "sourceUpsert": {
      if (!message) return state;
      if (operation.sourceType === "source-url") {
        const sourcePart: ChatMessagePart = {
          type: "source-url",
          id: operation.id,
          title: operation.title,
          url: operation.url,
          stepId: currentStepId,
        };
        return {
          ...state,
          message: {
            ...message,
            parts: upsertPart(message.parts, sourcePart),
          },
        };
      }

      const sourcePart: ChatMessagePart = {
        type: "source-document",
        id: operation.id,
        title: operation.title,
        excerpt: operation.excerpt,
        url: operation.url,
        stepId: currentStepId,
      };
      return {
        ...state,
        message: {
          ...message,
          parts: upsertPart(message.parts, sourcePart),
        },
      };
    }

    default:
      return state;
  }
}

function upsertPart(
  parts: ChatMessagePart[],
  part: ChatMessagePart,
): ChatMessagePart[] {
  const index = parts.findIndex((p) => p.id === part.id);
  if (index >= 0) {
    const updated = [...parts];
    updated[index] = part;
    return updated;
  }
  return [...parts, part];
}

function upsertTransientPart(
  transientParts: TransientPart[],
  part: TransientPart,
): TransientPart[] {
  const index = transientParts.findIndex((p) => p.id === part.id);
  if (index >= 0) {
    const updated = [...transientParts];
    updated[index] = part;
    return updated;
  }
  return [...transientParts, part];
}