import type { ExtensionAPI } from "@mariozechner/pi-coding-agent";
import { Editor, type EditorTheme, Key, matchesKey, Text, truncateToWidth } from "@mariozechner/pi-tui";
import { Type } from "@sinclair/typebox";
import {
  allQuestionsAnswered,
  createAnsweredResult,
  createCancelledResult,
  createCustomAnswer,
  createPredefinedAnswer,
  isSomethingElseOption,
  nextTabAfterAnswer,
  normalizeQuestions,
  summarizeAnswers,
  wrapPrefixedText,
} from "./question-core.mjs";

interface QuestionOption {
  value: string;
  label: string;
  description?: string;
}

interface Question {
  id: string;
  label: string;
  prompt: string;
  options: QuestionOption[];
}

interface Answer {
  id: string;
  value: string;
  label: string;
  wasCustom: boolean;
  index?: number;
}

interface QuestionResult {
  questions: Question[];
  answers: Answer[];
  cancelled: boolean;
}

const OptionSchema = Type.Object({
  value: Type.String({ description: "Machine-friendly value returned to the model" }),
  label: Type.String({ description: "Human-friendly label shown in the UI" }),
  description: Type.Optional(Type.String({ description: "Optional help text shown under the label" })),
});

const QuestionSchema = Type.Object({
  id: Type.String({ description: "Stable identifier for the answer" }),
  label: Type.Optional(Type.String({ description: "Short label for summaries and tabs" })),
  prompt: Type.String({ description: "Full question text shown to the user" }),
  options: Type.Array(OptionSchema, { description: "Predefined options for the user to choose from" }),
});

const QuestionParams = Type.Object({
  questions: Type.Array(QuestionSchema, { description: "One or more questions to ask the user" }),
});

function errorResult(message: string, questions: Question[] = []) {
  return {
    content: [{ type: "text" as const, text: message }],
    details: createCancelledResult(questions) as QuestionResult,
  };
}

