import { existsSync } from "node:fs";
import { isAbsolute, resolve } from "node:path";
import {
  TypeWhisperClient,
  TypeWhisperConnectionError,
  TypeWhisperError,
  TypeWhisperHttpError,
  discoverTypeWhisperAPI,
  type DictionaryCorrection,
  type DictionaryTermEntry,
  type DiscoveryOptions,
  type TranscribeFileOptions
} from "@typewhisper/mcp";
import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
import { Type } from "typebox";

export type TypeWhisperAPI = Pick<
  TypeWhisperClient,
  | "status"
  | "models"
  | "transcribeFile"
  | "searchHistory"
  | "listDictionaryTerms"
  | "upsertDictionaryTerms"
  | "deleteDictionaryTerm"
  | "listDictionaryCorrections"
  | "upsertDictionaryCorrection"
  | "deleteDictionaryCorrection"
>;

export interface TypeWhisperExtensionOptions {
  client?: TypeWhisperAPI;
  discovery?: DiscoveryOptions;
  fetch?: typeof fetch;
}

const noParameters = Type.Object({});

const transcribeParameters = Type.Object({
  path: Type.String({ description: "Audio or video file path, absolute or relative to Pi's current directory." }),
  language: Type.Optional(Type.String({ description: "Exact ISO 639-1 source language, for example de or en." })),
  languageHints: Type.Optional(Type.Array(Type.String(), { description: "Ordered language hints. Do not combine with language." })),
  task: Type.Optional(Type.Union([Type.Literal("transcribe"), Type.Literal("translate")], { description: "Defaults to transcribe." })),
  targetLanguage: Type.Optional(Type.String({ description: "Target language for Apple Translate output." })),
  engine: Type.Optional(Type.String({ description: "Optional TypeWhisper engine/provider id override." })),
  model: Type.Optional(Type.String({ description: "Optional model id override." })),
  awaitDownload: Type.Optional(Type.Boolean({ description: "Wait for model restore/download instead of failing fast." })),
  applyCorrections: Type.Optional(Type.Boolean({ description: "Apply TypeWhisper Dictionary Corrections. Defaults to true." }))
});

const historyParameters = Type.Object({
  query: Type.Optional(Type.String({ minLength: 1, description: "Search query for transcription history." })),
  limit: Type.Optional(Type.Integer({ minimum: 1, maximum: 50, description: "Maximum entries. Defaults to 10." })),
  offset: Type.Optional(Type.Integer({ minimum: 0, description: "Pagination offset. Defaults to 0." }))
});

const termEntry = Type.Object({
  term: Type.String({ minLength: 1 }),
  ctcMinSimilarity: Type.Optional(Type.Number({ minimum: 0, maximum: 1 }))
});

const upsertTermsParameters = Type.Object({
  terms: Type.Optional(Type.Array(Type.String({ minLength: 1 }))),
  termEntries: Type.Optional(Type.Array(termEntry)),
  replace: Type.Optional(Type.Boolean({ description: "Replace the full term list instead of merging. Defaults to false." }))
});

const deleteTermParameters = Type.Object({ term: Type.String({ minLength: 1 }) });
const upsertCorrectionParameters = Type.Object({
  original: Type.String({ minLength: 1 }),
  replacement: Type.String(),
  caseSensitive: Type.Optional(Type.Boolean())
});
const deleteCorrectionParameters = Type.Object({ original: Type.String({ minLength: 1 }) });

