import { mkdirSync } from "node:fs";
import { join } from "node:path";
import {
  type AgentRuntime,
  type AgentRuntimeConfig,
  type BotMessage,
  getSessionKey
} from "./core.js";
import { resolveRuntimeModel } from "./config.js";
import type { AgentSession } from "@mariozechner/pi-coding-agent";
import {
  AuthStorage,
  createAgentSession,
  ModelRegistry,
  SessionManager,
  SettingsManager
} from "@mariozechner/pi-coding-agent";
import { createPiResourceLoader } from "./pi-resources.js";
import { SessionStore } from "./session-store.js";
import { cozeWebSearchTool, cozeWebFetchTool } from "./tools/index.js";

export type PiAgentStreamHandlers = {
  onMeta?: (meta: { sessionKey: string }) => void;
  onDelta?: (delta: string) => void;
  onError?: (error: string) => void;
  /**
   * Optional hook for forwarding raw session events (useful for debugging/telemetry).
   * Keep this best-effort: callers should not rely on a stable event schema here.
   */
  onEvent?: (event: unknown) => void;
};

/**
 * Extension of AgentRuntime used by the dashboard.
 * Channels keep using the stable `run()` API.
 */
export interface PiAgentRuntime extends AgentRuntime {
  stream(message: BotMessage, handlers: PiAgentStreamHandlers): Promise<{ sessionKey: string; finalText: string }>;
  getSessionIfExists(sessionKey: string): AgentSession | undefined;
  listSessionKeys(): string[];
  ensureSessionLoaded(sessionKey: string): Promise<AgentSession | undefined>;
  abortSession(sessionKey: string): Promise<void>;
  resetSession(sessionKey: string): Promise<void>;
}

function extractAssistantText(session: AgentSession): string {
  const messages = [...session.state.messages].reverse();

  for (const message of messages) {
    if ((message as { role?: string }).role !== "assistant") {
      continue;
    }

    const content = (message as { content?: unknown }).content;
    if (typeof content === "string") {
      return content.trim();
    }

    if (!Array.isArray(content)) {
      return "";
    }

    return content
      .flatMap((part) => {
        if (!part || typeof part !== "object") {
          return [];
        }

        const typedPart = part as { type?: unknown; text?: unknown };
        return typedPart.type === "text" && typeof typedPart.text === "string" ? [typedPart.text] : [];
      })
      .join("")
      .trim();
  }

  return "";
}

function extractAssistantTextFromMessage(message: unknown): string {
  if (!message || typeof message !== "object") return "";
  if ((message as { role?: unknown }).role !== "assistant") return "";

  const content = (message as { content?: unknown }).content;
  if (typeof content === "string") {
    return content;
  }
  if (!Array.isArray(content)) {
    return "";
  }
  return content
    .flatMap((part) => {
      if (!part || typeof part !== "object") return [];
      const typedPart = part as { type?: unknown; text?: unknown };
      return typedPart.type === "text" && typeof typedPart.text === "string" ? [typedPart.text] : [];
    })
    .join("");
}

export function createMockAgentRuntime(): PiAgentRuntime {
  const sessionHistory = new Map<string, string[]>();

  return {
    async run(message: BotMessage): Promise<string> {
      const sessionKey = getSessionKey(message);
      const history = sessionHistory.get(sessionKey) ?? [];
      history.push(message.text);
      sessionHistory.set(sessionKey, history);
      return `mock:${sessionKey}: ${message.text}`;
    },
    async stream(message: BotMessage, handlers: PiAgentStreamHandlers) {
      const sessionKey = getSessionKey(message);
      handlers.onMeta?.({ sessionKey });
      const history = sessionHistory.get(sessionKey) ?? [];
      history.push(message.text);
      sessionHistory.set(sessionKey, history);
      const text = `mock:${sessionKey}: ${message.text}`;
      handlers.onDelta?.(text);
      return { sessionKey, finalText: text };
    },
    getSessionIfExists(_sessionKey: string): AgentSession | undefined {
      return undefined;
    },
    listSessionKeys(): string[] {
      return Array.from(sessionHistory.keys());
    },
    async ensureSessionLoaded(_sessionKey: string): Promise<AgentSession | undefined> {
      return undefined;
    },
    async abortSession(_sessionKey: string): Promise<void> {},
    async resetSession(_sessionKey: string): Promise<void> {},
    async dispose(): Promise<void> {}
  } satisfies PiAgentRuntime;
}

