import { existsSync, mkdirSync, readFileSync, readdirSync, statSync, writeFileSync } from "node:fs";
import type { ProviderId } from "./types.js";
import { homedir } from "node:os";
import { join } from "node:path";

interface LocalProviderStats {
  activeDates: string[];
  activeDays: number;
  hasLocalStats: boolean;
  hasPromptStats: boolean;
  modelUsage: Record<
    string,
    {
      cacheCreationInputTokens: number;
      cacheReadInputTokens: number;
      inputTokens: number;
      outputTokens: number;
    }
  >;
  recentDays: {
    date: string;
    messageCount: number;
  }[];
  todayPrompts: number;
  todaySessions: number;
  todayTokensByModel: Record<string, number>;
  todayTotalTokens: number;
  totalPrompts: number;
  totalSessions: number;
}

interface FileStatsEntry {
  mtime_ns: string | number;
  size: number;
  stats: Record<
    string,
    {
      model_usage?: Record<
        string,
        {
          cr: number;
          cw: number;
          inp: number;
          out: number;
        }
      >;
      models_by_day?: Record<
        string,
        Record<
          string,
          {
            cr: number;
            cw: number;
            inp: number;
            out: number;
            tot: number;
          }
        >
      >;
      prompts_by_day?: Record<string, number>;
      tokens_by_day?: Record<string, number>;
      total_prompts?: number;
    }
  >;
}

const CACHE_DIR = join(
  process.env.XDG_CACHE_HOME || join(homedir(), ".cache"),
  "pi-usage",
  "sessions",
);
const CACHE_FILE = join(CACHE_DIR, "session-tokens-cache.json");

// Strictly map official provider channels only (no unofficial proxies)
const OFFICIAL_PROVIDERS: Record<string, ProviderId> = {
  antigravity: "antigravity",
  codex: "openai-codex",
  deepseek: "deepseek",
  grok: "xai",
  kimi: "kimi-coding",
  "kimi-coding": "kimi-coding",
  openai: "openai-codex",
  "openai-codex": "openai-codex",
  opencode: "opencode-go",
  "opencode-go": "opencode-go",
  openrouter: "openrouter",
  xai: "xai",
};

function normalizeProvider(raw: string | undefined): ProviderId | undefined {
  if (!raw) {
    return undefined;
  }
  const clean = raw.toLowerCase().trim();
  return OFFICIAL_PROVIDERS[clean];
}

function toLocalDateString(d: Date): string {
  const year = d.getFullYear();
  const month = String(d.getMonth() + 1).padStart(2, "0");
  const day = String(d.getDate()).padStart(2, "0");
  return `${year}-${month}-${day}`;
}

function getRecentDates(days = 7): string[] {
  const dates: string[] = [];
  const now = new Date();
  for (let i = days - 1; i >= 0; i--) {
    const d = new Date(now.getFullYear(), now.getMonth(), now.getDate() - i);
    dates.push(toLocalDateString(d));
  }
  return dates;
}

function findJsonlFiles(dir: string): string[] {
  const results: string[] = [];
  if (!existsSync(dir)) {
    return results;
  }
  function walk(cur: string) {
    try {
      const entries = readdirSync(cur, { withFileTypes: true });
      for (const e of entries) {
        const full = join(cur, e.name);
        if (e.isDirectory()) {
          walk(full);
        } else if (e.isFile() && e.name.endsWith(".jsonl")) {
          results.push(full);
        }
      }
    } catch {
      // Ignore unreadable dirs
    }
  }
  walk(dir);
  return results;
}

function parseDayString(ts: unknown, defaultDay: string): string {
  if (!ts) {
    return defaultDay;
  }
  const str = String(ts).trim();
  if (/^\d{10,}$/.test(str)) {
    const ms = str.length > 10 ? Number(str) : Number(str) * 1000;
    return toLocalDateString(new Date(ms));
  }
  const parsed = new Date(str);
  if (!isNaN(parsed.getTime())) {
    return toLocalDateString(parsed);
  }
  return defaultDay;
}

