import { logError, logInfo, logWarn } from "@/logger";
import type CopilotPlugin from "@/main";
import { getSettings } from "@/settings/model";
import { App, FileSystemAdapter, Platform } from "obsidian";
import { MethodUnsupportedError } from "./errors";
import { backendStateSignature } from "./translateBackendState";
import type {
  BackendDescriptor,
  BackendId,
  BackendProcess,
  BackendState,
  SessionId,
} from "./types";

/**
 * Warm result of a successful preload — the still-running backend process,
 * the probe session id it owns, and the state snapshot it reported. Handed
 * off via `takeWarm()` so the manager can skip a fresh subprocess spawn +
 * `initialize` handshake on the first chat open, and (when the user opens
 * a fresh chat on that backend) adopt the probe session as the user's
 * session instead of paying another `newSession` round-trip.
 */
export interface WarmBackend {
  proc: BackendProcess;
  probeSessionId: SessionId;
  state: BackendState;
}

/**
 * Plugin-lifetime cache of per-backend session state and the running probe
 * subprocess that produced it. Backends expose `BackendState` only as a
 * side-effect of session creation / resume / load, so without this preload
 * the picker would show no entries for non-active backends and would blink
 * empty during the round-trip on a fresh session.
 *
 * Probes once per backend at startup: prefer resume of a persisted probe
 * sessionId, fall back to load, then to new (and persist the new id so the
 * next reload can reuse it — keeps the agent-side session store at one stale
 * entry per machine instead of growing with each reload).
 *
 * The probe subprocess is **kept warm** until the manager consumes it via
 * `takeWarm(backendId)`. That removes the warm subprocess spawn from the
 * critical path of the first chat-open: instead of preload booting a
 * subprocess just to read its catalog and immediately shutting it down,
 * the same subprocess becomes the manager's backend process on first use.
 */
export class AgentModelPreloader {
  private readonly warm = new Map<BackendId, WarmBackend>();
  // State-only cache for backends whose warm entry has already been
  // consumed (or that pushed updates after consumption via `setCached`).
  // The active session's `attachModelCacheSync` writes here.
  private readonly cache = new Map<BackendId, BackendState>();
  private readonly inflight = new Map<BackendId, Promise<void>>();
  private readonly listeners = new Set<() => void>();
  // Per-warm-entry exit-listener teardowns. Wired when the warm entry is
  // recorded so we can clear it if the probe subprocess dies before the
  // manager takes ownership.
  private readonly warmExitUnsubs = new Map<BackendId, () => void>();
  private disposed = false;

  constructor(
    private readonly app: App,
    private readonly plugin: CopilotPlugin,
    private readonly resolveDescriptor: (id: BackendId) => BackendDescriptor | undefined
  ) {}

  getCachedBackendState(backendId: BackendId): BackendState | null {
    return this.warm.get(backendId)?.state ?? this.cache.get(backendId) ?? null;
  }

  /**
   * Replace the cached entry for `backendId`. No-op when the signature
   * is unchanged, to avoid spurious picker rebuilds. Live sessions push
   * here via `attachModelCacheSync` after the warm entry is consumed.
   */
  setCached(backendId: BackendId, state: BackendState): void {
    if (this.disposed) return;
    const prev = this.getCachedBackendState(backendId);
    if (backendStateSignature(prev) === backendStateSignature(state)) return;
    this.cache.set(backendId, state);
    this.notify();
  }

  /**
   * Remove all cached state for `backendId` after its backend is restarted.
   * Drops the warm subprocess if it hasn't been taken yet so a fresh probe
   * runs on the next `preload(backendId)` call.
   */
  clearCached(backendId: BackendId): void {
    if (this.disposed) return;
    let changed = this.cache.delete(backendId);
    const warm = this.warm.get(backendId);
    if (warm) {
      this.warm.delete(backendId);
      this.warmExitUnsubs.get(backendId)?.();
      this.warmExitUnsubs.delete(backendId);
      // Best-effort shutdown of the abandoned warm proc.
      warm.proc.shutdown().catch((e) => {
        logWarn(`[AgentMode] preload clearCached: shutdown of warm ${backendId} failed`, e);
      });
      changed = true;
    }
    if (changed) this.notify();
  }

  /**
   * Hand the warm backend (process + probe session id + state snapshot) to
   * the manager. Single-shot: removes the entry so subsequent callers see
   * `null` and the manager owns lifetime of the process from here on.
   */
  takeWarm(backendId: BackendId): WarmBackend | null {
    const entry = this.warm.get(backendId);
    if (!entry) return null;
    this.warm.delete(backendId);
    this.warmExitUnsubs.get(backendId)?.();
    this.warmExitUnsubs.delete(backendId);
    // Keep the last-known state in `cache` so the picker still reads it
    // until the live session syncs its own state back in. Without this the
    // picker would briefly see `null` between `takeWarm` and the first
    // `attachModelCacheSync` write.
    this.cache.set(backendId, entry.state);
    return entry;
  }

  /** Best-effort probe; failures are logged and swallowed. Dedupes per backend. */
  preload(backendId: BackendId): Promise<void> {
    if (this.disposed) return Promise.resolve();
    const existing = this.inflight.get(backendId);
    if (existing) return existing;
    const promise = this.runProbe(backendId).finally(() => {
      this.inflight.delete(backendId);
    });
    this.inflight.set(backendId, promise);
    return promise;
  }

  subscribe(listener: () => void): () => void {
    this.listeners.add(listener);
    return () => this.listeners.delete(listener);
  }

