import { homedir } from "node:os";
import { StringDecoder } from "node:string_decoder";
import type { Usage } from "@earendil-works/pi-ai";
import type { ChildUsage, DedeActivity, DetailedUsage } from "./types.ts";

const MAX_JSON_LINE_BYTES = 2 * 1024 * 1024;
const MAX_PARTIAL_TEXT_BYTES = 32 * 1024;
const MAX_ACTIVITY = 100;

function zeroDetailedUsage(): DetailedUsage {
  return {
    input: 0,
    output: 0,
    cacheRead: 0,
    cacheWrite: 0,
    totalTokens: 0,
    cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
  };
}

export function sumUsage(target: DetailedUsage, usage: Partial<Usage> | undefined): void {
  if (!usage || typeof usage !== "object" || Array.isArray(usage)) return;
  const finite = (value: unknown): number => typeof value === "number" && Number.isFinite(value) && value >= 0 ? value : 0;
  usage = {
    input: finite(usage.input), output: finite(usage.output),
    cacheRead: finite(usage.cacheRead), cacheWrite: finite(usage.cacheWrite),
    totalTokens: finite(usage.totalTokens),
    cost: { input: finite(usage.cost?.input), output: finite(usage.cost?.output),
      cacheRead: finite(usage.cost?.cacheRead), cacheWrite: finite(usage.cost?.cacheWrite), total: finite(usage.cost?.total) },
  };
  target.input += usage.input ?? 0;
  target.output += usage.output ?? 0;
  target.cacheRead += usage.cacheRead ?? 0;
  target.cacheWrite += usage.cacheWrite ?? 0;
  target.totalTokens += usage.totalTokens ?? 0;
  target.cost.input += usage.cost?.input ?? 0;
  target.cost.output += usage.cost?.output ?? 0;
  target.cost.cacheRead += usage.cost?.cacheRead ?? 0;
  target.cost.cacheWrite += usage.cost?.cacheWrite ?? 0;
  target.cost.total += usage.cost?.total ?? 0;
}

function object(value: unknown): Record<string, any> | undefined {
  return value !== null && typeof value === "object" && !Array.isArray(value) ? value as Record<string, any> : undefined;
}

function textFromAssistant(message: Record<string, any>): string {
  if (!Array.isArray(message.content)) return "";
  return message.content
    .filter((part: unknown) => object(part)?.type === "text")
    .map((part: unknown) => typeof object(part)?.text === "string" ? object(part)!.text : "")
    .join("\n");
}

function messageIdentity(message: Record<string, any>): string {
  return JSON.stringify([message.id, message.responseId, message.provider, message.model,
    message.timestamp, message.stopReason, message.usage?.input, message.usage?.output]);
}

function short(value: unknown, max = 120): string {
  const raw = typeof value === "string" ? value : JSON.stringify(value ?? "");
  return raw.length > max ? `${raw.slice(0, max - 1)}…` : raw;
}

function displayPath(value: unknown): string {
  const raw = typeof value === "string" ? value : "";
  const home = homedir();
  return raw.startsWith(home) ? `~${raw.slice(home.length)}` : raw;
}

export interface CollectedProtocol {
  finalText: string;
  model?: string;
  stopReason?: string;
  errorMessage?: string;
  usage: DetailedUsage;
  turns: number;
  activity: DedeActivity[];
  sawAgentEnd: boolean;
  malformedLines: number;
  oversizedLines: number;
}

/** Bounded JSONL collector for Pi's event stream (json or rpc mode).
 *
 * `onProgress` receives human-readable activity hints; `onEvent` receives every
 * successfully-parsed event object, including protocol-control messages the
 * collector itself ignores (e.g. rpc `response`, `agent_settled`,
 * `extension_ui_request`). The collector keeps the authoritative state; callers
 * use `onEvent` only for transport control. */
