"use client";

import { useChat } from "@ai-sdk/react";
import { type ChatStatus, DefaultChatTransport, type UIMessage } from "ai";
import { CheckIcon, PenLineIcon } from "lucide-react";
import { useCallback, useEffect, useMemo, useRef, useState } from "react";
import { toast } from "sonner";
import {
  Levee,
  LeveeConversation,
  LeveeRail,
  LeveeStageChat,
  LeveeStageCollectedData,
  LeveeStageFormView,
  LeveeStageHeader,
  LeveeStageLockChecklist,
  LeveeStageView,
  useLevee,
} from "@/components/ai-elements/levee";
import { Loader } from "@/components/ai-elements/loader";
import {
  PromptInput,
  PromptInputFooter,
  PromptInputSubmit,
  PromptInputTextarea,
  PromptInputTools,
} from "@/components/ai-elements/prompt-input";
import { Button } from "@/components/ui/button";
import { ScrollArea } from "@/components/ui/scroll-area";
import { userResearchStages } from "../lib/stages";
import type {
  CompleteStageOutput,
  EvaluateLockOutput,
  LeveeContextData,
  LeveeStagePhase,
  LeveeUIMessage,
  SynthesisResult,
} from "../lib/types";
import { StageChat } from "./stage-chat";

const LEVEE_API = "/view/levee-user-research/api/levee";

// ---------------------------------------------------------------------------
// useLeveeStage -- centralised hook for all stage modes
// ---------------------------------------------------------------------------

type UseLeveeStageOptions = {
  stageConfig: (typeof userResearchStages)[number];
  priorData: LeveeContextData;
  formData: Record<string, unknown>;
  initialMessages: LeveeUIMessage[];
  onLockUpdate: (payload: EvaluateLockOutput) => void;
  onComplete: (payload: CompleteStageOutput) => void;
  onMessagesChange: (messages: LeveeUIMessage[]) => void;
  onPhaseChange: (phase: LeveeStagePhase) => void;
};

function useLeveeStage({
  stageConfig,
  priorData,
  formData,
  initialMessages,
  onLockUpdate,
  onComplete,
  onMessagesChange,
  onPhaseChange,
}: UseLeveeStageOptions) {
  const [input, setInput] = useState("");
  const [synthesisResult, setSynthesisResult] =
    useState<SynthesisResult | null>(null);
  const autoStartedRef = useRef(false);

  const refs = useRef({
    onLockUpdate,
    onComplete,
    onMessagesChange,
    onPhaseChange,
  });
  refs.current = { onLockUpdate, onComplete, onMessagesChange, onPhaseChange };

  const chatResult = useChat<LeveeUIMessage>({
    id: `levee-stage-${stageConfig.id}`,
    messages: initialMessages,
    transport: new DefaultChatTransport({
      api: LEVEE_API,
      credentials: "include",
    }),
    onData: (data) => {
      if (data.type === "data-lock-update") {
        refs.current.onLockUpdate(data.data as EvaluateLockOutput);
      }
      if (data.type === "data-stage-complete") {
        refs.current.onComplete(data.data as CompleteStageOutput);
      }
      if (data.type === "data-synthesis-result") {
        setSynthesisResult(data.data as SynthesisResult);
      }
    },
    onError: (err) => {
      toast.error(err.message);
    },
  });

  const { messages, sendMessage, status, addToolApprovalResponse, error } =
    chatResult;

  useEffect(() => {
    refs.current.onMessagesChange(messages as LeveeUIMessage[]);
    if (messages.length > 0) {
      refs.current.onPhaseChange("conversation");
    }
  }, [messages]);

  const buildRequestBody = useCallback(
    () => ({
      stageId: stageConfig.id,
      stageConfig,
      priorData,
      formData,
    }),
    [stageConfig, priorData, formData]
  );

  const sendStageMessage = useCallback(
    (message: { text: string }) =>
      sendMessage(message, { body: buildRequestBody() }),
    [sendMessage, buildRequestBody]
  );

  const sendStageMessageRef = useRef(sendStageMessage);
  sendStageMessageRef.current = sendStageMessage;

  useEffect(() => {
    if (autoStartedRef.current) return;
    if (initialMessages.length > 0) return;
    const mode = stageConfig.mode ?? "chat";
    if (mode === "form") return;
    autoStartedRef.current = true;
    refs.current.onPhaseChange("conversation");
    sendStageMessageRef.current({ text: "Begin this stage." });
  }, [initialMessages.length, stageConfig.mode]);

  return {
    messages,
    input,
    setInput,
    sendStageMessage,
    status,
    error,
    synthesisResult,
    addToolApprovalResponse,
  };
}

