import {
  chmodSync,
  existsSync,
  mkdirSync,
  readFileSync,
  renameSync,
  rmSync,
  writeFileSync,
} from "node:fs";
import { fileURLToPath } from "node:url";
import { dirname, join } from "node:path";
import { Effect, Schema, Semaphore } from "effect";
import {
  CHATTERBOX_MODEL_REVISION,
  CHATTERBOX_PACKAGE_VERSION,
} from "@local-studio/contracts/speech";
import { CommandTerminationError, resolveBinary, runCommandAsyncEffect } from "../../core/command";
import { prepareChatterboxStorage, secureSpeechDirectory } from "./storage";

export const CHATTERBOX_PACKAGE_SPEC = `chatterbox-tts==${CHATTERBOX_PACKAGE_VERSION}`;

const INSTALL_TIMEOUT_MS = 30 * 60_000;
const PREFETCH_TIMEOUT_MS = 60 * 60_000;

const InstallRecordSchema = Schema.Struct({
  packageVersion: Schema.Literal(CHATTERBOX_PACKAGE_VERSION),
  modelRevision: Schema.Literal(CHATTERBOX_MODEL_REVISION),
  gpuUuid: Schema.String,
  installedAt: Schema.String,
});

export type ChatterboxInstallStage =
  "preparing" | "creating_runtime" | "installing_package" | "prefetching_model";

export type ChatterboxRuntimeState =
  | { readonly status: "not_installed" }
  | {
      readonly status: "installing";
      readonly stage: ChatterboxInstallStage;
      readonly progress: number;
      readonly gpuUuid: string;
    }
  | {
      readonly status: "installed";
      readonly packageVersion: typeof CHATTERBOX_PACKAGE_VERSION;
      readonly modelRevision: typeof CHATTERBOX_MODEL_REVISION;
      readonly gpuUuid: string;
      readonly installedAt: string;
    }
  | { readonly status: "error"; readonly gpuUuid: string; readonly message: string };

export type ChatterboxRuntimePaths = {
  readonly runtimeDirectory: string;
  readonly pythonPath: string;
  readonly speechDirectory: string;
  readonly cacheDirectory: string;
  readonly voiceDirectory: string;
  readonly outputDirectory: string;
  readonly uploadDirectory: string;
  readonly installRecordPath: string;
  readonly workerPath: string;
};

export type SpeechRuntimeCommandResult = {
  readonly status: number | null;
  readonly stdout: string;
  readonly stderr: string;
  readonly timedOut: boolean;
  readonly exitConfirmed?: boolean | undefined;
};

export type SpeechRuntimeCommandOptions = {
  readonly env?: NodeJS.ProcessEnv | undefined;
  readonly signal?: AbortSignal | undefined;
  readonly timeoutMs: number;
};

export type ChatterboxInstallOptions = {
  readonly repair?: boolean | undefined;
};

export type SpeechRuntimeCommand = (
  command: string,
  args: string[],
  options: SpeechRuntimeCommandOptions,
) => Effect.Effect<SpeechRuntimeCommandResult, Error>;

export type ChatterboxRuntimeOptions = {
  readonly dataDirectory: string;
  readonly workerPath?: string | undefined;
  readonly environment?: NodeJS.ProcessEnv | undefined;
  readonly resolveBinary?: ((name: string) => string | null) | undefined;
  readonly runCommand?: SpeechRuntimeCommand | undefined;
  readonly now?: (() => Date) | undefined;
  readonly installTimeoutMs?: number | undefined;
  readonly prefetchTimeoutMs?: number | undefined;
};

type RuntimeDependencies = {
  readonly resolveBinary: (name: string) => string | null;
  readonly runCommand: SpeechRuntimeCommand;
  readonly now: () => Date;
  readonly installTimeoutMs: number;
  readonly prefetchTimeoutMs: number;
  readonly environment: NodeJS.ProcessEnv;
};

const defaultWorkerPath = fileURLToPath(new URL("worker.py", import.meta.url));

