import { SessionManager } from "@earendil-works/pi-coding-agent";
import type { ExtensionAPI, ExtensionContext } from "@earendil-works/pi-coding-agent";
import { Text } from "@earendil-works/pi-tui";
import { readFileSync } from "node:fs";

const CUSTOM_TYPE = "pi-played";
const DATA_VERSION = 1;
const DAY_MS = 24 * 60 * 60 * 1000;
const MINUTE_MS = 60 * 1000;
const IDLE_GAP_MS = 2 * 60 * 60 * 1000;

type JsonObject = Record<string, unknown>;

interface Interval {
  start: number;
  end: number;
}

interface SessionMetrics {
  path: string;
  tokens: number;
  cost: number;
  recordedIntervals: Interval[];
  estimatedIntervals: Interval[];
}

type PlayedData =
  | {
    version: 1;
    kind: "span_start";
    start: number;
    reason: string;
    sessionFile?: string;
  }
  | {
    version: 1;
    kind: "span";
    start: number;
    end: number;
    reason: string;
    sessionFile?: string;
  }
  | {
    version: 1;
    kind: "report";
    report: string;
    timestamp: number;
  };

function asObject(value: unknown): JsonObject | undefined {
  return value && typeof value === "object" && !Array.isArray(value) ? (value as JsonObject) : undefined;
}

function finiteNumber(value: unknown): number | undefined {
  return typeof value === "number" && Number.isFinite(value) ? value : undefined;
}

function parseIsoTime(value: unknown): number | undefined {
  if (typeof value !== "string") return undefined;
  const time = new Date(value).getTime();
  return Number.isFinite(time) ? time : undefined;
}

function getEntryTime(entry: JsonObject): number | undefined {
  const message = asObject(entry.message);
  const messageTime = finiteNumber(message?.timestamp);
  return messageTime ?? parseIsoTime(entry.timestamp);
}

function readSessionFile(filePath: string): JsonObject[] {
  let content: string;
  try {
    content = readFileSync(filePath, "utf8");
  } catch {
    return [];
  }

  // Parse line by line so one corrupt line (e.g. truncated by a crash) does
  // not drop the whole session.
  return content.split(/\r?\n/).flatMap((line) => {
    if (line.trim().length === 0) return [];
    try {
      return [JSON.parse(line) as JsonObject];
    } catch {
      return [];
    }
  });
}

function entriesFromSessionManager(ctx: ExtensionContext): JsonObject[] {
  return ctx.sessionManager.getEntries().map((entry) => entry as unknown as JsonObject);
}

function getAssistantUsage(message: JsonObject): { tokens: number; cost: number } {
  if (message.role !== "assistant") return { tokens: 0, cost: 0 };

  const usage = asObject(message.usage);
  if (!usage) return { tokens: 0, cost: 0 };

  return {
    tokens:
      (finiteNumber(usage.input) ?? 0) +
      (finiteNumber(usage.output) ?? 0) +
      (finiteNumber(usage.cacheRead) ?? 0) +
      (finiteNumber(usage.cacheWrite) ?? 0),
    cost: finiteNumber(asObject(usage.cost)?.total) ?? 0,
  };
}

function getMessageTimesAndUsage(entries: JsonObject[]): { messageTimes: number[]; tokens: number; cost: number } {
  const messageTimes: number[] = [];
  let tokens = 0;
  let cost = 0;

  for (const entry of entries) {
    if (entry.type !== "message") continue;

    const message = asObject(entry.message);
    if (!message) continue;

    const usage = getAssistantUsage(message);
    tokens += usage.tokens;
    cost += usage.cost;

    const time = getEntryTime(entry);
    if (time !== undefined) messageTimes.push(time);
  }

  messageTimes.sort((a, b) => a - b);
  return { messageTimes, tokens, cost };
}

