import assert from "node:assert/strict";
import { randomUUID } from "node:crypto";
import { mkdtemp, mkdir, writeFile } from "node:fs/promises";
import os from "node:os";
import path from "node:path";
import test from "node:test";

import { scanUsage } from "./index.js";

async function scanClaudeRows(rows) {
  const root = await mkdtemp(path.join(os.tmpdir(), "tokenmaxx-usage-test-"));
  const projectDir = path.join(root, "project");
  await mkdir(projectDir, { recursive: true });
  await writeFile(
    path.join(projectDir, "session.jsonl"),
    rows.map((row) => JSON.stringify(row)).join("\n")
  );

  const summary = await scanUsage({
    claudeRoot: root,
    codexRoot: path.join(root, "empty-codex")
  });
  return summary;
}

async function scanCodexRows(rows) {
  const root = await mkdtemp(path.join(os.tmpdir(), "tokenmaxx-usage-test-"));
  await writeFile(
    path.join(root, "session.jsonl"),
    rows.map((row) => JSON.stringify(row)).join("\n")
  );

  const summary = await scanUsage({
    claudeRoot: path.join(root, "empty-claude"),
    codexRoot: root
  });
  return summary;
}

async function scanCodexSessions(sessions) {
  const root = await mkdtemp(path.join(os.tmpdir(), "tokenmaxx-usage-test-"));
  for (const [name, rows] of Object.entries(sessions)) {
    await writeFile(
      path.join(root, `${name}.jsonl`),
      rows.map((row) => JSON.stringify(row)).join("\n")
    );
  }

  return scanUsage({
    claudeRoot: path.join(root, "empty-claude"),
    codexRoot: root
  });
}

function localDayKey(value) {
  const date = new Date(value);
  const y = date.getFullYear();
  const m = String(date.getMonth() + 1).padStart(2, "0");
  const d = String(date.getDate()).padStart(2, "0");
  return `${y}-${m}-${d}`;
}

function assistantRow(model, usage, timestamp = "2026-07-03T12:00:00.000Z") {
  return {
    type: "assistant",
    timestamp,
    requestId: randomUUID(),
    sessionId: randomUUID(),
    cwd: "/tmp/tokenmaxx-test",
    message: {
      model,
      usage
    }
  };
}

function codexRows(model, usage) {
  const timestamp = "2026-07-10T12:00:00.000Z";
  return [
    {
      type: "session_meta",
      timestamp,
      payload: {
        id: randomUUID(),
        cwd: "/tmp/tokenmaxx-test",
        model,
        model_provider: "openai"
      }
    },
    {
      type: "event_msg",
      timestamp,
      payload: {
        type: "token_count",
        info: { total_token_usage: usage }
      }
    }
  ];
}

function codexTokenRow(usage, timestamp) {
  return {
    type: "event_msg",
    timestamp,
    payload: {
      type: "token_count",
      info: {
        total_token_usage: usage,
        last_token_usage: usage
      }
    }
  };
}

for (const [model, expectedCost] of [
  ["gpt-5.6", 7.1],
  ["gpt-5.6-sol", 7.1],
  ["gpt-5.6-terra", 3.55],
  ["gpt-5.6-luna", 1.42]
]) {
  test(`prices OpenAI ${model} fresh, cached, and output tokens`, async () => {
    const summary = await scanCodexRows(
      codexRows(model, {
        total_tokens: 1_100_000,
        input_tokens: 1_000_000,
        cached_input_tokens: 200_000,
        output_tokens: 100_000
      })
    );

    assert.equal(summary.totals.pricedTokens, 1_100_000);
    assert.equal(summary.totals.unpricedTokens, 0);
    assert.equal(summary.totals.costUsd, expectedCost);
  });
}