export const chatterboxRuntimePaths = (
  dataDirectory: string,
  workerPath = defaultWorkerPath,
): ChatterboxRuntimePaths => {
  const runtimeDirectory = join(
    dataDirectory,
    "runtime",
    "venvs",
    `chatterbox-${CHATTERBOX_PACKAGE_VERSION}`,
  );
  const speechDirectory = join(dataDirectory, "runtime", "speech");
  return {
    runtimeDirectory,
    pythonPath: join(runtimeDirectory, "bin", "python"),
    speechDirectory,
    cacheDirectory: join(speechDirectory, "huggingface"),
    voiceDirectory: join(speechDirectory, "voices"),
    outputDirectory: join(speechDirectory, "outputs"),
    uploadDirectory: join(speechDirectory, "uploads"),
    installRecordPath: join(speechDirectory, `chatterbox-${CHATTERBOX_PACKAGE_VERSION}.json`),
    workerPath,
  };
};

export const chatterboxWorkerEnvironment = (
  paths: ChatterboxRuntimePaths,
  gpuUuid: string,
  environment: NodeJS.ProcessEnv = process.env,
): NodeJS.ProcessEnv => {
  const inherited = [
    "PATH",
    "HOME",
    "TMPDIR",
    "TMP",
    "TEMP",
    "LD_LIBRARY_PATH",
    "DYLD_LIBRARY_PATH",
  ].flatMap((name) => {
    const value = environment[name];
    return value === undefined ? [] : [[name, value] as const];
  });
  return {
    ...Object.fromEntries(inherited),
    CUDA_DEVICE_ORDER: "PCI_BUS_ID",
    CUDA_VISIBLE_DEVICES: gpuUuid,
    HF_HOME: paths.cacheDirectory,
    HF_HUB_DISABLE_TELEMETRY: "1",
    PYTHONNOUSERSITE: "1",
    PYTHONUNBUFFERED: "1",
  };
};

const defaultRunCommand: SpeechRuntimeCommand = (command, args, options) =>
  runCommandAsyncEffect(command, args, {
    timeoutMs: options.timeoutMs,
    maxOutputBytes: 64 * 1024,
    ...(options.env ? { env: options.env } : {}),
    ...(options.signal ? { signal: options.signal } : {}),
  });

const errorMessage = (error: unknown): string =>
  error instanceof Error ? error.message : String(error);

const validGpuUuid = (gpuUuid: string): boolean =>
  /^GPU-[0-9A-Fa-f]{8}(?:-[0-9A-Fa-f]{4}){3}-[0-9A-Fa-f]{12}$/.test(gpuUuid);

const readInstalledState = (paths: ChatterboxRuntimePaths): ChatterboxRuntimeState => {
  if (!existsSync(paths.pythonPath) || !existsSync(paths.installRecordPath)) {
    return { status: "not_installed" };
  }
  try {
    const record = Schema.decodeUnknownSync(InstallRecordSchema)(
      JSON.parse(readFileSync(paths.installRecordPath, "utf8")),
    );
    return { status: "installed", ...record };
  } catch {
    return { status: "not_installed" };
  }
};

const failedCommandMessage = (label: string, result: SpeechRuntimeCommandResult): string => {
  if (result.timedOut) return `${label} timed out`;
  return result.stderr.trim() || result.stdout.trim() || `${label} failed`;
};

export class ChatterboxRuntime {
  readonly paths: ChatterboxRuntimePaths;
  private readonly dependencies: RuntimeDependencies;
  private readonly installSemaphore = Semaphore.makeUnsafe(1);
  private state: ChatterboxRuntimeState;
  private installPromise: Promise<ChatterboxRuntimeState> | null = null;
  private installAbort: AbortController | null = null;
  private installCancelError: CommandTerminationError | null = null;

