import { debug } from "../utils/logger";
import { CacheManager } from "../utils/cache";
import type { RawEntryFields } from "../utils/claude";

export interface ModelPricing {
  name: string;
  input: number;
  cache_write_5m: number;
  cache_write_1h?: number;
  cache_read: number;
  output: number;
}

/**
 * Token counts grouped by the rate they are billed at. Totals states save sums
 * of these: a change here needs STATE_VERSION in transcript-totals.ts bumped.
 */
export interface BillableTokens {
  input: number;
  output: number;
  cacheRead: number;
  cacheWrite5m: number;
  cacheWrite1h: number;
}

interface EntryUsage {
  input_tokens?: number;
  output_tokens?: number;
  cache_creation_input_tokens?: number;
  cache_read_input_tokens?: number;
  cache_creation?: {
    ephemeral_1h_input_tokens?: number;
    ephemeral_5m_input_tokens?: number;
  };
}

const OFFLINE_PRICING_DATA: Record<string, ModelPricing> = {
  "claude-haiku-4-5-20251001": {
    name: "Claude Haiku 4.5",
    input: 1.0,
    output: 5.0,
    cache_write_5m: 1.25,
    cache_write_1h: 2.0,
    cache_read: 0.1,
  },
  "claude-haiku-4-5": {
    name: "Claude Haiku 4.5",
    input: 1.0,
    output: 5.0,
    cache_write_5m: 1.25,
    cache_write_1h: 2.0,
    cache_read: 0.1,
  },
  "claude-opus-4-20250514": {
    name: "Claude Opus 4",
    input: 15.0,
    output: 75.0,
    cache_write_5m: 18.75,
    cache_write_1h: 30.0,
    cache_read: 1.5,
  },
  "claude-opus-4-1": {
    name: "Claude Opus 4.1",
    input: 15.0,
    output: 75.0,
    cache_write_5m: 18.75,
    cache_write_1h: 30.0,
    cache_read: 1.5,
  },
  "claude-opus-4-1-20250805": {
    name: "Claude Opus 4.1",
    input: 15.0,
    output: 75.0,
    cache_write_5m: 18.75,
    cache_write_1h: 30.0,
    cache_read: 1.5,
  },
  "claude-sonnet-4-20250514": {
    name: "Claude Sonnet 4",
    input: 3.0,
    output: 15.0,
    cache_write_5m: 3.75,
    cache_write_1h: 6.0,
    cache_read: 0.3,
  },
  "claude-4-opus-20250514": {
    name: "Claude 4 Opus",
    input: 15.0,
    output: 75.0,
    cache_write_5m: 18.75,
    cache_write_1h: 30.0,
    cache_read: 1.5,
  },
  "claude-4-sonnet-20250514": {
    name: "Claude 4 Sonnet",
    input: 3.0,
    output: 15.0,
    cache_write_5m: 3.75,
    cache_write_1h: 6.0,
    cache_read: 0.3,
  },
  "claude-sonnet-4-5": {
    name: "Claude Sonnet 4.5",
    input: 3.0,
    output: 15.0,
    cache_write_5m: 3.75,
    cache_write_1h: 6.0,
    cache_read: 0.3,
  },
  "claude-sonnet-4-5-20250929": {
    name: "Claude Sonnet 4.5",
    input: 3.0,
    output: 15.0,
    cache_write_5m: 3.75,
    cache_write_1h: 6.0,
    cache_read: 0.3,
  },
  "claude-opus-4-5": {
    name: "Claude Opus 4.5",
    input: 5.0,
    output: 25.0,
    cache_write_5m: 6.25,
    cache_write_1h: 10.0,
    cache_read: 0.5,
  },
  "claude-opus-4-5-20251101": {
    name: "Claude Opus 4.5",
    input: 5.0,
    output: 25.0,
    cache_write_5m: 6.25,
    cache_write_1h: 10.0,
    cache_read: 0.5,
  },
  "claude-opus-4-6": {
    name: "Claude Opus 4.6",
    input: 5.0,
    output: 25.0,
    cache_write_5m: 6.25,
    cache_write_1h: 10.0,
    cache_read: 0.5,
  },
  "claude-opus-4-6-20260205": {
    name: "Claude Opus 4.6",
    input: 5.0,
    output: 25.0,
    cache_write_5m: 6.25,
    cache_write_1h: 10.0,
    cache_read: 0.5,
  },
  "claude-sonnet-4-6": {
    name: "Claude Sonnet 4.6",
    input: 3.0,
    output: 15.0,
    cache_write_5m: 3.75,
    cache_write_1h: 6.0,
    cache_read: 0.3,
  },
};