test("leaves unknown OpenAI models unpriced", async () => {
  const summary = await scanCodexRows(
    codexRows("gpt-not-real-1", {
      total_tokens: 2_000,
      input_tokens: 1_000,
      output_tokens: 1_000
    })
  );

  assert.equal(summary.totals.costUsd, 0);
  assert.equal(summary.totals.pricedTokens, 0);
  assert.equal(summary.totals.unpricedTokens, 2_000);
});

test("does not recount parent token history copied into a Codex fork", async () => {
  const parentId = randomUUID();
  const childId = randomUUID();
  const initial = {
    total_tokens: 100,
    input_tokens: 80,
    cached_input_tokens: 20,
    output_tokens: 20
  };
  const forkBaseline = {
    total_tokens: 250,
    input_tokens: 200,
    cached_input_tokens: 100,
    output_tokens: 50
  };
  const parentFinal = {
    total_tokens: 400,
    input_tokens: 320,
    cached_input_tokens: 180,
    output_tokens: 80
  };
  const childFinal = {
    total_tokens: 310,
    input_tokens: 250,
    cached_input_tokens: 130,
    output_tokens: 60
  };
  const parentMeta = {
    type: "session_meta",
    timestamp: "2026-07-10T12:00:00.000Z",
    payload: {
      id: parentId,
      cwd: "/tmp/tokenmaxx-test",
      model: "gpt-5.6-sol",
      model_provider: "openai"
    }
  };
  const childMeta = {
    type: "session_meta",
    timestamp: "2026-07-10T12:02:00.000Z",
    payload: {
      id: childId,
      forked_from_id: parentId,
      cwd: "/tmp/tokenmaxx-test",
      model: "gpt-5.6-sol",
      model_provider: "openai"
    }
  };

  const summary = await scanCodexSessions({
    parent: [
      parentMeta,
      codexTokenRow(initial, "2026-07-10T12:01:00.000Z"),
      codexTokenRow(forkBaseline, "2026-07-10T12:01:30.000Z"),
      codexTokenRow(parentFinal, "2026-07-10T12:03:00.000Z")
    ],
    child: [
      childMeta,
      parentMeta,
      codexTokenRow(initial, "2026-07-10T12:02:00.100Z"),
      codexTokenRow(forkBaseline, "2026-07-10T12:02:00.200Z"),
      codexTokenRow(childFinal, "2026-07-10T12:02:30.000Z")
    ]
  });

  assert.equal(summary.totals.totalTokens, 460);
  assert.equal(summary.totals.inputTokens, 370);
  assert.equal(summary.totals.cacheReadTokens, 210);
  assert.equal(summary.totals.outputTokens, 90);
  assert.equal(summary.sourceStats.codexEvents, 4);
  assert.equal(summary.sourceStats.codexForks, 1);
  assert.equal(summary.sourceStats.codexDedupedForks, 1);
  assert.equal(summary.sourceStats.codexFreshForks, 0);
  assert.equal(summary.sourceStats.codexUnresolvedForks, 0);

  const childSession = summary.sessions.find((session) => session.id === childId);
  assert.ok(childSession);
  assert.equal(childSession.totalTokens, 60);
  assert.equal(childSession.inputTokens, 50);
  assert.equal(childSession.cacheReadTokens, 30);
  assert.equal(childSession.outputTokens, 10);
});