export function registerTypeWhisperExtension(
  pi: ExtensionAPI,
  options: TypeWhisperExtensionOptions = {}
): void {
  const client = options.client ?? new TypeWhisperClient({
    connection: discoverTypeWhisperAPI(options.discovery),
    fetch: options.fetch
  });

  pi.registerTool({
    name: "typewhisper_status",
    label: "TypeWhisper Status",
    description: "Check whether the local TypeWhisper app is reachable and has a transcription model ready.",
    promptSnippet: "Check the local TypeWhisper transcription service",
    parameters: noParameters,
    async execute() {
      return runTool(async () => {
        const status = await client.status();
        const summary = status.status === "ready"
          ? `TypeWhisper is ready using ${status.engine ?? "unknown engine"} / ${status.model ?? "unknown model"}.`
          : "TypeWhisper is reachable, but no transcription model is ready.";
        return toolResult(summary, status);
      });
    }
  });

  pi.registerTool({
    name: "typewhisper_list_models",
    label: "TypeWhisper Models",
    description: "List available TypeWhisper transcription engines and models.",
    promptSnippet: "List local TypeWhisper engines and models",
    parameters: noParameters,
    async execute() {
      return runTool(async () => {
        const models = await client.models();
        return toolResult(`TypeWhisper returned ${models.models.length} model(s).`, models);
      });
    }
  });

  pi.registerTool({
    name: "typewhisper_transcribe_file",
    label: "TypeWhisper Transcribe",
    description: "Transcribe a local audio or video file through the running TypeWhisper app.",
    promptSnippet: "Transcribe local audio and video files with TypeWhisper",
    promptGuidelines: [
      "Use typewhisper_transcribe_file when the user asks to transcribe a local audio or video file with TypeWhisper."
    ],
    parameters: transcribeParameters,
    async execute(_toolCallId, input, signal, _onUpdate, ctx) {
      return runTool(async () => {
        if (signal?.aborted) throw new TypeWhisperError("Transcription cancelled.");
        const request = normalizeTranscribeRequest(input, ctx.cwd);
        const transcript = await client.transcribeFile(request);
        return toolResult(transcript.text || "TypeWhisper returned an empty transcription.", transcript);
      });
    }
  });

  pi.registerTool({
    name: "typewhisper_search_history",
    label: "TypeWhisper History",
    description: "Search recent TypeWhisper transcription history.",
    promptSnippet: "Search TypeWhisper transcription history",
    parameters: historyParameters,
    async execute(_toolCallId, input) {
      return runTool(async () => {
        const response = await client.searchHistory({
          query: input.query,
          limit: input.limit ?? 10,
          offset: input.offset ?? 0
        });
        const noun = response.entries.length === 1 ? "entry" : "entries";
        return toolResult(`Found ${response.entries.length} history ${noun} out of ${response.total}.`, response);
      });
    }
  });

  pi.registerTool({
    name: "typewhisper_list_dictionary_terms",
    label: "TypeWhisper Dictionary Terms",
    description: "List recognition terms configured in the TypeWhisper dictionary.",
    promptSnippet: "List TypeWhisper recognition dictionary terms",
    parameters: noParameters,
    async execute() {
      return runTool(async () => {
        const response = await client.listDictionaryTerms();
        return toolResult(`TypeWhisper has ${response.count} dictionary term(s).`, response);
      });
    }
  });

  pi.registerTool({
    name: "typewhisper_upsert_dictionary_terms",
    label: "TypeWhisper Upsert Terms",
    description: "Merge or replace TypeWhisper recognition terms.",
    promptSnippet: "Add or update TypeWhisper recognition dictionary terms",
    parameters: upsertTermsParameters,
    async execute(_toolCallId, input) {
      return runTool(async () => {
        const response = await client.upsertDictionaryTerms(normalizeTermUpsert(input));
        return toolResult(`TypeWhisper now has ${response.count} dictionary term(s).`, response);
      });
    }
  });

  pi.registerTool({
    name: "typewhisper_delete_dictionary_term",
    label: "TypeWhisper Delete Term",
    description: "Delete one recognition term from the TypeWhisper dictionary.",
    promptSnippet: "Delete a TypeWhisper recognition dictionary term",
    parameters: deleteTermParameters,
    async execute(_toolCallId, input) {
      return runTool(async () => {
        const response = await client.deleteDictionaryTerm(input.term);
        const summary = response.deleted
          ? `Deleted dictionary term "${input.term}".`
          : `Dictionary term "${input.term}" was not present.`;
        return toolResult(summary, response);
      });
    }
  });

  pi.registerTool({
    name: "typewhisper_list_dictionary_corrections",
    label: "TypeWhisper Dictionary Corrections",
    description: "List post-transcription corrections configured in the TypeWhisper dictionary.",
    promptSnippet: "List TypeWhisper post-transcription corrections",
    parameters: noParameters,
    async execute() {
      return runTool(async () => {
        const response = await client.listDictionaryCorrections();
        return toolResult(`TypeWhisper has ${response.count} dictionary correction(s).`, response);
      });
    }
  });

  pi.registerTool({
    name: "typewhisper_upsert_dictionary_correction",
    label: "TypeWhisper Upsert Correction",
    description: "Add or update one post-transcription correction in the TypeWhisper dictionary.",
    promptSnippet: "Add or update a TypeWhisper dictionary correction",
    parameters: upsertCorrectionParameters,
    async execute(_toolCallId, input) {
      return runTool(async () => {
        const correction: DictionaryCorrection = {
          original: input.original,
          replacement: input.replacement,
          caseSensitive: input.caseSensitive ?? false
        };
        const response = await client.upsertDictionaryCorrection(correction);
        return toolResult(`Upserted dictionary correction for "${input.original}".`, response);
      });
    }
  });

  pi.registerTool({
    name: "typewhisper_delete_dictionary_correction",
    label: "TypeWhisper Delete Correction",
    description: "Delete one post-transcription correction from the TypeWhisper dictionary.",
    promptSnippet: "Delete a TypeWhisper dictionary correction",
    parameters: deleteCorrectionParameters,
    async execute(_toolCallId, input) {
      return runTool(async () => {
        const response = await client.deleteDictionaryCorrection(input.original);
        const summary = response.deleted
          ? `Deleted dictionary correction for "${input.original}".`
          : `Dictionary correction for "${input.original}" was not present.`;
        return toolResult(summary, response);
      });
    }
  });
}