export class PricingService {
  private static executionCache: Record<string, ModelPricing> | null = null;
  // Holds the in-flight load, not just its resolved value. The tui style
  // fetches every segment's info concurrently (see powerline.ts's
  // Promise.allSettled), and session and today can each need pricing for an
  // entry with no recorded costUSD. On a cold cache, each one that arrived
  // before executionCache was set started its own GitHub fetch. Mirrors
  // claude.ts's parsedFileCache for the same reason.
  private static loadingCache: Promise<Record<string, ModelPricing>> | null =
    null;
  private static modelPricingCache = new Map<string, ModelPricing>();
  private static readonly GITHUB_PRICING_URL =
    "https://raw.githubusercontent.com/Owloops/claude-powerline/main/pricing.json";

  private static async loadDiskCache(): Promise<Record<
    string,
    ModelPricing
  > | null> {
    const TTL_24H = 24 * 60 * 60 * 1000;
    const minValidTime = Date.now() - TTL_24H;
    return (await CacheManager.getUsageCache(
      "pricing",
      minValidTime,
    )) as Record<string, ModelPricing> | null;
  }

  private static async saveDiskCache(
    data: Record<string, ModelPricing>,
  ): Promise<void> {
    await CacheManager.setUsageCache("pricing", data);
  }

  private static async fetchPricingData(): Promise<Record<
    string,
    ModelPricing
  > | null> {
    // Only reached when the 24h disk cache misses, so keep node:https off the
    // module graph that every status line render has to load.
    const { get } = await import("node:https");

    return new Promise((resolve) => {
      const parsedUrl = new URL(this.GITHUB_PRICING_URL);

      const request = get(
        {
          hostname: parsedUrl.hostname,
          path: parsedUrl.pathname,
          headers: {
            "User-Agent": "claude-powerline",
            "Cache-Control": "no-cache",
          },
          timeout: 5000,
        },
        (response) => {
          if (response.statusCode !== 200) {
            debug(`HTTP ${response.statusCode}: ${response.statusMessage}`);
            resolve(null);
            return;
          }

          let data = "";
          let size = 0;
          const MAX_SIZE = 1024 * 1024;

          response.on("data", (chunk) => {
            size += chunk.length;
            if (size > MAX_SIZE) {
              debug("Response too large");
              request.destroy();
              resolve(null);
              return;
            }
            data += chunk;
          });

          response.on("end", () => {
            try {
              const json = JSON.parse(data);
              const dataObj = json as Record<string, unknown>;
              const meta = dataObj._meta as { updated?: string } | undefined;

              const pricingData: Record<string, unknown> = {};
              for (const [key, value] of Object.entries(dataObj)) {
                if (key !== "_meta") {
                  pricingData[key] = value;
                }
              }

              if (this.validatePricingData(pricingData)) {
                debug(
                  `Fetched fresh pricing from GitHub for ${Object.keys(pricingData).length} models`,
                );
                debug(`Pricing last updated: ${meta?.updated || "unknown"}`);
                resolve(pricingData);
              } else {
                debug("Invalid pricing data structure");
                resolve(null);
              }
            } catch (error) {
              debug("Failed to parse JSON:", error);
              resolve(null);
            }
          });

          response.on("error", (error) => {
            debug("Response error:", error);
            resolve(null);
          });
        },
      );

      request.on("error", (error) => {
        debug("Request error:", error);
        resolve(null);
      });

      request.on("timeout", () => {
        debug("Request timeout");
        request.destroy();
        resolve(null);
      });

      request.end();
    });
  }

