import { createHash, randomUUID } from "node:crypto";
import { appendFile, mkdir, readFile, readdir, rm, stat, writeFile } from "node:fs/promises";
import { join } from "node:path";

export interface TextContent {
  type: "text";
  text: string;
}

export interface ToolResultInput {
  sessionId: string;
  toolCallId: string;
  toolName: string;
  input: Record<string, unknown>;
  content: TextContent[];
  details?: unknown;
  isError: boolean;
}

export interface ProcessingTelemetry {
  estimatedTokensBefore: number;
  estimatedTokensAfter: number;
  estimatedTokensSaved: number;
}

export interface ProcessedToolResult {
  changed: boolean;
  artifactId?: string;
  content: TextContent[];
  telemetry: ProcessingTelemetry;
}

export interface ProviderUsageSummary {
  requests: number;
  inputTokens: number;
  cacheReadTokens: number;
  cacheWriteTokens: number;
  outputTokens: number;
  cost: number;
}

export interface ProviderUsageInput {
  input: number;
  cacheRead: number;
  cacheWrite: number;
  output: number;
  cost: { total: number };
}

export interface ContextCostStatus {
  archivedResults: number;
  estimatedTokensBefore: number;
  estimatedTokensAfter: number;
  estimatedTokensSaved: number;
  provider?: ProviderUsageSummary;
}

export interface ContextCostAuditEvent {
  type: "tool-result-archived";
  timestamp: string;
  sessionId: string;
  toolCallId: string;
  toolName: string;
  artifactId: string;
  cacheImpact: "new-result-no-prefix-rewrite";
  estimatedTokensBefore: number;
  estimatedTokensAfter: number;
  estimatedTokensSaved: number;
}

export type ContextCostMode = "observe" | "safe" | "aggressive";

export interface ContextCostManager {
  processToolResult(input: ToolResultInput): Promise<ProcessedToolResult>;
  restore(input: { sessionId: string; artifactId: string }): Promise<string>;
  getAudit(input: { sessionId: string }): Promise<ContextCostAuditEvent[]>;
  purge(input: { sessionId: string }): Promise<void>;
  purgeExpired(input: { olderThan: Date }): Promise<number>;
  getStatus(): ContextCostStatus;
  recordUsage(usage: ProviderUsageInput): void;
  setMode(mode: ContextCostMode): void;
}

export interface ContextCostManagerConfig {
  artifactRoot: string;
  mode: ContextCostMode;
  initialResultTokenThreshold: number;
  previewTokenLimit: number;
}

const CHARACTERS_PER_TOKEN = 4;

function estimateTokens(text: string): number {
  return Math.ceil(text.length / CHARACTERS_PER_TOKEN);
}

function sessionDirectoryName(sessionId: string): string {
  return createHash("sha256").update(sessionId).digest("hex").slice(0, 24);
}

function sessionDirectory(artifactRoot: string, sessionId: string): string {
  return join(artifactRoot, sessionDirectoryName(sessionId));
}

function artifactPath(
  artifactRoot: string,
  sessionId: string,
  artifactId: string,
): string {
  if (!/^[0-9a-f-]{36}$/.test(artifactId)) {
    throw new Error("Invalid artifact ID");
  }
  return join(sessionDirectory(artifactRoot, sessionId), `${artifactId}.log`);
}

function buildPreview(
  source: string,
  previewTokenLimit: number,
): string {
  const characterLimit = Math.max(0, previewTokenLimit * CHARACTERS_PER_TOKEN);
  if (source.length <= characterLimit) return source;

  const headLength = Math.ceil(characterLimit / 2);
  const tailLength = Math.floor(characterLimit / 2);
  return `${source.slice(0, headLength)}\n… archived content omitted …\n${source.slice(-tailLength)}`;
}

function toolLabel(input: ToolResultInput, original: string): string {
  if (input.toolName === "read") {
    const path =
      typeof input.input.path === "string" ? input.input.path : "<unknown>";
    const offset =
      typeof input.input.offset === "number" ? input.input.offset : 1;
    const limit =
      typeof input.input.limit === "number" ? input.input.limit : undefined;
    const range = limit === undefined
      ? `from line ${offset}`
      : `lines ${offset}-${offset + limit - 1}`;
    return `path: ${path}\nrange: ${range}`;
  }

  if (input.toolName === "grep" || input.toolName === "find") {
    const pattern =
      typeof input.input.pattern === "string" ? input.input.pattern : "<unknown>";
    const root = typeof input.input.path === "string" ? input.input.path : ".";
    const resultCount = original.split("\n").filter((line) => line.length > 0).length;
    const lines = [`pattern: ${pattern}`, `root: ${root}`];
    if (input.toolName === "grep" && typeof input.input.glob === "string") {
      lines.push(`filter: ${input.input.glob}`);
    }
    lines.push(`results: ${resultCount}`);
    return lines.join("\n");
  }

  const command =
    typeof input.input.command === "string" ? input.input.command : "<unknown>";
  const details = input.details as { exitCode?: unknown } | undefined;
  const exitCode =
    typeof details?.exitCode === "number" ? String(details.exitCode) : "unknown";
  return `command: ${command}\nstatus: exit ${exitCode}`;
}