  constructor(options: ChatterboxRuntimeOptions) {
    this.paths = chatterboxRuntimePaths(options.dataDirectory, options.workerPath);
    prepareChatterboxStorage(this.paths);
    this.dependencies = {
      resolveBinary: options.resolveBinary ?? resolveBinary,
      runCommand: options.runCommand ?? defaultRunCommand,
      now: options.now ?? ((): Date => new Date()),
      installTimeoutMs: options.installTimeoutMs ?? INSTALL_TIMEOUT_MS,
      prefetchTimeoutMs: options.prefetchTimeoutMs ?? PREFETCH_TIMEOUT_MS,
      environment: options.environment ?? process.env,
    };
    this.state = readInstalledState(this.paths);
  }

  getState(): ChatterboxRuntimeState {
    return this.state;
  }

  startInstall(gpuUuid: string, options: ChatterboxInstallOptions = {}): ChatterboxRuntimeState {
    if (
      this.state.status === "installing" ||
      (this.state.status === "installed" && !options.repair)
    ) {
      return this.state;
    }
    if (!validGpuUuid(gpuUuid)) {
      this.state = { status: "error", gpuUuid, message: "A full NVIDIA GPU UUID is required" };
      return this.state;
    }
    if (options.repair) {
      rmSync(this.paths.installRecordPath, { force: true });
      rmSync(`${this.paths.installRecordPath}.tmp`, { force: true });
    }
    const abort = new AbortController();
    this.installAbort = abort;
    this.installCancelError = null;
    const installing: ChatterboxRuntimeState = {
      status: "installing",
      stage: "preparing",
      progress: 0.05,
      gpuUuid,
    };
    this.state = installing;
    const program = this.installSemaphore
      .withPermit(this.installEffect(gpuUuid, abort.signal))
      .pipe(
        Effect.match({
          onFailure: (error) => {
            if (abort.signal.aborted && error instanceof CommandTerminationError) {
              this.installCancelError = error;
            }
            this.state = {
              status: "error",
              gpuUuid,
              message: abort.signal.aborted ? "Chatterbox install cancelled" : errorMessage(error),
            };
            return this.state;
          },
          onSuccess: (installed) => {
            this.state = installed;
            return installed;
          },
        }),
      );
    const promise = Effect.runPromise(program);
    this.installPromise = promise;
    promise.then(() => {
      if (this.installPromise === promise) this.installAbort = null;
    });
    return installing;
  }

  waitForInstall(): Promise<ChatterboxRuntimeState> {
    return this.installPromise ?? Promise.resolve(this.state);
  }

  install(
    gpuUuid: string,
    options: ChatterboxInstallOptions = {},
  ): Promise<ChatterboxRuntimeState> {
    this.startInstall(gpuUuid, options);
    return this.waitForInstall();
  }

  cancelInstall(): Promise<void> {
    const abort = this.installAbort;
    const promise = this.installPromise;
    if (!abort || !promise) return Promise.resolve();
    abort.abort();
    return promise.then(() => {
      if (this.installCancelError) throw this.installCancelError;
    });
  }

  private setInstalling(gpuUuid: string, stage: ChatterboxInstallStage, progress: number): void {
    this.state = { status: "installing", stage, progress, gpuUuid };
  }

  private commandEffect(
    label: string,
    command: string,
    args: string[],
    options: SpeechRuntimeCommandOptions,
  ): Effect.Effect<void, Error> {
    return this.dependencies
      .runCommand(command, args, options)
      .pipe(
        Effect.flatMap((result) =>
          result.exitConfirmed === false
            ? Effect.fail(new CommandTerminationError())
            : result.status === 0
              ? Effect.void
              : Effect.fail(new Error(failedCommandMessage(label, result))),
        ),
      );
  }