  static getCurrentPricing(): Promise<Record<string, ModelPricing>> {
    if (this.executionCache !== null) {
      debug(
        `[CACHE-HIT] Pricing execution cache: ${Object.keys(this.executionCache).length} models`,
      );
      return Promise.resolve(this.executionCache);
    }

    if (!this.loadingCache) {
      this.loadingCache = this.loadCurrentPricing().finally(() => {
        this.loadingCache = null;
      });
    }
    return this.loadingCache;
  }

  private static async loadCurrentPricing(): Promise<
    Record<string, ModelPricing>
  > {
    const diskCached = await this.loadDiskCache();
    if (diskCached) {
      debug(
        `[CACHE-HIT] Pricing disk cache: ${Object.keys(diskCached).length} models`,
      );
      this.executionCache = diskCached;
      debug(
        `[CACHE-SET] Pricing execution cache stored: ${Object.keys(diskCached).length} models`,
      );
      return diskCached;
    }

    const freshData = await this.fetchPricingData();
    if (freshData) {
      await this.saveDiskCache(freshData);
      debug(
        `[CACHE-SET] Pricing disk cache stored: ${Object.keys(freshData).length} models`,
      );
      this.executionCache = freshData;
      debug(
        `[CACHE-SET] Pricing execution cache stored: ${Object.keys(freshData).length} models`,
      );
      return freshData;
    }

    debug(
      `[CACHE-FALLBACK] Using offline pricing data: ${Object.keys(OFFLINE_PRICING_DATA).length} models`,
    );
    this.executionCache = OFFLINE_PRICING_DATA;
    debug(
      `[CACHE-SET] Pricing execution cache stored: ${Object.keys(OFFLINE_PRICING_DATA).length} models`,
    );
    return OFFLINE_PRICING_DATA;
  }

  private static validatePricingData(
    data: unknown,
  ): data is Record<string, ModelPricing> {
    if (!data || typeof data !== "object") return false;

    for (const [, value] of Object.entries(data)) {
      if (!value || typeof value !== "object") return false;
      const pricing = value as Record<string, unknown>;

      if (
        typeof pricing.input !== "number" ||
        typeof pricing.output !== "number" ||
        typeof pricing.cache_read !== "number"
      ) {
        return false;
      }
    }

    return true;
  }

  static async getModelPricing(modelId: string): Promise<ModelPricing> {
    if (this.modelPricingCache.has(modelId)) {
      debug(`[CACHE-HIT] Model pricing cache: ${modelId}`);
      return this.modelPricingCache.get(modelId)!;
    }

    const allPricing = await this.getCurrentPricing();
    let pricing: ModelPricing;

    if (allPricing[modelId]) {
      pricing = allPricing[modelId];
    } else {
      pricing = this.fuzzyMatchModel(modelId, allPricing);
    }

    this.modelPricingCache.set(modelId, pricing);
    debug(`[CACHE-SET] Model pricing cache: ${modelId}`);
    return pricing;
  }