// ---------------------------------------------------------------------------
// Root component
// ---------------------------------------------------------------------------

export function LeveeUserResearch() {
  const [activeStageIndex, setActiveStageIndex] = useState(0);

  return (
    <Levee
      activeStageIndex={activeStageIndex}
      allowRevisit={true}
      config={userResearchStages}
      onActiveStageIndexChange={setActiveStageIndex}
      onComplete={() => {
        toast.success("Levee complete. Your research plan is ready.");
      }}
    >
      <LeveeSidebar />
      <LeveeWorkspace />
    </Levee>
  );
}

function LeveeSidebar() {
  const { stages, activeStageIndex, contextPanelOpen } = useLevee();
  const activeStageState = stages[activeStageIndex];

  return (
    <LeveeRail compact={!contextPanelOpen}>
      {activeStageState && (
        <>
          <LeveeStageLockChecklist
            items={activeStageState.lockItems}
            percentage={activeStageState.lockPercentage}
          />
          <LeveeStageCollectedData data={activeStageState.collectedData} />
        </>
      )}
    </LeveeRail>
  );
}

// ---------------------------------------------------------------------------
// Workspace -- routes to the correct view per stage mode
// ---------------------------------------------------------------------------

function LeveeWorkspace() {
  const {
    config,
    stages,
    activeStageIndex,
    setStagePhase,
    updateCollectedData,
    updateLockItems,
    updateFormData,
    advanceStage,
  } = useLevee();

  const activeStage = config[activeStageIndex];
  const activeStageState = stages[activeStageIndex];

  const priorData = useMemo<LeveeContextData>(
    () =>
      stages.reduce<LeveeContextData>((acc, stageState, index) => {
        if (index >= activeStageIndex) return acc;
        acc[stageState.id] = stageState.collectedData;
        return acc;
      }, {}),
    [stages, activeStageIndex]
  );

  const [messagesByStage, setMessagesByStage] = useState<
    Record<string, UIMessage[]>
  >({});

  const handleMessagesChange = useCallback(
    (messages: LeveeUIMessage[]) => {
      if (!activeStage) return;
      setMessagesByStage((prev) => ({
        ...prev,
        [activeStage.id]: messages,
      }));
    },
    [activeStage]
  );

  const handlePhaseChange = useCallback(
    (phase: LeveeStagePhase) => {
      setStagePhase(activeStageIndex, phase);
    },
    [activeStageIndex, setStagePhase]
  );

  const handleStageComplete = useCallback(
    (payload: CompleteStageOutput) => {
      updateCollectedData(activeStageIndex, payload.data);
      updateLockItems(activeStageIndex, payload.lockItems, payload.summary);
      advanceStage(payload.summary, payload.data);
    },
    [activeStageIndex, updateCollectedData, updateLockItems, advanceStage]
  );

  const handleStageLockUpdate = useCallback(
    (payload: EvaluateLockOutput) => {
      updateCollectedData(activeStageIndex, payload.collectedData);
      updateLockItems(activeStageIndex, payload.lockItems, payload.summary);
    },
    [activeStageIndex, updateCollectedData, updateLockItems]
  );

  if (!(activeStage && activeStageState)) {
    return null;
  }

  const stageMode = activeStage.mode ?? "chat";
  const initialMessages =
    (messagesByStage[activeStage.id] as LeveeUIMessage[] | undefined) ?? [];

  const sharedProps = {
    stageConfig: activeStage,
    priorData,
    formData: activeStageState.formData,
    initialMessages,
    onMessagesChange: handleMessagesChange,
    onPhaseChange: handlePhaseChange,
    onComplete: handleStageComplete,
    onLockUpdate: handleStageLockUpdate,
  } satisfies UseLeveeStageOptions;

  if (stageMode === "form") {
    return (
      <LeveeStageView>
        <LeveeStageHeader />
        <FormStageView
          {...sharedProps}
          key={activeStage.id}
          onFieldChange={(fieldId, value) => {
            updateFormData(activeStageIndex, { [fieldId]: value });
          }}
        />
      </LeveeStageView>
    );
  }

  if (stageMode === "synthesis") {
    return (
      <LeveeStageView>
        <LeveeStageHeader />
        <SynthesisStageView {...sharedProps} key={activeStage.id} />
      </LeveeStageView>
    );
  }

  if (stageMode === "qa") {
    return (
      <LeveeStageView>
        <LeveeStageHeader />
        <LeveeStageChat className="h-[calc(100%-57px)]">
          <QAStageView {...sharedProps} key={activeStage.id} />
        </LeveeStageChat>
      </LeveeStageView>
    );
  }

  return (
    <LeveeStageView>
      <LeveeStageHeader />
      <LeveeStageChat className="h-[calc(100%-57px)]">
        <ChatStageView {...sharedProps} key={activeStage.id} />
      </LeveeStageChat>
    </LeveeStageView>
  );
}