test("counts a Codex fork once when its usage counter restarts", async () => {
  const parentId = randomUUID();
  const childId = randomUUID();
  const parentMeta = {
    type: "session_meta",
    timestamp: "2026-07-10T12:00:00.000Z",
    payload: {
      id: parentId,
      cwd: "/tmp/tokenmaxx-test",
      model: "gpt-5.6-sol",
      model_provider: "openai"
    }
  };
  const childMeta = {
    type: "session_meta",
    timestamp: "2026-07-10T12:02:00.000Z",
    payload: {
      id: childId,
      forked_from_id: parentId,
      cwd: "/tmp/tokenmaxx-test",
      model: "gpt-5.6-sol",
      model_provider: "openai"
    }
  };

  const summary = await scanCodexSessions({
    parent: [
      parentMeta,
      codexTokenRow(
        { total_tokens: 250, input_tokens: 200, output_tokens: 50 },
        "2026-07-10T12:01:00.000Z"
      ),
      codexTokenRow(
        { total_tokens: 300, input_tokens: 240, output_tokens: 60 },
        "2026-07-10T12:03:00.000Z"
      )
    ],
    child: [
      childMeta,
      codexTokenRow(
        { total_tokens: 10, input_tokens: 8, output_tokens: 2 },
        "2026-07-10T12:02:10.000Z"
      ),
      codexTokenRow(
        { total_tokens: 40, input_tokens: 32, output_tokens: 8 },
        "2026-07-10T12:02:30.000Z"
      )
    ]
  });

  assert.equal(summary.totals.totalTokens, 340);
  assert.equal(summary.sourceStats.codexForks, 1);
  assert.equal(summary.sourceStats.codexDedupedForks, 0);
  assert.equal(summary.sourceStats.codexFreshForks, 1);
  assert.equal(summary.sourceStats.codexUnresolvedForks, 0);
  assert.equal(
    summary.sessions.find((session) => session.id === childId)?.totalTokens,
    40
  );
});

test("prices Claude Fable 5 cache and token usage", async () => {
  const summary = await scanClaudeRows([
    assistantRow("claude-fable-5", {
      input_tokens: 1_000,
      output_tokens: 2_000,
      cache_read_input_tokens: 3_000,
      cache_creation: {
        ephemeral_5m_input_tokens: 4_000,
        ephemeral_1h_input_tokens: 5_000
      }
    })
  ]);

  assert.equal(summary.totals.totalTokens, 15_000);
  assert.equal(summary.totals.unpricedTokens, 0);
  assert.equal(summary.totals.pricedEvents, 1);
  assert.equal(summary.totals.costUsd, 0.263);
});

test("prices Claude Opus 4.8 at current Opus rates", async () => {
  const summary = await scanClaudeRows([
    assistantRow("claude-opus-4-8", {
      input_tokens: 1_000_000,
      output_tokens: 1_000_000
    })
  ]);

  assert.equal(summary.totals.unpricedTokens, 0);
  assert.equal(summary.totals.costUsd, 30);
});

test("leaves unknown Claude models unpriced", async () => {
  const summary = await scanClaudeRows([
    assistantRow("claude-not-real-1", {
      input_tokens: 1_000,
      output_tokens: 1_000
    })
  ]);

  assert.equal(summary.totals.costUsd, 0);
  assert.equal(summary.totals.pricedTokens, 0);
  assert.equal(summary.totals.unpricedTokens, 2_000);
});

test("rolls up contribution days by model", async () => {
  const timestamp = "2026-07-03T12:00:00.000Z";
  const day = localDayKey(timestamp);
  const summary = await scanClaudeRows([
    assistantRow("claude-fable-5", { input_tokens: 1_000, output_tokens: 500 }, timestamp),
    assistantRow("claude-fable-5", { input_tokens: 500, output_tokens: 500 }, timestamp),
    assistantRow("claude-sonnet-5", { input_tokens: 700 }, timestamp)
  ]);

  const fable = summary.modelContributionDays.find(
    (row) => row.date === day && row.model === "claude-fable-5"
  );
  const sonnet = summary.modelContributionDays.find(
    (row) => row.date === day && row.model === "claude-sonnet-5"
  );

  assert.ok(fable);
  assert.ok(sonnet);
  assert.equal(fable.provider, "anthropic");
  assert.equal(fable.tool, "claude");
  assert.equal(fable.totalTokens, 2_500);
  assert.equal(fable.inputTokens, 1_500);
  assert.equal(fable.outputTokens, 1_000);
  assert.equal(fable.events, 2);
  assert.equal(sonnet.totalTokens, 700);
  assert.equal(
    summary.contributionDays.find((row) => row.date === day).totalTokens,
    3_200
  );
});