async function runQuestionFlow(ctx: any, questions: Question[]): Promise<QuestionResult> {
  return ctx.ui.custom<QuestionResult>((tui, theme, _kb, done) => {
    const isMulti = questions.length > 1;
    let currentTab = 0;
    let optionIndex = 0;
    let inputMode = false;
    let cachedLines: string[] | undefined;
    const answers = new Map<string, Answer>();

    const editorTheme: EditorTheme = {
      borderColor: (text) => theme.fg("accent", text),
      selectList: {
        selectedPrefix: (text) => theme.fg("accent", text),
        selectedText: (text) => theme.fg("accent", text),
        description: (text) => theme.fg("muted", text),
        scrollInfo: (text) => theme.fg("dim", text),
        noMatch: (text) => theme.fg("warning", text),
      },
    };

    const editor = new Editor(tui, editorTheme);

    function refresh() {
      cachedLines = undefined;
      tui.requestRender();
    }

    function currentQuestion(): Question | undefined {
      return questions[currentTab];
    }

    function currentOptions(): QuestionOption[] {
      return currentQuestion()?.options ?? [];
    }

    function finish(cancelled: boolean) {
      if (cancelled) {
        done(createCancelledResult(questions) as QuestionResult);
        return;
      }
      done(createAnsweredResult(questions, Array.from(answers.values())) as QuestionResult);
    }

    editor.onSubmit = (value) => {
      const question = currentQuestion();
      const trimmed = value.trim();
      if (!question || trimmed.length === 0) {
        refresh();
        return;
      }

      answers.set(question.id, createCustomAnswer(question.id, trimmed) as Answer);
      inputMode = false;
      editor.setText("");

      if (!isMulti) {
        finish(false);
        return;
      }

      currentTab = nextTabAfterAnswer(currentTab, questions.length);
      optionIndex = 0;
      refresh();
    };

    function handleInput(data: string) {
      if (inputMode) {
        if (matchesKey(data, Key.escape)) {
          inputMode = false;
          editor.setText("");
          refresh();
          return;
        }
        editor.handleInput(data);
        refresh();
        return;
      }

      if (isMulti) {
        if (matchesKey(data, Key.tab) || matchesKey(data, Key.right)) {
          currentTab = (currentTab + 1) % (questions.length + 1);
          optionIndex = 0;
          refresh();
          return;
        }
        if (matchesKey(data, Key.shift("tab")) || matchesKey(data, Key.left)) {
          currentTab = (currentTab - 1 + questions.length + 1) % (questions.length + 1);
          optionIndex = 0;
          refresh();
          return;
        }
        if (currentTab === questions.length) {
          if (matchesKey(data, Key.enter) && allQuestionsAnswered(questions, answers)) {
            finish(false);
            return;
          }
          if (matchesKey(data, Key.escape)) {
            finish(true);
            return;
          }
          return;
        }
      }

      const question = currentQuestion();
      const options = currentOptions();
      if (!question || options.length === 0) {
        return;
      }

      if (matchesKey(data, Key.up)) {
        optionIndex = Math.max(0, optionIndex - 1);
        refresh();
        return;
      }

      if (matchesKey(data, Key.down)) {
        optionIndex = Math.min(options.length - 1, optionIndex + 1);
        refresh();
        return;
      }

      if (matchesKey(data, Key.enter)) {
        const selected = options[optionIndex]!;
        if (isSomethingElseOption(selected)) {
          inputMode = true;
          editor.setText("");
          refresh();
          return;
        }

        answers.set(question.id, createPredefinedAnswer(question.id, selected, optionIndex + 1) as Answer);
        if (!isMulti) {
          finish(false);
          return;
        }

        currentTab = nextTabAfterAnswer(currentTab, questions.length);
        optionIndex = 0;
        refresh();
        return;
      }

      if (matchesKey(data, Key.escape)) {
        finish(true);
      }
    }

    function render(width: number): string[] {
      if (cachedLines) return cachedLines;

      const lines: string[] = [];
      const add = (line: string) => lines.push(truncateToWidth(line, width));
      const question = currentQuestion();
      const options = currentOptions();

      function addWrapped(text: string, color: string, firstPrefix = "", continuationPrefix = firstPrefix) {
        for (const line of wrapPrefixedText(text, width, firstPrefix, continuationPrefix)) {
          add(theme.fg(color, line));
        }
      }

      function addWrappedOption(option: QuestionOption, index: number, selected: boolean) {
        const firstPrefix = `${selected ? "> " : "  "}${index + 1}. `;
        const continuationPrefix = " ".repeat(firstPrefix.length);
        addWrapped(option.label, selected ? "accent" : "text", firstPrefix, continuationPrefix);
        if (option.description) {
          addWrapped(option.description, "muted", "     ");
        }
      }

      function addWrappedReviewAnswer(questionLabel: string, value: string) {
        const firstPrefix = ` ${questionLabel}: `;
        const continuationPrefix = " ".repeat(firstPrefix.length);
        const wrapped = wrapPrefixedText(value, width, firstPrefix, continuationPrefix);
        for (let index = 0; index < wrapped.length; index += 1) {
          const prefix = index === 0 ? firstPrefix : continuationPrefix;
          const line = wrapped[index]!;
          add(theme.fg("muted", prefix) + theme.fg("text", line.slice(prefix.length)));
        }
      }

      add(theme.fg("accent", "─".repeat(width)));

      if (isMulti) {
        const tabs: string[] = [];
        for (let index = 0; index < questions.length; index += 1) {
          const tabQuestion = questions[index]!;
          const active = index === currentTab;
          const answered = answers.has(tabQuestion.id);
          const box = answered ? "■" : "□";
          const text = ` ${box} ${tabQuestion.label} `;
          tabs.push(active ? theme.bg("selectedBg", theme.fg("text", text)) : theme.fg(answered ? "success" : "muted", text));
        }

        const submitText = " ✓ Submit ";
        const submitActive = currentTab === questions.length;
        const submitReady = allQuestionsAnswered(questions, answers);
        tabs.push(
          submitActive
            ? theme.bg("selectedBg", theme.fg("text", submitText))
            : theme.fg(submitReady ? "success" : "dim", submitText),
        );

        add(` ${tabs.join(" ")}`);
        lines.push("");
      }

      if (inputMode && question) {
        addWrapped(question.prompt, "text", " ");
        lines.push("");
        for (let index = 0; index < options.length; index += 1) {
          addWrappedOption(options[index]!, index, index === optionIndex);
        }
        lines.push("");
        add(theme.fg("muted", " Your answer:"));
        for (const line of editor.render(Math.max(1, width - 2))) {
          add(` ${line}`);
        }
      } else if (isMulti && currentTab === questions.length) {
        add(theme.fg("accent", theme.bold(" Ready to submit")));
        lines.push("");
        for (const reviewQuestion of questions) {
          const answer = answers.get(reviewQuestion.id);
          if (!answer) continue;
          const label = answer.wasCustom ? `(wrote) ${answer.label}` : `${answer.index}. ${answer.label}`;
          addWrappedReviewAnswer(reviewQuestion.label, label);
        }
        lines.push("");
        if (allQuestionsAnswered(questions, answers)) {
          add(theme.fg("success", " Press Enter to submit"));
        } else {
          add(theme.fg("warning", " All questions must be answered before submit"));
        }
      } else if (question) {
        addWrapped(question.prompt, "text", " ");
        lines.push("");
        for (let index = 0; index < options.length; index += 1) {
          addWrappedOption(options[index]!, index, index === optionIndex);
        }
      }

      lines.push("");
      if (inputMode) {
        add(theme.fg("dim", " Enter to submit • Esc to go back"));
      } else if (isMulti) {
        add(theme.fg("dim", " Tab/←→ navigate • ↑↓ select • Enter confirm • Esc cancel"));
      } else {
        add(theme.fg("dim", " ↑↓ navigate • Enter select • Esc cancel"));
      }
      add(theme.fg("accent", "─".repeat(width)));

      cachedLines = lines;
      return lines;
    }

    return {
      render,
      invalidate: () => {
        cachedLines = undefined;
      },
      handleInput,
    };
  });
}