// ---------------------------------------------------------------------------
// FormStageView
// ---------------------------------------------------------------------------

type FormStageViewProps = UseLeveeStageOptions & {
  onFieldChange: (fieldId: string, value: unknown) => void;
};

function FormStageView({ onFieldChange, ...props }: FormStageViewProps) {
  const { messages, sendStageMessage, status } = useLeveeStage(props);

  const [confirmationText, setConfirmationText] = useState("");
  const isSubmitting = status === "streaming" || status === "submitted";

  useEffect(() => {
    if (messages.length === 0) return;
    const lastMessage = messages.at(-1);
    if (lastMessage?.role !== "assistant") return;
    const textParts = lastMessage.parts.filter((p) => p.type === "text");
    const text = textParts.map((p) => p.text).join("");
    if (text) {
      setConfirmationText(text);
    }
  }, [messages]);

  const handleFormSubmit = () => {
    props.onPhaseChange("submitting");

    const fieldSummary = props.stageConfig.initialInputs
      .map((field) => {
        const val = props.formData[field.id];
        const display = Array.isArray(val) ? val.join(", ") : String(val ?? "");
        return `${field.label}: ${display}`;
      })
      .filter((line) => !line.endsWith(": "))
      .join("\n");

    sendStageMessage({
      text: `Here is my information for the ${props.stageConfig.title} stage:\n\n${fieldSummary}\n\nPlease confirm what you received, evaluate the lock items, and complete this stage.`,
    });
  };

  return (
    <LeveeStageFormView
      confirmationText={confirmationText}
      fields={props.stageConfig.initialInputs}
      isSubmitting={isSubmitting}
      onFieldChange={onFieldChange}
      onSubmit={handleFormSubmit}
      values={props.formData}
    />
  );
}

// ---------------------------------------------------------------------------
// ChatStageView
// ---------------------------------------------------------------------------

