import type { UIMessage } from "ai";
import type { Dispatch, SetStateAction } from "react";
import {
  parseAnswerUiBlock,
  type AnswerUiPart,
  applyPermissionResponse,
  agentTaskContextReferenceSchema,
  applyQuestionResponse,
  toTodoUpdate,
  upsertPermission,
  upsertProgress,
  upsertQuestion,
  type AgentChatMessage,
  type AgentPermissionReply,
  type AgentPermissionStatus,
  type AgentQuestionRequest,
  type AgentSessionStreamEvent,
  type AgentStreamRenderState
} from "@/features/agent";

type AgentUiDataParts = {
  session_created: Record<string, unknown>;
  agent_message: Record<string, unknown>;
  task_context: Record<string, unknown>;
  execution_plan: Record<string, unknown>;
  progress: Record<string, unknown>;
  final_answer: Record<string, unknown>;
  credential_refresh_required: Record<string, unknown>;
  credential_refreshed: Record<string, unknown>;
  credential_refresh_failed: Record<string, unknown>;
  todo_update: Record<string, unknown>;
  permission_request: Record<string, unknown>;
  permission_response: Record<string, unknown>;
  question_request: Record<string, unknown>;
  question_response: Record<string, unknown>;
  stream_token: Record<string, unknown>;
  tool_call: Record<string, unknown>;
  frontend_action: Record<string, unknown>;
  frontend_action_result: Record<string, unknown>;
  answer_ui: Record<string, unknown>;
  ui_envelope: Record<string, unknown>;
  session_title: Record<string, unknown>;
  user_input: Record<string, unknown>;
};

export type AgentUiMessage = UIMessage<unknown, AgentUiDataParts>;

export type AgentDataPart = Extract<AgentUiMessage["parts"][number], { type: `data-${string}` }>;

export type PermissionOverride = {
  status: AgentPermissionStatus;
  error?: string;
};

export type QuestionOverride = {
  status: AgentQuestionRequest["status"];
  answers?: string[][];
  error?: string;
};

type StreamRenderStateSetter = Dispatch<SetStateAction<AgentStreamRenderState>>;

export function markStreamRenderPending(
  data: Record<string, unknown>,
  messages: AgentUiMessage[],
  setStreamRenderState: StreamRenderStateSetter
) {
  const text = getDataString(data, "content");
  if (!text) {
    return;
  }

  const messageId = getDataString(data, "message_id") ?? getLastAssistantUiMessageId(messages);
  if (!messageId) {
    return;
  }

  setStreamRenderState((current) => {
    if (current[messageId]?.done === false) {
      return current;
    }

    return {
      ...current,
      [messageId]: { done: false }
    };
  });
}

export function createCompletedStreamRenderState(messages: AgentUiMessage[]): AgentStreamRenderState {
  return messages.reduce<AgentStreamRenderState>((next, message) => {
    if (message.role === "assistant") {
      next[message.id] = { done: true };
    }
    return next;
  }, {});
}

export function markLastAssistantStreamDone(
  messages: AgentUiMessage[],
  setStreamRenderState: StreamRenderStateSetter
) {
  const messageId = getLastAssistantUiMessageId(messages);
  setStreamRenderState((current) => {
    if (!messageId) {
      return Object.fromEntries(
        Object.entries(current).map(([id, state]) => [id, { ...state, done: true }])
      );
    }

    return {
      ...current,
      [messageId]: { done: true }
    };
  });
}

