import * as path from "path";
import * as vscode from "vscode";
import { Backend, nextQuestion } from "./llm";

const SECRET_KEY = "nextObviousQuestion.openrouterApiKey";

export function activate(context: vscode.ExtensionContext) {
  // Last question per document, so switching files keeps each one's question.
  const questions = new Map<string, string>();
  // Document version each question was asked for, to tell whether it has been edited since.
  const askedAt = new Map<string, number>();
  // Questions rejected with "Different question" since the last fresh ask or edit.
  const rejected = new Map<string, string[]>();
  const pending = new Set<string>();

  const item = vscode.window.createStatusBarItem(vscode.StatusBarAlignment.Left, 100);
  item.command = "nextObviousQuestion.menu";
  context.subscriptions.push(item);

  const render = () => {
    const doc = vscode.window.activeTextEditor?.document;
    if (!doc || !isSupported(doc)) {
      item.hide();
      return;
    }
    const key = doc.uri.toString();
    const q = questions.get(key);
    if (pending.has(key)) {
      item.text = "$(sync~spin) Thinking of the next question…";
      item.tooltip = undefined;
    } else if (q) {
      item.text = `$(question) ${q}`;
      item.tooltip =
        askedAt.get(key) === doc.version
          ? `${q}\n\nClick to ask again or get a different question.`
          : `${q}\n\nThe document has changed. Click for a new question.`;
    } else {
      item.text = "$(question) Ask next question";
      item.tooltip = "Ask what a reader would want to know next.";
    }
    item.show();
  };

  const ask = async (different: boolean) => {
    const doc = vscode.window.activeTextEditor?.document;
    if (!doc || !isSupported(doc)) {
      vscode.window.showInformationMessage("Open a Markdown (.md) or LaTeX (.tex) file first.");
      return;
    }
    const key = doc.uri.toString();
    if (pending.has(key)) return;

    const shown = questions.get(key);
    const avoid = different && shown ? [...(rejected.get(key) ?? []), shown] : [];
    if (avoid.length === 0) rejected.delete(key);

    const text = doc.getText();
    if (text.trim() === "" && avoid.length === 0) {
      questions.set(key, "What is this?");
      askedAt.set(key, doc.version);
      render();
      return;
    }

    const backend = await getBackend(context);
    if (!backend) return;

    const version = doc.version;
    pending.add(key);
    render();
    try {
      const q = await nextQuestion(backend, path.basename(doc.fileName), text, avoid);
      if (q) {
        questions.set(key, q);
        askedAt.set(key, version);
        // An edit while waiting starts a new cycle, so only keep the list if the text is unchanged.
        if (avoid.length > 0 && doc.version === version) rejected.set(key, avoid);
      }
    } catch (err) {
      vscode.window.showErrorMessage(`Next Obvious Question: ${(err as Error).message}`);
    } finally {
      pending.delete(key);
      render();
    }
  };

  const menu = async () => {
    const doc = vscode.window.activeTextEditor?.document;
    if (!doc || pending.has(doc.uri.toString())) return;
    // No question yet, or the document was edited since: just ask, no menu.
    if (askedAt.get(doc.uri.toString()) !== doc.version) return ask(false);
    const choice = await vscode.window.showQuickPick(
      [
        { label: "$(arrow-swap) Different question", different: true },
        { label: "$(refresh) Ask again", description: "start over", different: false },
      ],
      { placeHolder: questions.get(doc.uri.toString()) },
    );
    if (choice) await ask(choice.different);
  };

  context.subscriptions.push(
    vscode.commands.registerCommand("nextObviousQuestion.ask", () => ask(false)),
    vscode.commands.registerCommand("nextObviousQuestion.askDifferent", () => ask(true)),
    vscode.commands.registerCommand("nextObviousQuestion.menu", menu),
    vscode.commands.registerCommand("nextObviousQuestion.setOpenRouterKey", () => promptForKey(context)),
    vscode.window.onDidChangeActiveTextEditor(render),
    vscode.workspace.onDidChangeTextDocument((e) => {
      if (e.contentChanges.length === 0) return;
      rejected.delete(e.document.uri.toString());
      if (e.document === vscode.window.activeTextEditor?.document) render();
    }),
    vscode.workspace.onDidCloseTextDocument((doc) => {
      questions.delete(doc.uri.toString());
      askedAt.delete(doc.uri.toString());
      rejected.delete(doc.uri.toString());
    }),
  );
  render();
}

export function deactivate() {}

function isSupported(doc: vscode.TextDocument): boolean {
  if (["markdown", "latex", "tex"].includes(doc.languageId)) return true;
  const ext = path.extname(doc.fileName).toLowerCase();
  return ext === ".md" || ext === ".tex";
}

async function getBackend(context: vscode.ExtensionContext): Promise<Backend | undefined> {
  const config = vscode.workspace.getConfiguration("nextObviousQuestion");
  if (config.get<string>("backend") === "openrouter") {
    const apiKey =
      (await context.secrets.get(SECRET_KEY)) || process.env.OPENROUTER_API_KEY || (await promptForKey(context));
    if (!apiKey) return undefined;
    return { kind: "openrouter", apiKey, model: config.get<string>("openrouterModel") || "anthropic/claude-sonnet-5.5" };
  }
  return {
    kind: "claudeCode",
    path: config.get<string>("claudeCodePath") || "claude",
    model: config.get<string>("claudeCodeModel") || "sonnet",
  };
}

async function promptForKey(context: vscode.ExtensionContext): Promise<string | undefined> {
  const key = await vscode.window.showInputBox({
    title: "OpenRouter API key",
    prompt: "Stored in VS Code's secret storage.",
    password: true,
    ignoreFocusOut: true,
  });
  if (key) await context.secrets.store(SECRET_KEY, key.trim());
  return key?.trim() || undefined;
}