function ChatStageView(props: UseLeveeStageOptions) {
  const {
    messages,
    input,
    setInput,
    sendStageMessage,
    status,
    error,
    addToolApprovalResponse,
  } = useLeveeStage(props);

  const handleSubmit = (
    _message: { text?: string; files?: unknown[] },
    e: React.FormEvent
  ) => {
    e.preventDefault();
    if (!input.trim()) return;
    sendStageMessage({ text: input });
    setInput("");
  };

  return (
    <>
      <LeveeConversation className="h-full pb-2">
        {messages.length === 0 && !error && (
          <div className="flex h-full items-center justify-center">
            <Loader />
          </div>
        )}
        <StageChat
          addToolApprovalResponse={addToolApprovalResponse}
          messages={messages}
          sendMessage={sendStageMessage}
          status={status as ChatStatus}
        />
      </LeveeConversation>

      <div className="hidden">
        <PromptInput onSubmit={handleSubmit} />
      </div>
    </>
  );
}

// ---------------------------------------------------------------------------
// QAStageView
// ---------------------------------------------------------------------------

function QAStageView(props: UseLeveeStageOptions) {
  const { messages, sendStageMessage, status, error, addToolApprovalResponse } =
    useLeveeStage(props);

  return (
    <LeveeConversation className="h-full pb-2">
      {messages.length === 0 && !error && (
        <div className="flex h-full items-center justify-center">
          <Loader />
        </div>
      )}
      <StageChat
        addToolApprovalResponse={addToolApprovalResponse}
        messages={messages}
        mode="qa"
        sendMessage={sendStageMessage}
        status={status as ChatStatus}
      />
    </LeveeConversation>
  );
}

// ---------------------------------------------------------------------------
// SynthesisStageView
// ---------------------------------------------------------------------------

function SynthesisStageView(props: UseLeveeStageOptions) {
  const {
    messages,
    input,
    setInput,
    sendStageMessage,
    status,
    error,
    synthesisResult,
    addToolApprovalResponse,
  } = useLeveeStage(props);

  const [showFeedback, setShowFeedback] = useState(false);

  if (synthesisResult) {
    return <LeveeSynthesisDeliverable result={synthesisResult} />;
  }

  const isReady = status === "ready" && messages.length > 0;

  const handleApprove = () => {
    sendStageMessage({ text: "Approved" });
  };

  const handleSubmitFeedback = (
    _message: { text?: string; files?: unknown[] },
    e: React.FormEvent
  ) => {
    e.preventDefault();
    if (!input.trim()) return;
    sendStageMessage({ text: input });
    setInput("");
    setShowFeedback(false);
  };

  return (
    <LeveeStageChat className="h-[calc(100%-57px)]">
      <LeveeConversation className="h-full pb-2">
        {messages.length === 0 && !error && (
          <div className="flex h-full items-center justify-center">
            <Loader />
          </div>
        )}
        <StageChat
          addToolApprovalResponse={addToolApprovalResponse}
          messages={messages}
          sendMessage={sendStageMessage}
          status={status as ChatStatus}
        />
      </LeveeConversation>

      <div className="border-t p-3">
        {isReady && !showFeedback ? (
          <div className="flex items-center gap-2">
            <Button className="gap-1.5" onClick={handleApprove} type="button">
              <CheckIcon className="size-4" />
              Approve Plan
            </Button>
            <Button
              className="gap-1.5"
              onClick={() => setShowFeedback(true)}
              type="button"
              variant="outline"
            >
              <PenLineIcon className="size-3.5" />
              Request Changes
            </Button>
          </div>
        ) : (
          <PromptInput onSubmit={handleSubmitFeedback}>
            <PromptInputTextarea
              onChange={(e) => setInput(e.target.value)}
              placeholder={
                showFeedback
                  ? "Describe the changes you'd like..."
                  : "Waiting for synthesis..."
              }
              value={input}
            />
            <PromptInputFooter>
              {showFeedback && (
                <Button
                  onClick={() => {
                    setShowFeedback(false);
                    setInput("");
                  }}
                  size="sm"
                  type="button"
                  variant="ghost"
                >
                  Cancel
                </Button>
              )}
              <PromptInputTools className="ml-auto">
                <PromptInputSubmit disabled={!input.trim()} status={status} />
              </PromptInputTools>
            </PromptInputFooter>
          </PromptInput>
        )}
      </div>
    </LeveeStageChat>
  );
}