  private installEffect(
    gpuUuid: string,
    signal: AbortSignal,
  ): Effect.Effect<ChatterboxRuntimeState, Error> {
    const paths = this.paths;
    const dependencies = this.dependencies;
    const runtime = this;
    return Effect.gen(function* () {
      if (!existsSync(paths.workerPath)) {
        return yield* Effect.fail(new Error("Chatterbox worker resource is unavailable"));
      }
      yield* Effect.try({
        try: () => {
          mkdirSync(dirname(paths.runtimeDirectory), { recursive: true });
          prepareChatterboxStorage(paths);
        },
        catch: (error) => new Error(`Could not prepare Chatterbox storage: ${errorMessage(error)}`),
      });

      const uv = dependencies.resolveBinary("uv");
      const python = dependencies.resolveBinary("python3.11");
      const environment = chatterboxWorkerEnvironment(paths, gpuUuid, dependencies.environment);
      if (!uv && !python) {
        return yield* Effect.fail(new Error("Python 3.11 is required to install Chatterbox"));
      }

      if (!existsSync(paths.pythonPath)) {
        runtime.setInstalling(gpuUuid, "creating_runtime", 0.15);
        if (uv) {
          yield* runtime.commandEffect(
            "Creating the Chatterbox runtime",
            uv,
            ["venv", "--python", "3.11", paths.runtimeDirectory],
            { timeoutMs: dependencies.installTimeoutMs, env: environment, signal },
          );
        } else if (python) {
          yield* runtime.commandEffect(
            "Creating the Chatterbox runtime",
            python,
            ["-m", "venv", paths.runtimeDirectory],
            { timeoutMs: dependencies.installTimeoutMs, env: environment, signal },
          );
        }
      }
      secureSpeechDirectory(paths.runtimeDirectory);

      runtime.setInstalling(gpuUuid, "installing_package", 0.35);
      if (uv) {
        yield* runtime.commandEffect(
          "Installing Chatterbox",
          uv,
          [
            "pip",
            "install",
            "--python",
            paths.pythonPath,
            "--torch-backend=cu124",
            "--upgrade",
            CHATTERBOX_PACKAGE_SPEC,
          ],
          { timeoutMs: dependencies.installTimeoutMs, env: environment, signal },
        );
      } else {
        yield* runtime.commandEffect("Checking pip", paths.pythonPath, ["-m", "pip", "--version"], {
          timeoutMs: 10_000,
          env: environment,
          signal,
        });
        yield* runtime.commandEffect(
          "Installing the CUDA 12.4 PyTorch runtime",
          paths.pythonPath,
          [
            "-m",
            "pip",
            "install",
            "--upgrade",
            "torch==2.6.0+cu124",
            "torchaudio==2.6.0+cu124",
            "--index-url",
            "https://download.pytorch.org/whl/cu124",
          ],
          { timeoutMs: dependencies.installTimeoutMs, env: environment, signal },
        );
        yield* runtime.commandEffect(
          "Installing Chatterbox",
          paths.pythonPath,
          ["-m", "pip", "install", "--upgrade", CHATTERBOX_PACKAGE_SPEC],
          { timeoutMs: dependencies.installTimeoutMs, env: environment, signal },
        );
      }

      runtime.setInstalling(gpuUuid, "prefetching_model", 0.75);
      yield* runtime.commandEffect(
        "Prefetching the pinned Chatterbox Turbo model",
        paths.pythonPath,
        [paths.workerPath, "--prefetch"],
        {
          timeoutMs: dependencies.prefetchTimeoutMs,
          env: environment,
          signal,
        },
      );

      if (signal.aborted) return yield* Effect.fail(new Error("Chatterbox install cancelled"));

      const installedAt = dependencies.now().toISOString();
      const installed: ChatterboxRuntimeState = {
        status: "installed",
        packageVersion: CHATTERBOX_PACKAGE_VERSION,
        modelRevision: CHATTERBOX_MODEL_REVISION,
        gpuUuid,
        installedAt,
      };
      yield* Effect.try({
        try: () => {
          const temporaryPath = `${paths.installRecordPath}.tmp`;
          writeFileSync(
            temporaryPath,
            JSON.stringify({
              packageVersion: CHATTERBOX_PACKAGE_VERSION,
              modelRevision: CHATTERBOX_MODEL_REVISION,
              gpuUuid,
              installedAt,
            }),
            { mode: 0o600 },
          );
          renameSync(temporaryPath, paths.installRecordPath);
          chmodSync(paths.installRecordPath, 0o600);
        },
        catch: (error) =>
          new Error(`Could not record the Chatterbox install: ${errorMessage(error)}`),
      });
      return installed;
    });
  }
}