function scanLocalSessionStats(): Record<ProviderId, LocalProviderStats> {
  const recentDaysList = getRecentDates(7);
  const todayStr = recentDaysList[recentDaysList.length - 1];

  const sessionDirs = [
    join(homedir(), ".pi", "agent", "sessions"),
    join(homedir(), ".omp", "agent", "sessions"),
  ];

  const files: string[] = [];
  for (const d of sessionDirs) {
    files.push(...findJsonlFiles(d));
  }

  let cache: { files: Record<string, FileStatsEntry>; version: number } = {
    files: {},
    version: 11,
  };

  if (existsSync(CACHE_FILE)) {
    try {
      const parsed = JSON.parse(readFileSync(CACHE_FILE, "utf8"));
      if (parsed?.version === 11 && typeof parsed?.files === "object") {
        cache = parsed;
      }
    } catch {
      // Reset corrupted cache
    }
  }

  let cacheDirty = false;
  const outFiles: Record<string, FileStatsEntry> = {};

  for (const f of files) {
    try {
      const st = statSync(f);
      const mtime = st.mtimeMs;
      const cached = cache.files[f];

      if (cached && String(cached.mtime_ns) === String(mtime) && cached.size === st.size) {
        outFiles[f] = cached;
        continue;
      }

      // Fast line-by-line scanning with JSON parsing on assistant lines
      const content = readFileSync(f, "utf8");
      const lines = content.split("\n");
      const fileStats: FileStatsEntry["stats"] = {};

      for (const line of lines) {
        if (!line.includes('"usage":') || !line.includes('"role":"assistant"')) {
          continue;
        }

        try {
          const entry = JSON.parse(line);
          const msg = entry.message || {};
          const rawProv = msg.provider || entry.provider;
          const canonical = normalizeProvider(rawProv);
          if (!canonical) {
            continue;
          }

          const model = msg.model || entry.model || "unknown";
          const rawTs = entry.timestamp || msg.timestamp;
          const day = parseDayString(rawTs, todayStr);

          const usage = msg.usage || entry.usage || {};
          const inp = Number(usage.input || 0);
          const out = Number(usage.output || 0);
          const cr = Number(usage.cacheRead || 0);
          const cw = Number(usage.cacheWrite || 0);
          const total = Number(usage.totalTokens || inp + out + cr + cw);

          if (total <= 0) {
            continue;
          }

          if (!fileStats[canonical]) {
            fileStats[canonical] = {
              model_usage: {},
              models_by_day: {},
              prompts_by_day: {},
              tokens_by_day: {},
              total_prompts: 0,
            };
          }
          const rec = fileStats[canonical];

          rec.total_prompts = (rec.total_prompts || 0) + 1;
          if (!rec.prompts_by_day) {
            rec.prompts_by_day = {};
          }
          rec.prompts_by_day[day] = (rec.prompts_by_day[day] || 0) + 1;

          if (!rec.tokens_by_day) {
            rec.tokens_by_day = {};
          }
          rec.tokens_by_day[day] = (rec.tokens_by_day[day] || 0) + total;

          if (!rec.models_by_day) {
            rec.models_by_day = {};
          }
          if (!rec.models_by_day[day]) {
            rec.models_by_day[day] = {};
          }
          const dayModels = rec.models_by_day[day];
          if (!dayModels[model]) {
            dayModels[model] = { cr: 0, cw: 0, inp: 0, out: 0, tot: 0 };
          }
          const m = dayModels[model];
          m.inp += inp;
          m.out += out;
          m.cr += cr;
          m.cw += cw;
          m.tot += total;

          if (!rec.model_usage) {
            rec.model_usage = {};
          }
          if (!rec.model_usage[model]) {
            rec.model_usage[model] = { cr: 0, cw: 0, inp: 0, out: 0 };
          }
          const mu = rec.model_usage[model];
          mu.inp += inp;
          mu.out += out;
          mu.cr += cr;
          mu.cw += cw;
        } catch {
          // Skip unparseable line
        }
      }

      outFiles[f] = {
        mtime_ns: mtime,
        size: st.size,
        stats: fileStats,
      };
      cacheDirty = true;
    } catch {
      // Skip unreadable file
    }
  }

  if (cacheDirty) {
    try {
      if (!existsSync(CACHE_DIR)) {
        mkdirSync(CACHE_DIR, { recursive: true });
      }
      writeFileSync(CACHE_FILE, JSON.stringify({ files: outFiles, version: 11 }), "utf8");
    } catch {
      // Ignore cache write failures
    }
  }

  // Aggregate across all files per provider
  const aggregated: Partial<
    Record<
      ProviderId,
      {
        activeDates: Set<string>;
        modelUsage: Record<string, { cr: number; cw: number; inp: number; out: number }>;
        recentTokensByDay: Record<string, number>;
        todayPrompts: number;
        todaySessions: Set<string>;
        todayTokens: number;
        todayTokensByModel: Record<string, number>;
        totalPrompts: number;
        totalSessions: Set<string>;
      }
    >
  > = {};

  for (const [filePath, entry] of Object.entries(outFiles)) {
    for (const [providerKey, st] of Object.entries(entry.stats)) {
      const canonical = providerKey as ProviderId;
      if (!aggregated[canonical]) {
        aggregated[canonical] = {
          activeDates: new Set<string>(),
          modelUsage: {},
          recentTokensByDay: Object.fromEntries(recentDaysList.map((d) => [d, 0])),
          todayPrompts: 0,
          todaySessions: new Set<string>(),
          todayTokens: 0,
          todayTokensByModel: {},
          totalPrompts: 0,
          totalSessions: new Set<string>(),
        };
      }
      const agg = aggregated[canonical]!;

      const prompts = st.total_prompts || 0;
      if (prompts > 0) {
        agg.totalPrompts += prompts;
        agg.totalSessions.add(filePath);
      }

      for (const [day, count] of Object.entries(st.prompts_by_day || {})) {
        agg.activeDates.add(day);
        if (day === todayStr) {
          agg.todayPrompts += count;
          agg.todaySessions.add(filePath);
        }
      }

      for (const [day, tokens] of Object.entries(st.tokens_by_day || {})) {
        if (day in agg.recentTokensByDay) {
          agg.recentTokensByDay[day] += tokens;
        }
        if (day === todayStr) {
          agg.todayTokens += tokens;
        }
      }

      for (const [model, stats] of Object.entries(st.model_usage || {})) {
        if (!agg.modelUsage[model]) {
          agg.modelUsage[model] = { cr: 0, cw: 0, inp: 0, out: 0 };
        }
        const modelEntry = agg.modelUsage[model];
        modelEntry.inp += stats.inp;
        modelEntry.out += stats.out;
        modelEntry.cr += stats.cr;
        modelEntry.cw += stats.cw;
      }

      const todayModels = st.models_by_day?.[todayStr] || {};
      for (const [model, stats] of Object.entries(todayModels)) {
        agg.todayTokensByModel[model] = (agg.todayTokensByModel[model] || 0) + stats.tot;
      }
    }
  }

  // Convert to clean LocalProviderStats
  const finalStats: Record<ProviderId, LocalProviderStats> = {
    antigravity: createEmptyStats(recentDaysList),
    deepseek: createEmptyStats(recentDaysList),
    "kimi-coding": createEmptyStats(recentDaysList),
    "openai-codex": createEmptyStats(recentDaysList),
    "opencode-go": createEmptyStats(recentDaysList),
    openrouter: createEmptyStats(recentDaysList),
    xai: createEmptyStats(recentDaysList),
  };

  for (const [providerKey, agg] of Object.entries(aggregated)) {
    const canonical = providerKey as ProviderId;
    if (!agg) {
      continue;
    }
    const activeDatesSorted = [...agg.activeDates].toSorted((a, b) => a.localeCompare(b));
    const recentDays = recentDaysList.map((date) => ({
      date,
      messageCount: agg.recentTokensByDay[date] || 0,
    }));

    const modelUsageFormatted: LocalProviderStats["modelUsage"] = {};
    for (const [model, stats] of Object.entries(agg.modelUsage)) {
      modelUsageFormatted[model] = {
        cacheCreationInputTokens: stats.cw,
        cacheReadInputTokens: stats.cr,
        inputTokens: stats.inp,
        outputTokens: stats.out,
      };
    }

    finalStats[canonical] = {
      activeDates: activeDatesSorted,
      activeDays: activeDatesSorted.length,
      hasLocalStats: true,
      hasPromptStats: true,
      modelUsage: modelUsageFormatted,
      recentDays,
      todayPrompts: agg.todayPrompts,
      todaySessions: agg.todaySessions.size,
      todayTokensByModel: agg.todayTokensByModel,
      todayTotalTokens: agg.todayTokens,
      totalPrompts: agg.totalPrompts,
      totalSessions: agg.totalSessions.size,
    };
  }

  return finalStats;
}

function createEmptyStats(recentDaysList = getRecentDates(7)): LocalProviderStats {
  return {
    activeDates: [],
    activeDays: 0,
    hasLocalStats: true,
    hasPromptStats: true,
    modelUsage: {},
    recentDays: recentDaysList.map((date) => ({ date, messageCount: 0 })),
    todayPrompts: 0,
    todaySessions: 0,
    todayTokensByModel: {},
    todayTotalTokens: 0,
    totalPrompts: 0,
    totalSessions: 0,
  };
}

export type { LocalProviderStats };
export { createEmptyStats, scanLocalSessionStats };