function getRecordedIntervals(entries: JsonObject[], messageTimes: number[]): Interval[] {
  const intervals: Interval[] = [];
  const openStarts: number[] = [];
  const closedStarts = new Set<number>();

  for (const entry of entries) {
    if (entry.type !== "custom" || entry.customType !== CUSTOM_TYPE) continue;

    const data = asObject(entry.data);
    if (!data || data.version !== DATA_VERSION) continue;

    if (data.kind === "span") {
      const start = finiteNumber(data.start);
      const end = finiteNumber(data.end);
      if (start === undefined || end === undefined || end <= start) continue;

      intervals.push({ start, end });
      closedStarts.add(start);
    } else if (data.kind === "span_start") {
      const start = finiteNumber(data.start);
      if (start !== undefined) openStarts.push(start);
    }
  }

  // A span_start without a matching span means the session did not shut down
  // cleanly (crash, killed terminal). Recover the span using the session's
  // last message activity.
  const lastMessageTime = messageTimes[messageTimes.length - 1];
  for (const start of openStarts) {
    if (closedStarts.has(start)) continue;
    if (lastMessageTime !== undefined && lastMessageTime > start) {
      intervals.push({ start, end: lastMessageTime });
    }
  }

  return intervals;
}

function estimateIntervalsFromMessages(messageTimes: number[]): Interval[] {
  const intervals: Interval[] = [];
  const uniqueTimes = [...new Set(messageTimes)];

  for (let i = 1; i < uniqueTimes.length; i++) {
    const start = uniqueTimes[i - 1];
    const end = uniqueTimes[i];
    const gap = end - start;

    if (gap > 0 && gap <= IDLE_GAP_MS) {
      intervals.push({ start, end });
    }
  }

  return intervals;
}

function mergeIntervals(intervals: Interval[]): Interval[] {
  const sorted = intervals
    .filter((interval) => interval.end > interval.start)
    .sort((a, b) => a.start - b.start || a.end - b.end);

  const merged: Interval[] = [];
  for (const interval of sorted) {
    const previous = merged[merged.length - 1];
    if (!previous || interval.start > previous.end) {
      merged.push({ ...interval });
    } else {
      previous.end = Math.max(previous.end, interval.end);
    }
  }

  return merged;
}

function sumIntervals(intervals: Interval[]): number {
  return intervals.reduce((sum, interval) => sum + Math.max(0, interval.end - interval.start), 0);
}

function currentOpenInterval(currentStart: number | undefined, now: number): Interval[] {
  if (currentStart === undefined || now <= currentStart) return [];
  return [{ start: currentStart, end: now }];
}

function sessionMetrics(path: string, entries: JsonObject[]): SessionMetrics {
  const { messageTimes, tokens, cost } = getMessageTimesAndUsage(entries);
  return {
    path,
    tokens,
    cost,
    recordedIntervals: getRecordedIntervals(entries, messageTimes),
    estimatedIntervals: estimateIntervalsFromMessages(messageTimes),
  };
}

async function collectAllSessionMetrics(currentSessionFile: string | undefined): Promise<SessionMetrics[]> {
  const sessions = await SessionManager.listAll();
  const paths = new Set(sessions.map((session) => session.path));
  if (currentSessionFile) paths.add(currentSessionFile);

  return [...paths].map((path) => sessionMetrics(path, readSessionFile(path)));
}

function formatDuration(ms: number): string {
  const totalMinutes = Math.max(0, Math.floor(ms / MINUTE_MS));
  const days = Math.floor(totalMinutes / (24 * 60));
  const hours = Math.floor((totalMinutes % (24 * 60)) / 60);
  const minutes = totalMinutes % 60;

  const parts: string[] = [];
  if (days > 0) parts.push(`${days} ${days === 1 ? "day" : "days"}`);
  if (hours > 0 || days > 0) parts.push(`${hours} ${hours === 1 ? "hour" : "hours"}`);
  parts.push(`${minutes} ${minutes === 1 ? "minute" : "minutes"}`);

  return parts.join(" ");
}

function formatTokens(tokens: number): string {
  return Math.round(tokens).toLocaleString();
}

function formatCost(cost: number): string {
  return `$${cost.toFixed(2)}`;
}

function formatTokenCost(tokens: number, cost: number): string {
  return `${formatTokens(tokens)} (${formatCost(cost)})`;
}