  private static fuzzyMatchModel(
    modelId: string,
    allPricing: Record<string, ModelPricing>,
  ): ModelPricing {
    const lowerModelId = modelId.toLowerCase();

    for (const [key, pricing] of Object.entries(allPricing)) {
      if (key.toLowerCase() === lowerModelId) {
        return pricing;
      }
    }
    const patterns = [
      {
        pattern: ["opus-4-6", "claude-opus-4-6"],
        fallback: "claude-opus-4-6-20260205",
      },
      {
        pattern: ["opus-4-5", "claude-opus-4-5"],
        fallback: "claude-opus-4-5-20251101",
      },
      {
        pattern: ["opus-4-1", "claude-opus-4-1"],
        fallback: "claude-opus-4-1-20250805",
      },
      {
        pattern: ["opus-4", "claude-opus-4"],
        fallback: "claude-opus-4-20250514",
      },
      {
        pattern: ["sonnet-4-6", "sonnet-4.6", "claude-sonnet-4-6"],
        fallback: "claude-sonnet-4-6",
      },
      {
        pattern: ["sonnet-4.5", "4-5-sonnet", "sonnet-4-5"],
        fallback: "claude-sonnet-4-5-20250929",
      },
      {
        pattern: ["sonnet-4", "claude-sonnet-4"],
        fallback: "claude-sonnet-4-20250514",
      },
      {
        pattern: ["haiku-4.5", "4-5-haiku", "haiku-4-5"],
        fallback: "claude-haiku-4-5-20251001",
      },
      { pattern: ["haiku"], fallback: "claude-haiku-4-5-20251001" },
      { pattern: ["opus"], fallback: "claude-opus-4-20250514" },
      { pattern: ["sonnet"], fallback: "claude-sonnet-4-5-20250929" },
    ];

    for (const { pattern, fallback } of patterns) {
      if (pattern.some((p) => lowerModelId.includes(p))) {
        if (allPricing[fallback]) {
          return allPricing[fallback];
        }
      }
    }

    return (
      allPricing["claude-sonnet-4-5-20250929"] || {
        name: `${modelId} (Unknown Model)`,
        input: 3.0,
        cache_write_5m: 3.75,
        cache_write_1h: 6.0,
        cache_read: 0.3,
        output: 15.0,
      }
    );
  }

  static async calculateCostForEntry(entry: RawEntryFields): Promise<number> {
    const usage: EntryUsage | undefined = entry.message?.usage;
    if (!usage) {
      return 0;
    }

    return this.calculateCost(
      this.extractModelId(entry),
      this.billableTokens(usage),
    );
  }

  /**
   * @info Splits one entry's usage by rate. The split is per entry because the
   * uncategorized cache writes are clamped at zero per entry, so sums of these
   * price exactly like the entries they came from.
   */
  static billableTokens(usage: EntryUsage): BillableTokens {
    const cacheCreationTokens = usage.cache_creation_input_tokens || 0;
    const cacheCreation = usage.cache_creation;
    const oneHourTokens = cacheCreation?.ephemeral_1h_input_tokens ?? 0;
    const explicitFiveMinuteTokens =
      cacheCreation?.ephemeral_5m_input_tokens ?? 0;
    const uncategorizedFiveMinuteTokens = Math.max(
      0,
      cacheCreationTokens - oneHourTokens - explicitFiveMinuteTokens,
    );

    return {
      input: usage.input_tokens || 0,
      output: usage.output_tokens || 0,
      cacheRead: usage.cache_read_input_tokens || 0,
      cacheWrite5m: explicitFiveMinuteTokens + uncategorizedFiveMinuteTokens,
      cacheWrite1h: oneHourTokens,
    };
  }

  static async calculateCost(
    modelId: string,
    tokens: BillableTokens,
  ): Promise<number> {
    const pricing = await this.getModelPricing(modelId);
    const cacheWrite1hRate = pricing.cache_write_1h ?? pricing.cache_write_5m;

    const inputCost = (tokens.input / 1_000_000) * pricing.input;
    const outputCost = (tokens.output / 1_000_000) * pricing.output;
    const cacheReadCost = (tokens.cacheRead / 1_000_000) * pricing.cache_read;
    const cacheCreationCost =
      (tokens.cacheWrite1h / 1_000_000) * cacheWrite1hRate +
      (tokens.cacheWrite5m / 1_000_000) * pricing.cache_write_5m;

    return inputCost + outputCost + cacheCreationCost + cacheReadCost;
  }

  static extractModelId(entry: RawEntryFields): string {
    if (entry.model && typeof entry.model === "string") {
      return entry.model;
    }

    if (entry.message?.model) {
      const model = entry.message.model;
      if (typeof model === "string") {
        return model;
      }
      const modelObj = model as Record<string, unknown> | undefined;
      return (
        (typeof modelObj?.id === "string" ? modelObj.id : null) ||
        "claude-sonnet-4-5-20250929"
      );
    }

    if (entry.model_id && typeof entry.model_id === "string") {
      return entry.model_id;
    }

    return "claude-sonnet-4-5-20250929";
  }
}