export function toAgentChatMessages(
  messages: AgentUiMessage[],
  permissionOverrides: Record<string, PermissionOverride>,
  questionOverrides: Record<string, QuestionOverride>
): AgentChatMessage[] {
  return messages.flatMap((message) => {
    if (message.role !== "user" && message.role !== "assistant") {
      return [];
    }

    const answerParts = collectAnswerParts(message);
    let next: AgentChatMessage = {
      id: message.id,
      role: message.role,
      content: answerParts ? answerParts.map((part) => part.type === "text" ? part.text : part.block.fallback_text).join("\n\n") : collectMessageText(message),
      ...(answerParts ? { answerParts } : {})
    };

    for (const part of message.parts) {
      if (!isAgentDataPart(part)) {
        continue;
      }
      if (part.type === "data-task_context") {
        const parsed = agentTaskContextReferenceSchema.safeParse(part.data);
        if (parsed.success) next = { ...next, taskContext: parsed.data };
      } else if (part.type === "data-final_answer") {
        const content = getDataString(part.data, "content");
        if (content?.trim()) {
          const refs = part.data?.ui_block_ids;
          const uiBlockIds = Array.isArray(refs) ? [...new Set(refs.filter((id): id is string => typeof id === "string"))].slice(0, 48) : [];
          next = { ...next, finalAnswer: { content, uiBlockIds } };
        }
      } else if (part.type === "data-progress") {
        next = {
          ...next,
          progress: upsertProgress(next.progress, part.data)
        };
      } else if (part.type === "data-todo_update") {
        const todoUpdate = toTodoUpdate(part.data);
        if (todoUpdate) {
          next = {
            ...next,
            todos: todoUpdate
          };
        }
      } else if (part.type === "data-permission_request") {
        next = {
          ...next,
          permissions: upsertPermission(next.permissions, part.data)
        };
      } else if (part.type === "data-permission_response") {
        next = {
          ...next,
          permissions: applyPermissionResponse(next.permissions, part.data)
        };
      } else if (part.type === "data-question_request") {
        next = {
          ...next,
          questions: upsertQuestion(next.questions, part.data)
        };
      } else if (part.type === "data-question_response") {
        next = {
          ...next,
          questions: applyQuestionResponse(next.questions, part.data)
        };
      }
    }

    if (next.finalAnswer) {
      const selected = new Set(next.finalAnswer.uiBlockIds);
      next.content = [next.finalAnswer.content, ...(next.answerParts ?? []).flatMap((part) => part.type === "answer-ui" && selected.has(part.block.block_id) ? [part.block.fallback_text] : [])].join("\n\n");
    }
    return [applyQuestionOverrides(applyPermissionOverrides(next, permissionOverrides), questionOverrides)];
  });
}

export function toAgentUiMessages(messages: unknown[]): AgentUiMessage[] {
  return messages.flatMap((message) => {
    if (isAgentUiMessage(message)) {
      return [message];
    }
    const legacyMessage = toAgentUiMessageFromLegacy(message);
    return legacyMessage ? [legacyMessage] : [];
  });
}

export function applySessionStreamEvent(messages: AgentUiMessage[], event: AgentSessionStreamEvent): AgentUiMessage[] {
  if (event.type === "state") {
    return toAgentUiMessages(event.messages);
  }

  if (event.type === "user_input") {
    const id = getDataString(event.data, "id");
    const text = getDataString(event.data, "text");
    if (event.data.status !== "delivered" || !id || !text || messages.some((message) => message.id === id)) return messages;
    const input: AgentUiMessage = { id, role: "user", parts: [{ type: "text", text }] };
    const tail = messages.at(-1);
    return tail?.role === "assistant" ? [...messages.slice(0, -1), input, tail] : [...messages, input];
  }

  if (event.type === "token") {
    const token = getDataString(event.data, "content");
    return token ? updateLastAssistantUiMessage(messages, (message) => appendTextToUiMessage(message, token)) : messages;
  }

  if (
    event.type === "final_answer" ||
    event.type === "progress" ||
    event.type === "execution_plan" ||
    event.type === "todo_update" ||
    event.type === "permission_request" ||
    event.type === "permission_response" ||
    event.type === "question_request" ||
    event.type === "question_response" ||
    event.type === "answer_ui" ||
    event.type === "ui_envelope" ||
    event.type === "session_title"
  ) {
    return updateLastAssistantUiMessage(messages, (message) =>
      upsertUiDataPart(message, event.type, event.data)
    );
  }

  if (event.type === "error") {
    const message = getDataString(event.data, "message") ?? "Agent stream failed";
    return updateLastAssistantUiMessage(messages, (item) =>
      appendTextToUiMessage(item, item.parts.some((part) => part.type === "text") ? `\n\n错误：${message}` : `错误：${message}`)
    );
  }

  return messages;
}