export function createContextCostManager(
  config: ContextCostManagerConfig,
): ContextCostManager {
  let mode = config.mode;
  const status: ContextCostStatus = {
    archivedResults: 0,
    estimatedTokensBefore: 0,
    estimatedTokensAfter: 0,
    estimatedTokensSaved: 0,
  };

  return {
    async processToolResult(input) {
      const original = input.content.map((block) => block.text).join("\n");
      const estimatedTokensBefore = estimateTokens(original);

      if (
        input.isError ||
        !["bash", "read", "grep", "find"].includes(input.toolName) ||
        estimatedTokensBefore <= config.initialResultTokenThreshold
      ) {
        return {
          changed: false,
          content: input.content,
          telemetry: {
            estimatedTokensBefore,
            estimatedTokensAfter: estimatedTokensBefore,
            estimatedTokensSaved: 0,
          },
        };
      }

      const artifactId = randomUUID();
      const preview = buildPreview(original, config.previewTokenLimit);
      const replacement = [
        `[${input.toolName} output archived]`,
        toolLabel(input, original),
        "preview:",
        preview,
        `full output: cost-artifact://${artifactId}`,
      ].join("\n");
      const estimatedTokensAfter = estimateTokens(replacement);
      const estimatedTokensSaved = Math.max(
        0,
        estimatedTokensBefore - estimatedTokensAfter,
      );

      status.estimatedTokensBefore += estimatedTokensBefore;
      status.estimatedTokensAfter += estimatedTokensAfter;
      status.estimatedTokensSaved += estimatedTokensSaved;

      if (mode === "observe") {
        return {
          changed: false,
          content: input.content,
          telemetry: {
            estimatedTokensBefore,
            estimatedTokensAfter,
            estimatedTokensSaved,
          },
        };
      }

      const path = artifactPath(config.artifactRoot, input.sessionId, artifactId);
      const directory = sessionDirectory(config.artifactRoot, input.sessionId);
      await mkdir(directory, { recursive: true, mode: 0o700 });
      await writeFile(path, original, { encoding: "utf8", mode: 0o600 });
      const auditEvent: ContextCostAuditEvent = {
        type: "tool-result-archived",
        timestamp: new Date().toISOString(),
        sessionId: input.sessionId,
        toolCallId: input.toolCallId,
        toolName: input.toolName,
        artifactId,
        cacheImpact: "new-result-no-prefix-rewrite",
        estimatedTokensBefore,
        estimatedTokensAfter,
        estimatedTokensSaved,
      };
      await appendFile(join(directory, "audit.jsonl"), `${JSON.stringify(auditEvent)}\n`, {
        encoding: "utf8",
        mode: 0o600,
      });
      status.archivedResults += 1;

      return {
        changed: true,
        artifactId,
        content: [{ type: "text", text: replacement }],
        telemetry: {
          estimatedTokensBefore,
          estimatedTokensAfter,
          estimatedTokensSaved,
        },
      };
    },

    restore({ sessionId, artifactId }) {
      return readFile(artifactPath(config.artifactRoot, sessionId, artifactId),
        "utf8");
    },

    async getAudit({ sessionId }) {
      const contents = await readFile(
        join(sessionDirectory(config.artifactRoot, sessionId), "audit.jsonl"),
        "utf8",
      );
      return contents
        .split("\n")
        .filter(Boolean)
        .map((line) => JSON.parse(line) as ContextCostAuditEvent);
    },

    async purge({ sessionId }) {
      await rm(sessionDirectory(config.artifactRoot, sessionId), {
        recursive: true,
        force: true,
      });
    },

    async purgeExpired({ olderThan }) {
      let entries;
      try {
        entries = await readdir(config.artifactRoot, { withFileTypes: true });
      } catch {
        return 0;
      }
      let removed = 0;
      for (const entry of entries) {
        if (!entry.isDirectory()) continue;
        const directory = join(config.artifactRoot, entry.name);
        const metadata = await stat(directory);
        if (metadata.mtime >= olderThan) continue;
        await rm(directory, { recursive: true, force: true });
        removed += 1;
      }
      return removed;
    },

    getStatus() {
      return status.provider
        ? { ...status, provider: { ...status.provider } }
        : { ...status };
    },

    recordUsage(usage) {
      const provider = status.provider ?? {
        requests: 0,
        inputTokens: 0,
        cacheReadTokens: 0,
        cacheWriteTokens: 0,
        outputTokens: 0,
        cost: 0,
      };
      provider.requests += 1;
      provider.inputTokens += usage.input;
      provider.cacheReadTokens += usage.cacheRead;
      provider.cacheWriteTokens += usage.cacheWrite;
      provider.outputTokens += usage.output;
      provider.cost += usage.cost.total;
      status.provider = provider;
    },

    setMode(nextMode) {
      mode = nextMode;
    },
  };
}