  shutdown(): void {
    this.disposed = true;
    this.cache.clear();
    this.inflight.clear();
    this.listeners.clear();
    for (const [backendId, warm] of this.warm) {
      this.warmExitUnsubs.get(backendId)?.();
      warm.proc.shutdown().catch((e) => {
        logWarn(`[AgentMode] preload shutdown: warm ${backendId} shutdown failed`, e);
      });
    }
    this.warm.clear();
    this.warmExitUnsubs.clear();
  }

  private notify(): void {
    for (const l of this.listeners) {
      try {
        l();
      } catch (e) {
        logWarn("[AgentMode] preload listener threw", e);
      }
    }
  }

  private async runProbe(backendId: BackendId): Promise<void> {
    if (Platform.isMobile) return;
    const adapter = this.app.vault.adapter;
    if (!(adapter instanceof FileSystemAdapter)) return;
    const cwd = adapter.getBasePath();

    const descriptor = this.resolveDescriptor(backendId);
    if (!descriptor) {
      logWarn(`[AgentMode] preload skipped: unknown backend ${backendId}`);
      return;
    }
    if (descriptor.getInstallState(getSettings()).kind !== "ready") return;

    const proc = descriptor.createBackendProcess({
      plugin: this.plugin,
      app: this.app,
      clientVersion: this.plugin.manifest.version,
      descriptor,
    });

    let probe: { sessionId: SessionId; state: BackendState } | null = null;
    try {
      await proc.start?.();
      const storedId = descriptor.getProbeSessionId?.(getSettings());
      probe = await this.fetchInitialState(proc, descriptor, backendId, storedId, cwd);
    } catch (err) {
      logError(`[AgentMode] preload ${backendId} failed`, err);
    }

    if (this.disposed || !probe || (!probe.state.model && !probe.state.mode)) {
      if (probe) {
        logInfo(`[AgentMode] preload ${backendId}: agent did not report any initial state`);
      }
      try {
        await proc.shutdown();
      } catch (e) {
        logWarn(`[AgentMode] preload ${backendId}: shutdown failed`, e);
      }
      return;
    }

    // Probe succeeded — retain the running subprocess as a warm entry so
    // the first chat-open can adopt it instead of paying another spawn +
    // initialize round-trip.
    const warm: WarmBackend = {
      proc,
      probeSessionId: probe.sessionId,
      state: probe.state,
    };
    const exitUnsub = proc.onExit(() => {
      if (this.disposed) return;
      // Subprocess died before the manager claimed it. Drop the warm
      // entry; next createSession will spawn a fresh one through the
      // descriptor.
      if (this.warm.get(backendId) === warm) {
        this.warm.delete(backendId);
        this.warmExitUnsubs.delete(backendId);
        this.notify();
      }
    });
    this.warm.set(backendId, warm);
    this.warmExitUnsubs.set(backendId, exitUnsub);
    logProbeResult(backendId, "session probe", probe.state);
    this.notify();
  }

  private async fetchInitialState(
    proc: BackendProcess,
    descriptor: BackendDescriptor,
    backendId: BackendId,
    storedId: string | undefined,
    cwd: string
  ): Promise<{ sessionId: SessionId; state: BackendState }> {
    type Strategy = {
      label: string;
      sessionId: string;
      run: () => Promise<{ sessionId: string; state: BackendState }>;
    };
    const strategies: Strategy[] = [];
    if (storedId) {
      strategies.push({
        label: `resumed probe session ${storedId}`,
        sessionId: storedId,
        run: () => proc.resumeSession({ sessionId: storedId, cwd, mcpServers: [] }),
      });
      strategies.push({
        label: `loaded probe session ${storedId}`,
        sessionId: storedId,
        run: () => proc.loadSession({ sessionId: storedId, cwd, mcpServers: [] }),
      });
    }

    for (const { label, sessionId, run } of strategies) {
      try {
        // Register a no-op handler before the call so updates emitted
        // during the call are demuxed against this sessionId rather than
        // buffered indefinitely. The manager overrides this with the real
        // handler when it adopts the session.
        proc.registerSessionHandler(sessionId, () => {});
        const resp = await run();
        logInfo(`[AgentMode] preload ${backendId}: ${label}`);
        return { sessionId: resp.sessionId, state: resp.state };
      } catch (err) {
        if (!(err instanceof MethodUnsupportedError)) {
          logWarn(`[AgentMode] preload ${backendId}: ${label} failed (will fall back)`, err);
        }
      }
    }

    const resp = await proc.newSession({ cwd, mcpServers: [] });
    proc.registerSessionHandler(resp.sessionId, () => {});
    logInfo(`[AgentMode] preload ${backendId}: created probe session ${resp.sessionId}`);
    if (descriptor.persistProbeSessionId) {
      try {
        await descriptor.persistProbeSessionId(resp.sessionId, this.plugin);
      } catch (e) {
        logWarn(`[AgentMode] preload ${backendId}: persistProbeSessionId failed`, e);
      }
    }
    return { sessionId: resp.sessionId, state: resp.state };
  }
}

function logProbeResult(backendId: BackendId, label: string, state: BackendState): void {
  const ids = state.model?.availableModels.map((m) => m.baseModelId).join(", ") ?? "";
  const modeOpts = state.mode?.options.map((o) => o.value).join(", ") ?? "";
  const currentBaseId = state.model?.current.baseModelId ?? "-";
  const currentEntry = state.model?.availableModels.find((e) => e.baseModelId === currentBaseId);
  const effortOpts = currentEntry?.effortOptions.map((o) => o.value ?? "default").join(", ") ?? "";
  logInfo(
    `[AgentMode] preload ${backendId} (${label}): models=[${ids}] (current=${currentBaseId}), ` +
      `mode=[${modeOpts}] effort=[${effortOpts}]`
  );
}