export function toPermissionStatus(reply: AgentPermissionReply): AgentPermissionStatus {
  if (reply === "always") {
    return "approved_always";
  }
  if (reply === "once") {
    return "approved_once";
  }
  return "rejected";
}

export function getDataString(data: unknown, key: string) {
  if (typeof data !== "object" || data === null) {
    return undefined;
  }
  const value = (data as Record<string, unknown>)[key];
  return typeof value === "string" ? value : undefined;
}

export function readBodyString(body: unknown, key: string) {
  if (typeof body !== "object" || body === null) {
    return undefined;
  }
  const value = (body as Record<string, unknown>)[key];
  return typeof value === "string" ? value : undefined;
}

function getLastAssistantUiMessageId(messages: AgentUiMessage[]) {
  for (let index = messages.length - 1; index >= 0; index -= 1) {
    if (messages[index].role === "assistant") {
      return messages[index].id;
    }
  }
  return undefined;
}

function isAgentUiMessage(value: unknown): value is AgentUiMessage {
  if (typeof value !== "object" || value === null) {
    return false;
  }
  const message = value as Record<string, unknown>;
  return (
    typeof message.id === "string" &&
    (message.role === "user" || message.role === "assistant") &&
    Array.isArray(message.parts)
  );
}

function toAgentUiMessageFromLegacy(value: unknown): AgentUiMessage | null {
  if (typeof value !== "object" || value === null) {
    return null;
  }

  const message = value as Record<string, unknown>;
  if (
    typeof message.id !== "string" ||
    (message.role !== "user" && message.role !== "assistant")
  ) {
    return null;
  }

  const parts: AgentUiMessage["parts"] = [];
  if (typeof message.content === "string" && message.content) {
    parts.push({ type: "text", text: message.content });
  }
  appendLegacyDataParts(parts, "progress", message.progress);
  appendLegacyDataParts(parts, "permission_request", message.permissions);
  appendLegacyDataParts(parts, "question_request", message.questions);
  appendLegacyDataParts(parts, "todo_update", message.todos);
  appendLegacyDataParts(parts, "execution_plan", message.executionPlans);

  return {
    id: message.id,
    role: message.role,
    parts
  } as AgentUiMessage;
}

function appendLegacyDataParts(
  parts: AgentUiMessage["parts"],
  eventType: string,
  value: unknown
) {
  if (!value) {
    return;
  }

  const values = Array.isArray(value) ? value : [value];
  values.forEach((item, index) => {
    if (typeof item !== "object" || item === null) {
      return;
    }
    const data = normalizeLegacyDataPart(eventType, item as Record<string, unknown>);
    const id =
      getDataString(data, "id") ??
      getDataString(data, "request_id") ??
      `${eventType}-${index}`;
    parts.push({
      type: `data-${eventType}`,
      id,
      data
    } as AgentDataPart);
  });
}

function normalizeLegacyDataPart(eventType: string, value: Record<string, unknown>) {
  if (eventType === "progress") {
    return {
      ...value,
      started_at: value.started_at ?? value.startedAt,
      ended_at: value.ended_at ?? value.endedAt,
      elapsed_ms: value.elapsed_ms ?? value.elapsedMs,
      duration_ms: value.duration_ms ?? value.durationMs
    };
  }

  if (eventType === "permission_request" || eventType === "question_request") {
    return {
      ...value,
      session_id: value.session_id ?? value.sessionId,
      request_id: value.request_id ?? value.requestId,
      created_at: value.created_at ?? value.createdAt
    };
  }

  if (eventType === "todo_update") {
    return {
      ...value,
      session_id: value.session_id ?? value.sessionId,
      message_id: value.message_id ?? value.messageId,
      created_at: value.created_at ?? value.createdAt
    };
  }

  return value;
}

function updateLastAssistantUiMessage(
  messages: AgentUiMessage[],
  updater: (message: AgentUiMessage) => AgentUiMessage
) {
  for (let index = messages.length - 1; index >= 0; index -= 1) {
    if (messages[index].role === "assistant") {
      const next = [...messages];
      next[index] = updater(messages[index]);
      return next;
    }
  }

  return [...messages, updater(createAssistantUiMessage())];
}