export default function question(pi: ExtensionAPI) {
  pi.registerTool({
    name: "question",
    label: "Question",
    description:
      "Ask the user one or more multiple-choice questions. Every question automatically gets a final Something else… option for free-text answers.",
    parameters: QuestionParams,

    async execute(_toolCallId, params, _signal, _onUpdate, ctx) {
      if (!ctx.hasUI) {
        return errorResult("Error: UI not available (running in non-interactive mode)");
      }

      if (params.questions.length === 0) {
        return errorResult("Error: No questions provided");
      }

      const questions = normalizeQuestions(params.questions) as Question[];
      const result = await runQuestionFlow(ctx, questions);
      if (result.cancelled) {
        return {
          content: [{ type: "text", text: "User cancelled the question flow" }],
          details: result,
        };
      }

      return {
        content: [{ type: "text", text: summarizeAnswers(result.questions, result.answers).join("\n") }],
        details: result,
      };
    },

    renderCall(args, theme) {
      const questions = Array.isArray(args.questions) ? args.questions : [];
      const labels = questions
        .map((question: { label?: string; id?: string }) => question.label || question.id)
        .filter(Boolean)
        .join(", ");

      let text = theme.fg("toolTitle", theme.bold("question "));
      text += theme.fg("muted", `${questions.length} question${questions.length === 1 ? "" : "s"}`);
      if (labels) {
        text += theme.fg("dim", ` (${labels})`);
      }
      return new Text(text, 0, 0);
    },

    renderResult(result, _options, theme) {
      const details = result.details as QuestionResult | undefined;
      if (!details) {
        const first = result.content[0];
        return new Text(first?.type === "text" ? first.text : "", 0, 0);
      }

      if (details.cancelled) {
        return new Text(theme.fg("warning", "Cancelled"), 0, 0);
      }

      const lines = summarizeAnswers(details.questions, details.answers).map(
        (line) => `${theme.fg("success", "✓ ")}${line}`,
      );
      return new Text(lines.join("\n"), 0, 0);
    },
  });
}