export class PiJsonCollector {
  private readonly decoder = new StringDecoder("utf8");
  private buffer = "";
  private partialText = "";
  private discardingOversizedLine = false;
  private readonly seenAssistantMessages = new Set<string>();
  private readonly state: CollectedProtocol = {
    finalText: "",
    usage: zeroDetailedUsage(),
    turns: 0,
    activity: [],
    sawAgentEnd: false,
    malformedLines: 0,
    oversizedLines: 0,
  };

  constructor(
    private readonly onProgress?: (text: string) => void,
    private readonly onEvent?: (event: Record<string, any>) => void,
  ) {}

  push(chunk: Buffer | string): void {
    this.consume(typeof chunk === "string" ? chunk : this.decoder.write(chunk));
  }

  end(): CollectedProtocol {
    this.consume(this.decoder.end());
    if (!this.discardingOversizedLine && this.buffer.trim()) this.processLine(this.buffer);
    this.buffer = "";
    if (this.partialText.trim()) {
      this.state.finalText = `${this.partialText}\n\n[Partial response interrupted before finalization.]`;
    }
    return {
      ...this.state,
      usage: { ...this.state.usage, cost: { ...this.state.usage.cost } },
      activity: [...this.state.activity],
    };
  }

  snapshot(): CollectedProtocol {
    return {
      ...this.state,
      usage: { ...this.state.usage, cost: { ...this.state.usage.cost } },
      activity: [...this.state.activity],
    };
  }

  private consume(text: string): void {
    let remaining = text;
    while (remaining.length > 0) {
      if (this.discardingOversizedLine) {
        const newline = remaining.indexOf("\n");
        if (newline < 0) return;
        this.discardingOversizedLine = false;
        remaining = remaining.slice(newline + 1);
        continue;
      }

      const newline = remaining.indexOf("\n");
      if (newline >= 0) {
        const line = this.buffer + remaining.slice(0, newline).replace(/\r$/, "");
        this.buffer = "";
        remaining = remaining.slice(newline + 1);
        if (Buffer.byteLength(line, "utf8") > MAX_JSON_LINE_BYTES) {
          this.oversized();
        } else {
          this.processLine(line);
        }
        continue;
      }

      this.buffer += remaining;
      if (Buffer.byteLength(this.buffer, "utf8") > MAX_JSON_LINE_BYTES) {
        this.buffer = "";
        this.discardingOversizedLine = true;
        this.oversized();
      }
      return;
    }
  }

  private oversized(): void {
    this.state.oversizedLines++;
    this.addActivity("status", "ignored oversized protocol line");
  }