function createAssistantUiMessage(): AgentUiMessage {
  return {
    id: `assistant-${Date.now().toString(36)}`,
    role: "assistant",
    parts: []
  } as AgentUiMessage;
}

function appendTextToUiMessage(message: AgentUiMessage, text: string): AgentUiMessage {
  let textIndex = -1;
  for (let index = message.parts.length - 1; index >= 0; index -= 1) {
    if (message.parts[index].type === "text") { textIndex = index; break; }
  }
  if (message.parts.slice(textIndex + 1).some((part) => part.type === "data-answer_ui")) textIndex = -1;
  if (textIndex === -1) {
    return {
      ...message,
      parts: [...message.parts, { type: "text", text }]
    };
  }

  return {
    ...message,
    parts: message.parts.map((part, index) =>
      index === textIndex && part.type === "text"
        ? { ...part, text: `${part.text}${text}` }
        : part
    )
  };
}

function upsertUiDataPart(
  message: AgentUiMessage,
  eventType: string,
  data: Record<string, unknown>
): AgentUiMessage {
  const type = `data-${eventType}`;
  const id =
    getDataString(data, "id") ??
    getDataString(data, "request_id") ??
    getDataString(data, "envelope_id") ??
    getDataString(data, "block_id");
  const existingIndex =
    id === undefined
      ? -1
      : message.parts.findIndex((part) => part.type === type && "id" in part && part.id === id);
  const part = (id ? { type, id, data } : { type, data }) as AgentDataPart;

  if (existingIndex === -1) {
    return {
      ...message,
      parts: [...message.parts, part]
    };
  }

  return {
    ...message,
    parts: message.parts.map((item, index) => (index === existingIndex ? part : item))
  };
}

function applyPermissionOverrides(
  message: AgentChatMessage,
  permissionOverrides: Record<string, PermissionOverride>
) {
  if (!message.permissions?.length) {
    return message;
  }

  return {
    ...message,
    permissions: message.permissions.map((permission) => {
      const override = permissionOverrides[permission.requestId];
      return override
        ? {
            ...permission,
            status: override.status,
            error: override.error,
            repliedAt: override.status === "submitting" || override.status === "error" ? permission.repliedAt : Date.now()
          }
        : permission;
    })
  };
}

function applyQuestionOverrides(
  message: AgentChatMessage,
  questionOverrides: Record<string, QuestionOverride>
) {
  if (!message.questions?.length) {
    return message;
  }

  return {
    ...message,
    questions: message.questions.map((question) => {
      const override = questionOverrides[question.requestId];
      return override
        ? {
            ...question,
            status: override.status,
            answers: override.answers ?? question.answers,
            error: override.error,
            repliedAt: override.status === "submitting" || override.status === "error" ? question.repliedAt : Date.now()
          }
        : question;
    })
  };
}

function collectMessageText(message: AgentUiMessage) {
  return message.parts
    .flatMap((part) => (part.type === "text" ? [part.text] : []))
    .join("");
}

function isAgentDataPart(part: AgentUiMessage["parts"][number]): part is AgentDataPart {
  return typeof part.type === "string" && part.type.startsWith("data-") && "data" in part;
}

function collectAnswerParts(message: AgentUiMessage): AnswerUiPart[] | undefined {
  if (!message.parts.some((part) => part.type === "data-answer_ui" || part.type === "data-final_answer")) return undefined;
  const result: AnswerUiPart[] = [];
  for (const part of message.parts) {
    if (part.type === "text" && part.text) result.push({ type: "text", text: part.text });
    if (part.type === "data-answer_ui") {
      const data = typeof part.data === "object" && part.data !== null ? part.data : {};
      const { version, block_id, fallback_text, spec } = data;
      const block = parseAnswerUiBlock({ version, block_id, fallback_text, spec });
      if (block) result.push({ type: "answer-ui", block });
      else result.push({ type: "text", text: typeof fallback_text === "string" ? fallback_text.slice(0, 4000) : "这段结构化内容无法显示。" });
    }

  }
  return result;
}