function formatReport(options: {
  totalMs: number;
  sessionMs: number;
  totalTokens: number;
  sessionTokens: number;
  totalCost: number;
  sessionCost: number;
  includesEstimates: boolean;
  sessionEstimated: boolean;
}): string {
  const estimateNote = options.includesEstimates
    ? " (est)"
    : "";
  const sessionEstimateNote = options.sessionEstimated ? " (est)" : "";

  return [
    `Total time played: ${formatDuration(options.totalMs)}${estimateNote}`,
    `Total tokens: ${formatTokenCost(options.totalTokens, options.totalCost)}`,
    `Time played this session: ${formatDuration(options.sessionMs)}${sessionEstimateNote}`,
    `Tokens this session: ${formatTokenCost(options.sessionTokens, options.sessionCost)}`,
  ].join("\n");
}

export default function playedExtension(pi: ExtensionAPI) {
  let currentStart: number | undefined;
  let currentSessionFile: string | undefined;

  pi.registerEntryRenderer<PlayedData>(CUSTOM_TYPE, (entry, _options, theme) => {
    const data = entry.data;
    if (!data || data.kind !== "report") return undefined;

    return new Text(theme.fg("thinkingHigh", data.report), 1, 0);
  });

  pi.on("session_start", (event, ctx) => {
    currentStart = Date.now();
    currentSessionFile = ctx.sessionManager.getSessionFile();

    const data: PlayedData = {
      version: DATA_VERSION,
      kind: "span_start",
      start: currentStart,
      reason: event.reason,
      sessionFile: currentSessionFile,
    };
    pi.appendEntry(CUSTOM_TYPE, data);
  });

  pi.on("session_shutdown", (event, ctx) => {
    if (currentStart !== undefined) {
      const end = Date.now();
      const data: PlayedData = {
        version: DATA_VERSION,
        kind: "span",
        start: currentStart,
        end,
        reason: event.reason,
        sessionFile: currentSessionFile ?? ctx.sessionManager.getSessionFile(),
      };
      pi.appendEntry(CUSTOM_TYPE, data);
    }

    currentStart = undefined;
    currentSessionFile = undefined;
  });

  pi.registerCommand("played", {
    description: "World of Warcraft style /played time and token totals.",
    handler: async (_args, ctx) => {
      const now = Date.now();
      const activeSessionFile = ctx.sessionManager.getSessionFile() ?? currentSessionFile;
      const allMetrics = await collectAllSessionMetrics(activeSessionFile);
      const currentEntries = entriesFromSessionManager(ctx);
      const currentMetrics = sessionMetrics(activeSessionFile ?? "<in-memory>", currentEntries);

      const ongoingIntervals = currentOpenInterval(currentStart, now);
      const recordedIntervals = allMetrics.flatMap((metrics) => metrics.recordedIntervals).concat(ongoingIntervals);
      const recordedTotalMs = sumIntervals(mergeIntervals(recordedIntervals));
      const includesEstimates = recordedTotalMs < DAY_MS;
      const totalIntervals = includesEstimates
        ? recordedIntervals.concat(allMetrics.flatMap((metrics) => metrics.estimatedIntervals))
        : recordedIntervals;

      const currentRecordedIntervals = currentMetrics.recordedIntervals.concat(ongoingIntervals);
      const hasCurrentRecordedSpan = currentRecordedIntervals.length > 0;
      const currentTimeIntervals = hasCurrentRecordedSpan
        ? currentRecordedIntervals
        : currentMetrics.estimatedIntervals;

      const totalTokens = allMetrics.reduce((sum, metrics) => sum + metrics.tokens, 0);
      const sessionTokens = currentMetrics.tokens;
      const totalCost = allMetrics.reduce((sum, metrics) => sum + metrics.cost, 0);
      const sessionCost = currentMetrics.cost;

      const report = formatReport({
        totalMs: sumIntervals(mergeIntervals(totalIntervals)),
        sessionMs: sumIntervals(mergeIntervals(currentTimeIntervals)),
        totalTokens,
        sessionTokens,
        totalCost,
        sessionCost,
        includesEstimates,
        sessionEstimated: !hasCurrentRecordedSpan,
      });

      if (ctx.mode === "tui") {
        pi.appendEntry<PlayedData>(CUSTOM_TYPE, {
          version: DATA_VERSION,
          kind: "report",
          report,
          timestamp: now,
        });
      } else if (ctx.hasUI) {
        ctx.ui.notify(report, "info");
      } else {
        console.log(report);
      }
    },
  });
}