export default function typeWhisperExtension(pi: ExtensionAPI): void {
  registerTypeWhisperExtension(pi);
}

function normalizeTranscribeRequest(
  input: {
    path: string;
    language?: string;
    languageHints?: string[];
    task?: "transcribe" | "translate";
    targetLanguage?: string;
    engine?: string;
    model?: string;
    awaitDownload?: boolean;
    applyCorrections?: boolean;
  },
  cwd: string
): TranscribeFileOptions {
  const cleanPath = input.path.startsWith("@") ? input.path.slice(1) : input.path;
  const absolutePath = isAbsolute(cleanPath) ? cleanPath : resolve(cwd, cleanPath);

  if (!existsSync(absolutePath)) {
    throw new TypeWhisperError(`File not found: ${absolutePath}`);
  }
  if (input.language && input.languageHints?.length) {
    throw new TypeWhisperError("Use either language or languageHints, not both.");
  }

  return {
    path: absolutePath,
    language: input.language,
    languageHints: input.languageHints,
    task: input.task,
    targetLanguage: input.targetLanguage,
    engine: input.engine,
    model: input.model,
    awaitDownload: input.awaitDownload ?? false,
    applyCorrections: input.applyCorrections ?? true
  };
}

function normalizeTermUpsert(input: {
  terms?: string[];
  termEntries?: Array<{ term: string; ctcMinSimilarity?: number }>;
  replace?: boolean;
}): { terms?: string[]; termEntries?: DictionaryTermEntry[]; replace?: boolean } {
  const hasTerms = Boolean(input.terms?.length);
  const hasTermEntries = Boolean(input.termEntries?.length);
  if (hasTerms === hasTermEntries) {
    throw new TypeWhisperError("Provide exactly one of terms or termEntries.");
  }

  return {
    terms: input.terms,
    termEntries: input.termEntries?.map((entry) => ({
      term: entry.term,
      ctc_min_similarity: entry.ctcMinSimilarity
    })),
    replace: input.replace ?? false
  };
}

async function runTool<T>(fn: () => Promise<T>): Promise<T | ReturnType<typeof errorResult>> {
  try {
    return await fn();
  } catch (error) {
    return errorResult(error);
  }
}

function toolResult(summary: string, payload: unknown) {
  return {
    content: [{ type: "text" as const, text: summary }],
    details: payload
  };
}

function errorResult(error: unknown) {
  return {
    content: [{ type: "text" as const, text: formatError(error) }],
    details: {},
    isError: true
  };
}

function formatError(error: unknown): string {
  if (error instanceof TypeWhisperHttpError) {
    if (error.status === 401) {
      return "TypeWhisper API authentication failed. Restart TypeWhisper so the discovery token refreshes, or pass TYPEWHISPER_API_TOKEN.";
    }
    if (error.status === 503) {
      return "TypeWhisper has no transcription model ready. Load or select a model in TypeWhisper first.";
    }
    return `TypeWhisper API returned HTTP ${error.status}: ${error.message}`;
  }
  if (error instanceof TypeWhisperConnectionError || error instanceof TypeWhisperError) {
    return error.message;
  }
  if (error instanceof Error) return error.message;
  return "Unknown TypeWhisper error.";
}