import { type Router } from "express";
import { z } from "zod";

import { getAgentAuthContext } from "../auth/agentAuth.js";
import { type CredentialRefreshCoordinator } from "../auth/credentialRefresh.js";
import { logger } from "../logger.js";
import { type OpencodeRuntimeAdapter } from "../runtime/opencode.js";
import {
  getRuntimeSessionContext,
  setRuntimeSessionContext,
} from "../runtime/sessionContext.js";
import { type SessionMetadataStore } from "../sessions/metadataStore.js";
import { type SessionUiStateStore } from "../sessions/uiStateStore.js";
import { toActorKey, toProjectKey } from "../utils/fileStore.js";
import {
  type ActiveRun,
  toPermissionStatus,
  updateLastAssistantPermission,
  updateLastAssistantQuestion,
} from "./chatUiState.js";

const permissionReplyPayloadSchema = z.object({
  request_id: z.string().min(1),
  reply: z.enum(["once", "always", "reject"]),
  message: z.string().max(1000).optional(),
});

const questionReplyPayloadSchema = z.object({
  request_id: z.string().min(1),
  action: z.enum(["reply", "reject"]).default("reply"),
  answers: z.array(z.array(z.string().max(2000))).default([]),
});

const credentialRefreshPayloadSchema = z.object({
  request_id: z.string().min(1).max(128),
});

type RegisterInteractionRoutesOptions = {
  activeRuns: Map<string, ActiveRun>;
  credentialRefreshCoordinator: CredentialRefreshCoordinator;
  runtime: OpencodeRuntimeAdapter;
  sessionMetadataStore: SessionMetadataStore;
  sessionUiStateStore: SessionUiStateStore;
};

const toSessionUiStateContext = (sessionId: string) => ({
  sessionId,
});