  private processLine(line: string): void {
    if (!line.trim()) return;
    let event: Record<string, any>;
    try {
      const parsed: unknown = JSON.parse(line);
      const record = object(parsed);
      if (!record || typeof record.type !== "string" || !record.type.trim()) throw new Error("Invalid protocol event");
      event = record;
    } catch {
      this.state.malformedLines++;
      if (this.state.malformedLines <= 3) this.addActivity("status", "ignored malformed protocol line");
      return;
    }

    const malformed = () => {
      this.state.malformedLines++;
      if (this.state.malformedLines <= 3) this.addActivity("status", "ignored malformed event fields");
    };
    if ((["message_start", "message_end"].includes(event.type) && !object(event.message)) ||
        (event.type === "message_update" && !object(event.assistantMessageEvent)) ||
        (event.type.startsWith("tool_execution_") && typeof event.toolName !== "string") ||
        (event.type === "response" && (typeof event.command !== "string" || typeof event.success !== "boolean")) ||
        (event.type === "extension_ui_request" && (typeof event.id !== "string" || typeof event.method !== "string"))) {
      malformed(); return;
    }
    if (event.type === "message_end") {
      const message = event.message;
      if (["role", "provider", "model", "stopReason", "errorMessage"].some((key) => message[key] !== undefined && typeof message[key] !== "string")) {
        malformed(); return;
      }
    }

    switch (event.type) {
      case "message_start": {
        const message = object(event.message);
        if (message?.role === "assistant") this.partialText = "";
        break;
      }
      case "message_update": {
        // Keep a bounded fallback so a timed-out child can show how close it was, but never
        // emit answer deltas as progress or trigger per-token renders.
        const delta = object(event.assistantMessageEvent);
        if (delta?.type === "text_delta" && typeof delta.delta === "string") {
          const combined = Buffer.from(this.partialText + delta.delta, "utf8");
          if (combined.length <= MAX_PARTIAL_TEXT_BYTES) {
            this.partialText += delta.delta;
          } else {
            let end = MAX_PARTIAL_TEXT_BYTES;
            while (end > 0 && (combined[end] & 0xc0) === 0x80) end--;
            this.partialText = combined.subarray(0, end).toString("utf8");
          }
        }
        break;
      }
      case "message_end": {
        const message = object(event.message);
        if (message?.role !== "assistant") break;
        const identity = messageIdentity(message);
        if (!this.seenAssistantMessages.has(identity)) {
          this.seenAssistantMessages.add(identity);
          this.state.turns++;
          sumUsage(this.state.usage, message.usage);
        }
        const text = textFromAssistant(message);
        // An empty abort/error must not erase the evidence from the last completed response.
        if (text || (message.stopReason !== "aborted" && message.stopReason !== "error")) this.state.finalText = text;
        this.state.model = message.model ? `${message.provider ? `${message.provider}/` : ""}${message.model}` : this.state.model;
        if (typeof message.stopReason === "string") this.state.stopReason = message.stopReason;
        if (typeof message.errorMessage === "string") this.state.errorMessage = message.errorMessage;
        this.partialText = "";
        this.onProgress?.("responded");
        break;
      }
      case "tool_execution_start": {
        const name = String(event.toolName ?? "tool");
        const args = object(event.args) ?? {};
        let description = name;
        if (name === "read") description = `reading ${short(displayPath(args.path ?? args.file_path ?? "file"), 100)}`;
        else if (name === "grep") description = `grep /${short(args.pattern ?? "", 70)}/`;
        else if (name === "find") description = `find ${short(args.pattern ?? "*", 80)}`;
        else if (name === "ls") description = `listing ${short(displayPath(args.path ?? "."), 100)}`;
        else if (name === "bash") description = `$ ${short(args.command ?? "", 100)}`;
        else if (name === "edit" || name === "write") description = `${name} ${short(displayPath(args.path ?? args.file_path ?? "file"), 100)}`;
        this.addActivity("tool", description);
        this.onProgress?.(description);
        break;
      }
      case "tool_execution_update":
        // Tool output can be large or sensitive. Start/end events provide enough progress.
        break;
      case "tool_execution_end":
        this.onProgress?.(`${String(event.toolName ?? "tool")} finished`);
        break;
      case "agent_end":
        this.state.sawAgentEnd = true;
        this.addActivity("status", "agent finished");
        break;
      case "auto_retry_start":
      case "retry_start":
        this.addActivity("status", "model retrying");
        this.onProgress?.("model retrying");
        break;
      case "auto_compaction_start":
      case "compaction_start":
        this.addActivity("status", "context compacting");
        this.onProgress?.("context compacting");
        break;
    }

    // Notify transport observers after state is updated. Ignored event types
    // (rpc `response`, `agent_settled`, `extension_ui_request`, `queue_update`,
    // ...) still pass through so the RPC controller can react to them.
    this.onEvent?.(event);
  }

  private addActivity(type: DedeActivity["type"], text: string): void {
    if (!text) return;
    if (this.state.activity.length >= MAX_ACTIVITY) this.state.activity.shift();
    this.state.activity.push({ type, text: short(text, 240) });
  }
}

export function childUsage(protocol: CollectedProtocol): ChildUsage {
  return {
    input: protocol.usage.input,
    output: protocol.usage.output,
    cacheRead: protocol.usage.cacheRead,
    cacheWrite: protocol.usage.cacheWrite,
    cost: protocol.usage.cost.total,
    totalTokens: protocol.usage.totalTokens,
    turns: protocol.turns,
  };
}

export function aggregateUsages(usages: DetailedUsage[]): DetailedUsage {
  const result = zeroDetailedUsage();
  for (const usage of usages) sumUsage(result, usage);
  return result;
}