// ---------------------------------------------------------------------------
// LeveeSynthesisDeliverable
// ---------------------------------------------------------------------------

function LeveeSynthesisDeliverable({ result }: { result: SynthesisResult }) {
  const isStreaming = result.status === "streaming";

  if (isStreaming && !result.fullPlan) {
    return (
      <div className="flex flex-1 items-center justify-center">
        <Loader />
      </div>
    );
  }

  if (isStreaming) {
    return (
      <ScrollArea className="h-[calc(100vh-57px)] flex-1">
        <div className="mx-auto max-w-2xl px-6 py-8">
          <div className="mb-6 flex items-center gap-2">
            <Loader />
            <span className="text-muted-foreground text-sm">
              Generating your research plan...
            </span>
          </div>
          <article className="prose prose-sm dark:prose-invert max-w-none whitespace-pre-wrap text-pretty">
            {result.fullPlan}
          </article>
        </div>
      </ScrollArea>
    );
  }

  return (
    <ScrollArea className="h-[calc(100vh-57px)] flex-1">
      <div className="mx-auto max-w-2xl px-6 py-8">
        <h2 className="mb-6 text-balance font-semibold text-lg">
          User Research Plan
        </h2>

        {result.researchObjective && (
          <section className="mb-6">
            <h3 className="mb-2 font-medium text-muted-foreground text-sm uppercase tracking-wide">
              Research Objective
            </h3>
            <p className="text-pretty leading-relaxed">
              {result.researchObjective}
            </p>
          </section>
        )}

        {result.participantProfile && (
          <section className="mb-6">
            <h3 className="mb-2 font-medium text-muted-foreground text-sm uppercase tracking-wide">
              Participant Profile
            </h3>
            <p className="text-pretty text-sm leading-relaxed">
              {result.participantProfile}
            </p>
          </section>
        )}

        {result.interviewGuide.length > 0 && (
          <section className="mb-6">
            <h3 className="mb-2 font-medium text-muted-foreground text-sm uppercase tracking-wide">
              Interview Guide
            </h3>
            <ul className="space-y-1.5">
              {result.interviewGuide.map((question) => (
                <li
                  className="flex gap-2 text-sm leading-relaxed"
                  key={question}
                >
                  <span className="mt-1.5 size-1.5 shrink-0 rounded-full bg-primary" />
                  <span className="text-pretty">{question}</span>
                </li>
              ))}
            </ul>
          </section>
        )}

        {result.analysisFramework.length > 0 && (
          <section className="mb-6">
            <h3 className="mb-2 font-medium text-muted-foreground text-sm uppercase tracking-wide">
              Analysis Framework
            </h3>
            <ul className="space-y-1.5">
              {result.analysisFramework.map((approach) => (
                <li
                  className="flex gap-2 text-sm leading-relaxed"
                  key={approach}
                >
                  <span className="mt-1.5 size-1.5 shrink-0 rounded-full bg-primary" />
                  <span className="text-pretty">{approach}</span>
                </li>
              ))}
            </ul>
          </section>
        )}

        {result.timeline.length > 0 && (
          <section className="mb-6">
            <h3 className="mb-2 font-medium text-muted-foreground text-sm uppercase tracking-wide">
              Timeline and Logistics
            </h3>
            <ol className="space-y-1.5">
              {result.timeline.map((item, index) => (
                <li className="flex gap-2 text-sm leading-relaxed" key={item}>
                  <span className="shrink-0 font-medium text-muted-foreground tabular-nums">
                    {index + 1}.
                  </span>
                  <span className="text-pretty">{item}</span>
                </li>
              ))}
            </ol>
          </section>
        )}
      </div>
    </ScrollArea>
  );
}