export const registerChatInteractionRoutes = (
  chatRouter: Router,
  {
    activeRuns,
    credentialRefreshCoordinator,
    runtime,
    sessionMetadataStore,
    sessionUiStateStore,
  }: RegisterInteractionRoutesOptions,
) => {
  chatRouter.post("/sessions/:session_id/credential-refreshes", async (req, res) => {
    const parsed = credentialRefreshPayloadSchema.safeParse(req.body);
    if (!parsed.success) {
      res.status(400).json({
        message: "invalid request payload",
        detail: parsed.error.flatten(),
      });
      return;
    }

    const authContext = getAgentAuthContext(req);
    const actorKey = toActorKey(authContext.userId);
    const projectKey = toProjectKey(authContext.projectId);
    const sessionRecord = await sessionMetadataStore.get(
      {
        actorKey,
        projectId: authContext.projectId,
        projectKey,
        userId: authContext.userId,
      },
      req.params.session_id,
    );
    if (!sessionRecord) {
      res.status(404).json({ message: "session not found" });
      return;
    }

    const current = getRuntimeSessionContext(sessionRecord.sessionId);
    if (!current || current.actorKey !== actorKey || current.projectKey !== projectKey) {
      res.status(409).json({ message: "runtime session context unavailable" });
      return;
    }
    if (
      credentialRefreshCoordinator.getPendingRequestId(sessionRecord.sessionId) !==
      parsed.data.request_id
    ) {
      res.status(409).json({ message: "credential refresh request is no longer pending" });
      return;
    }
    const refreshedContext = {
      ...current,
      accessToken: authContext.accessToken,
      authExpired: undefined,
      network: authContext.network,
      projectId: authContext.projectId,
      tokenExpiresAt: authContext.tokenExpiresAt,
      traceId: req.header("x-trace-id")?.trim() || current.traceId,
    };
    setRuntimeSessionContext(refreshedContext);
    credentialRefreshCoordinator.resolve(
      sessionRecord.sessionId,
      parsed.data.request_id,
      refreshedContext,
    );
    res.status(202).json({
      session_id: sessionRecord.sessionId,
      request_id: parsed.data.request_id,
      status: "accepted",
    });
  });

  chatRouter.post("/sessions/:session_id/permission-responses", async (req, res) => {
    const parsed = permissionReplyPayloadSchema.safeParse(req.body);
    if (!parsed.success) {
      res.status(400).json({
        message: "invalid request payload",
        detail: parsed.error.flatten(),
      });
      return;
    }

    try {
      const authContext = getAgentAuthContext(req);
      const projectId = authContext.projectId;
      const userId = authContext.userId;
      const actorKey = toActorKey(userId);
      const projectKey = toProjectKey(projectId);
      const requestId = parsed.data.request_id;
      const sessionRecord = await sessionMetadataStore.get(
        { actorKey, projectId, projectKey, userId },
        req.params.session_id,
      );
      if (!sessionRecord) {
        res.status(404).json({ message: "session not found" });
        return;
      }

      const run = activeRuns.get(sessionRecord.sessionId);
      if (!run || run.status !== "running") {
        res.status(409).json({ message: "session is not waiting for permissions" });
        return;
      }

      const pendingPermission = run.pendingPermissions.get(requestId);
      if (!pendingPermission) {
        res.status(404).json({ message: "permission request not found" });
        return;
      }
      const persistPermissionState = async () => {
        const currentState = await sessionUiStateStore.read(
          toSessionUiStateContext(sessionRecord.sessionId),
        );
        await sessionUiStateStore.write(toSessionUiStateContext(sessionRecord.sessionId), {
          sessionId: sessionRecord.sessionId,
          isTitleManuallyEdited: currentState?.isTitleManuallyEdited ?? false,
          messages: run.messages,
        });
      };

      try {
        await runtime.replyPermission({
          requestId,
          sessionId: sessionRecord.sessionId,
          reply: parsed.data.reply,
          message: parsed.data.message,
        });
      } catch (error) {
        run.messages = updateLastAssistantPermission(
          run.messages,
          requestId,
          (permission) => ({
            ...permission,
            status: "error",
            error:
              error instanceof Error
                ? error.message
                : "failed to reply permission",
          }),
        );
        await persistPermissionState().catch((persistError) => {
          logger.warn(
            { err: persistError, sessionId: sessionRecord.sessionId },
            "failed to persist permission error state",
          );
        });
        res.status(502).json({
          message: "permission reply failed",
          detail: error instanceof Error ? error.message : String(error),
        });
        return;
      }

      run.pendingPermissions.delete(requestId);
      const status = toPermissionStatus(parsed.data.reply);
      run.messages = updateLastAssistantPermission(
        run.messages,
        requestId,
        (permission) => ({
          ...permission,
          status,
          repliedAt: Date.now(),
        }),
      );
      await persistPermissionState().catch((persistError) => {
        logger.warn(
          { err: persistError, sessionId: sessionRecord.sessionId },
          "failed to persist permission reply state",
        );
      });
      for (const subscriber of run.subscribers) {
        subscriber.write("permission_response", {
          session_id: sessionRecord.sessionId,
          request_id: requestId,
          reply: parsed.data.reply,
        });
      }

      res.status(202).json({
        session_id: sessionRecord.sessionId,
        request_id: requestId,
        reply: parsed.data.reply,
      });
    } catch (error) {
      const detail = error instanceof Error ? error.message : String(error);
      logger.error({ err: error }, "permission reply route failed");
      res.status(500).json({
        message: "permission reply route failed",
        detail,
      });
    }
  });

  chatRouter.post("/sessions/:session_id/question-responses", async (req, res) => {
    const parsed = questionReplyPayloadSchema.safeParse(req.body);
    if (!parsed.success) {
      res.status(400).json({
        message: "invalid request payload",
        detail: parsed.error.flatten(),
      });
      return;
    }

    try {
      const authContext = getAgentAuthContext(req);
      const projectId = authContext.projectId;
      const userId = authContext.userId;
      const actorKey = toActorKey(userId);
      const projectKey = toProjectKey(projectId);
      const requestId = parsed.data.request_id;
      const sessionRecord = await sessionMetadataStore.get(
        { actorKey, projectId, projectKey, userId },
        req.params.session_id,
      );
      if (!sessionRecord) {
        res.status(404).json({ message: "session not found" });
        return;
      }

      const run = activeRuns.get(sessionRecord.sessionId);
      if (!run) {
        res.status(409).json({ message: "session is not waiting for questions" });
        return;
      }

      const pendingQuestion = run.pendingQuestions.get(requestId);
      if (!pendingQuestion) {
        res.status(404).json({ message: "question request not found" });
        return;
      }
      const persistQuestionState = async () => {
        const currentState = await sessionUiStateStore.read(
          toSessionUiStateContext(sessionRecord.sessionId),
        );
        await sessionUiStateStore.write(toSessionUiStateContext(sessionRecord.sessionId), {
          sessionId: sessionRecord.sessionId,
          isTitleManuallyEdited: currentState?.isTitleManuallyEdited ?? false,
          messages: run.messages,
        });
      };

      try {
        if (parsed.data.action === "reject") {
          await runtime.rejectQuestion({
            requestId,
            sessionId: sessionRecord.sessionId,
          });
        } else {
          await runtime.replyQuestion({
            requestId,
            sessionId: sessionRecord.sessionId,
            answers: parsed.data.answers,
          });
        }
      } catch (error) {
        run.messages = updateLastAssistantQuestion(
          run.messages,
          requestId,
          (question) => ({
            ...question,
            status: "error",
            error:
              error instanceof Error
                ? error.message
                : `failed to ${parsed.data.action} question`,
          }),
        );
        await persistQuestionState().catch((persistError) => {
          logger.warn(
            { err: persistError, sessionId: sessionRecord.sessionId },
            "failed to persist question error state",
          );
        });
        res.status(502).json({
          message: `question ${parsed.data.action} failed`,
          detail: error instanceof Error ? error.message : String(error),
        });
        return;
      }

      run.pendingQuestions.delete(requestId);
      run.messages = updateLastAssistantQuestion(
        run.messages,
        requestId,
        (question) => ({
          ...question,
          status: parsed.data.action === "reject" ? "rejected" : "answered",
          answers:
            parsed.data.action === "reject" ? question.answers : parsed.data.answers,
          repliedAt: Date.now(),
          error: undefined,
        }),
      );
      await persistQuestionState().catch((persistError) => {
        logger.warn(
          { err: persistError, sessionId: sessionRecord.sessionId },
          "failed to persist question reply state",
        );
      });
      for (const subscriber of run.subscribers) {
        subscriber.write("question_response", {
          session_id: pendingQuestion.session_id,
          request_id: requestId,
          ...(parsed.data.action === "reject"
            ? { rejected: true }
            : { answers: parsed.data.answers }),
        });
      }
      if (
        run.status !== "running" &&
        run.pendingPermissions.size === 0 &&
        run.pendingQuestions.size === 0
      ) {
        activeRuns.delete(sessionRecord.sessionId);
      }

      res.status(202).json({
        session_id: pendingQuestion.session_id,
        request_id: requestId,
        ...(parsed.data.action === "reject"
          ? { rejected: true }
          : { answers: parsed.data.answers }),
      });
    } catch (error) {
      const detail = error instanceof Error ? error.message : String(error);
      logger.error({ err: error }, "question response route failed");
      res.status(500).json({
        message: "question response route failed",
        detail,
      });
    }
  });
};