export async function createPiAgentRuntime(config: AgentRuntimeConfig): Promise<PiAgentRuntime> {
  const cwd = config.cwd ?? process.cwd();
  const agentDir = config.agentDir;
  const sessionRootDir = agentDir ?? cwd;
  const authStorage = agentDir ? AuthStorage.create(join(agentDir, "auth.json")) : AuthStorage.create();
  const modelRegistry = new ModelRegistry(
    authStorage,
    agentDir ? join(agentDir, "models.json") : undefined
  );
  const settingsManager = SettingsManager.create(cwd, agentDir);
  const resourceLoader = createPiResourceLoader({ cwd, agentDir, settingsManager });
  const resolvedModel = resolveRuntimeModel({
    provider: config.provider,
    model: config.model,
    baseUrl: config.baseUrl,
    configPath: config.configPath
  });
  const sessions = new Map<string, AgentSession>();
  const pendingSessionLoads = new Map<string, Promise<AgentSession>>();
  const sessionGenerations = new Map<string, number>();
  const sessionStore = new SessionStore(sessionRootDir);

  if (agentDir) {
    mkdirSync(agentDir, { recursive: true });
  }

  await resourceLoader.reload();
  sessionStore.load();

  if (resolvedModel?.apiKey) {
    authStorage.setRuntimeApiKey(resolvedModel.model.provider, resolvedModel.apiKey);
  }

  async function getOrCreateSession(sessionKey: string): Promise<AgentSession> {
    const existing = sessions.get(sessionKey);
    if (existing) return existing;

    const pending = pendingSessionLoads.get(sessionKey);
    if (pending) return pending;

    const capturedGen = sessionGenerations.get(sessionKey) ?? 0;
    const sessionPromise = (async () => {
      const record = sessionStore.ensureSession(sessionKey);
      const sessionManager = SessionManager.open(record.sessionFile);
      const { session } = await createAgentSession({
        authStorage,
        cwd,
        agentDir,
        model: resolvedModel?.model,
        modelRegistry,
        thinkingLevel: config.thinkingLevel,
        customTools: [cozeWebSearchTool, cozeWebFetchTool],
        resourceLoader,
        settingsManager,
        sessionManager
      });

      const currentGen = sessionGenerations.get(sessionKey) ?? 0;
      if (currentGen !== capturedGen) {
        // Session was reset while we were creating it; discard this instance.
        try {
          (session as unknown as { dispose?: () => void }).dispose?.();
        } catch {
          // ignore
        }
        throw new Error("stale session");
      }

      sessions.set(sessionKey, session);
      return session;
    })();

    pendingSessionLoads.set(sessionKey, sessionPromise);

    try {
      return await sessionPromise;
    } finally {
      pendingSessionLoads.delete(sessionKey);
    }
  }

  return {
    async run(message: BotMessage): Promise<string> {
      const sessionKey = getSessionKey(message);
      const session = await getOrCreateSession(sessionKey);
      if (session.isStreaming) {
        await session.prompt(message.text, { streamingBehavior: "followUp" });
      } else {
        await session.prompt(message.text);
      }
      return extractAssistantText(session);
    },
    async stream(message: BotMessage, handlers: PiAgentStreamHandlers) {
      const sessionKey = getSessionKey(message);
      handlers.onMeta?.({ sessionKey });

      const session = await (async () => {
        try {
          return await getOrCreateSession(sessionKey);
        } catch (err) {
          if (String(err).includes("stale session")) {
            // Race with reset; retry once against the latest transcript.
            return await getOrCreateSession(sessionKey);
          }
          throw err;
        }
      })();
      if (session.isStreaming) {
        // Dashboard UI expects a single in-flight stream per session.
        throw new Error(`Session is busy: ${sessionKey}`);
      }

      let lastText = "";
      let sawTextDelta = false;

      const unsubscribe = session.subscribe((event) => {
        handlers.onEvent?.(event);

        // Primary streaming signal in pi-coding-agent is `assistantMessageEvent.text_delta`.
        // Keep a fallback for older/alternate event shapes.
        if (!event || typeof event !== "object") return;
        const type = (event as { type?: unknown }).type;

        if (type === "message_update") {
          const assistantMessageEvent = (event as { assistantMessageEvent?: unknown }).assistantMessageEvent;
          if (assistantMessageEvent && typeof assistantMessageEvent === "object") {
            const evtType = (assistantMessageEvent as { type?: unknown }).type;
            const delta = (assistantMessageEvent as { delta?: unknown }).delta;
            if (evtType === "text_delta" && typeof delta === "string" && delta) {
              sawTextDelta = true;
              handlers.onDelta?.(delta);
              return;
            }
          }
        }

        // Fallback: stream by diffing assistant full text snapshots.
        // IMPORTANT: once we have real text_delta events, do not mix snapshot-based
        // fallback deltas into the same output stream, otherwise we may duplicate
        // or corrupt the accumulated text.
        if (sawTextDelta) return;
        if (type !== "message_start" && type !== "message_update" && type !== "message_end") return;

        const msg = (event as { message?: unknown }).message;
        if (!msg || typeof msg !== "object") return;
        if ((msg as { role?: unknown }).role !== "assistant") return;

        const fullText = extractAssistantTextFromMessage(msg);
        if (fullText.length < lastText.length || !fullText.startsWith(lastText)) {
          lastText = fullText;
          if (fullText) {
            console.log("[pi-bot][agent-stream] fallback_fulltext", fullText);
            handlers.onDelta?.(fullText);
          }
          return;
        }
        const delta = fullText.slice(lastText.length);
        lastText = fullText;
        if (delta) {
          handlers.onDelta?.(delta);
        }
      });

      try {
        await session.prompt(message.text);
        const finalText = extractAssistantText(session);
        return { sessionKey, finalText };
      } catch (err) {
        handlers.onError?.(String(err));
        throw err;
      } finally {
        unsubscribe();
      }
    },
    getSessionIfExists(sessionKey: string): AgentSession | undefined {
      return sessions.get(sessionKey);
    },
    listSessionKeys(): string[] {
      const all = new Set<string>([...sessionStore.listSessionKeys(), ...sessions.keys()]);
      return Array.from(all);
    },
    async ensureSessionLoaded(sessionKey: string): Promise<AgentSession | undefined> {
      try {
        return await getOrCreateSession(sessionKey);
      } catch (err) {
        if (String(err).includes("stale session")) {
          return await getOrCreateSession(sessionKey);
        }
        throw err;
      }
    },
    async abortSession(sessionKey: string): Promise<void> {
      const session = sessions.get(sessionKey);
      if (!session) return;
      await session.abort();
    },
    async resetSession(sessionKey: string): Promise<void> {
      const nextGen = (sessionGenerations.get(sessionKey) ?? 0) + 1;
      sessionGenerations.set(sessionKey, nextGen);
      sessionStore.resetSession(sessionKey);
      const session = sessions.get(sessionKey);
      if (!session) return;
      try {
        if (session.isStreaming) {
          await session.abort();
        }
      } catch {
        // ignore
      }
      try {
        (session as unknown as { dispose?: () => void }).dispose?.();
      } catch {
        // ignore
      }
      sessions.delete(sessionKey);
    },
    async dispose(): Promise<void> {
      sessions.clear();
    }
  } satisfies PiAgentRuntime;
}

export async function createAgentRuntime(config: AgentRuntimeConfig): Promise<PiAgentRuntime> {
  if (config.mode === "mock") {
    return createMockAgentRuntime();
  }

  return createPiAgentRuntime(config);
}