// =======================
// AI Proxy Usage Widget
// =======================
// Scriptable widget for https://ai.ein-softworks.com/usage
// Supports: Small, Medium, Large sizes

const API_URL = "https://ai.ein-softworks.com/usage";

// -- Color Palette --
const COLORS = {
  bg:      new Color("#0d1117"),
  surface: new Color("#161b22"),
  border:  new Color("#30363d"),
  title:   new Color("#e6edf3"),
  label:   new Color("#8b949e"),
  value:   new Color("#e6edf3"),
  accent:  new Color("#58a6ff"),
  green:   new Color("#3fb950"),
  yellow:  new Color("#d29922"),
  red:     new Color("#f85149"),
  purple:  new Color("#bc8cff"),
  dim:     new Color("#484f58"),
};

// -- Fetch --
async function fetchUsage() {
  try {
    let req = new Request(API_URL);
    req.timeoutInterval = 10;
    return await req.loadJSON();
  } catch (e) {
    return null;
  }
}

// -- Helpers --
function fmt$(n) {
  if (n == null) return "—";
  if (n === 0) return "$0.00";
  if (n < 0.01) return "<$0.01";
  if (n < 1) return "$" + n.toFixed(3);
  return "$" + n.toFixed(2);
}

function fmtTok(n) {
  if (n == null) return "—";
  if (n >= 1_000_000) return (n / 1_000_000).toFixed(1) + "M";
  if (n >= 1_000) return (n / 1_000).toFixed(1) + "k";
  return String(n);
}

function fmtPeriod(p) {
  if (!p) return "—";
  let parts = p.split("-");
  if (parts.length === 2) return parseInt(parts[1], 10) + "/" + parts[0];
  return p;
}

function fmtTime(iso) {
  let d = iso ? new Date(iso) : new Date();
  return d.toLocaleTimeString([], { hour: "2-digit", minute: "2-digit" });
}

function provColor(name) {
  return { anthropic: COLORS.accent, deepseek: COLORS.green, zai: COLORS.purple }[name] || COLORS.label;
}

function pl(n, s) { return n + " " + s + (n === 1 ? "" : "s"); }

// Aggregate monthly stats for a provider by summing all its models
function providerMonthly(prov) {
  let cost = 0, reqs = 0, inTok = 0, outTok = 0;
  for (let m of Object.values(prov.models || {})) {
    let mo = m.monthly || {};
    reqs += mo.requests || 0;
    inTok += mo.inputTokens || 0;
    outTok += mo.outputTokens || 0;
    cost += mo.costUsd || 0;
  }
  return { cost, reqs, inTok, outTok };
}

// Aggregate monthly/today stats for all models that route to a given target
function findModelStats(data, providerName, modelName) {
  let prov = data.providers?.[providerName];
  if (!prov) return null;

  // Sum all models whose routesTo matches, OR whose key matches
  let monthly = { inputTokens: 0, outputTokens: 0, requests: 0, costUsd: 0 };
  let today = { inputTokens: 0, outputTokens: 0, requests: 0, costUsd: 0 };
  let total = { inputTokens: 0, outputTokens: 0, requests: 0, costUsd: 0 };
  let found = false;

  for (let [key, m] of Object.entries(prov.models || {})) {
    if (m.routesTo === modelName || key === modelName) {
      found = true;
      let mo = m.monthly || {};
      monthly.inputTokens += mo.inputTokens || 0;
      monthly.outputTokens += mo.outputTokens || 0;
      monthly.requests += mo.requests || 0;
      monthly.costUsd += mo.costUsd || 0;

      let td = m.today || {};
      today.inputTokens += td.inputTokens || 0;
      today.outputTokens += td.outputTokens || 0;
      today.requests += td.requests || 0;
      today.costUsd += td.costUsd || 0;

      let tt = m.total || {};
      total.inputTokens += tt.inputTokens || 0;
      total.outputTokens += tt.outputTokens || 0;
      total.requests += tt.requests || 0;
      total.costUsd += tt.costUsd || 0;
    }
  }

  if (!found) return null;
  return { monthly, today, total };
}

// ========================================
// SMALL WIDGET
// ========================================
function buildSmall(w, data) {
  w.setPadding(12, 12, 12, 12);

  // Header
  let hdr = w.addStack();
  hdr.layoutHorizontally();
  hdr.centerAlignContent();
  addText(hdr, "AI Proxy", Font.boldMonospacedSystemFont(13), COLORS.title);
  hdr.addSpacer();
  addText(hdr, fmtPeriod(data.currentPeriod), Font.mediumMonospacedSystemFont(10), COLORS.label);

  w.addSpacer(6);

  // Monthly cost
  let c = w.addText(fmt$(data.totals?.monthlyCostUsd));
  c.font = Font.boldMonospacedSystemFont(26);
  c.textColor = COLORS.accent;
  c.minimumScaleFactor = 0.6;
  addText(w, "monthly spend", Font.systemFont(10), COLORS.label);

  w.addSpacer(4);

  // Today
  let todayReqs = data.totals?.todayRequests ?? 0;
  let tRow = w.addStack();
  tRow.layoutHorizontally();
  addText(tRow, "today", Font.systemFont(10), COLORS.label);
  tRow.addSpacer();
  addText(tRow, todayReqs > 0 ? fmt$(data.totals.todayCostUsd) : "no usage",
    Font.mediumMonospacedSystemFont(11), todayReqs > 0 ? COLORS.value : COLORS.dim);

  w.addSpacer(3);

  // DeepSeek balance
  let dsBal = data.providers?.deepseek?.balance?.totalBalance;
  let bRow = w.addStack();
  bRow.layoutHorizontally();
  addText(bRow, "deepseek", Font.systemFont(10), COLORS.label);
  bRow.addSpacer();
  addText(bRow, dsBal ? "$" + dsBal : "—", Font.mediumMonospacedSystemFont(11), COLORS.green);

  w.addSpacer();
  addFooter(w, data);
}

// ========================================
// MEDIUM WIDGET
// ========================================
function buildMedium(w, data) {
  w.setPadding(12, 14, 12, 14);

  // Header
  let hdr = w.addStack();
  hdr.layoutHorizontally();
  hdr.centerAlignContent();
  addText(hdr, "AI Proxy Usage", Font.boldMonospacedSystemFont(13), COLORS.title);
  hdr.addSpacer();
  addText(hdr, fmtPeriod(data.currentPeriod), Font.mediumMonospacedSystemFont(11), COLORS.label);

  w.addSpacer(6);

  let body = w.addStack();
  body.layoutHorizontally();

  // -- Left: cost + stats --
  let left = body.addStack();
  left.layoutVertically();

  let c = left.addText(fmt$(data.totals?.monthlyCostUsd));
  c.font = Font.boldMonospacedSystemFont(26);
  c.textColor = COLORS.accent;
  c.minimumScaleFactor = 0.6;
  addText(left, "monthly spend", Font.systemFont(10), COLORS.label);

  left.addSpacer(5);
  addText(left, pl(data.totals?.monthlyRequests ?? 0, "request"),
    Font.mediumMonospacedSystemFont(11), COLORS.value);

  left.addSpacer(3);

  // Today summary
  let tReqs = data.totals?.todayRequests ?? 0;
  let tStr = tReqs > 0
    ? "today: " + fmt$(data.totals.todayCostUsd) + " / " + pl(tReqs, "req")
    : "today: no usage";
  addText(left, tStr, Font.regularMonospacedSystemFont(9), COLORS.dim);

  body.addSpacer();

  // -- Right: providers --
  let right = body.addStack();
  right.layoutVertically();

  let provs = data.providers || {};
  for (let name of ["zai", "anthropic", "deepseek"]) {
    let prov = provs[name];
    if (!prov) continue;

    let row = right.addStack();
    row.layoutHorizontally();
    row.centerAlignContent();

    addText(row, "●", Font.systemFont(8), provColor(name));
    row.addSpacer(4);
    let nm = addText(row, name, Font.mediumMonospacedSystemFont(11), COLORS.value);
    nm.lineLimit = 1;
    row.addSpacer(6);

    let detail = "";
    let stats = providerMonthly(prov);
    if (name === "zai") {
      detail = "$" + (prov.monthlyCostUsd || 10) + "/mo flat";
    } else if (stats.reqs > 0) {
      detail = fmt$(stats.cost);
    } else if (name === "deepseek" && prov.balance?.totalBalance) {
      detail = "$" + prov.balance.totalBalance + " bal";
    } else if (name === "anthropic") {
      detail = "see console";
    } else {
      detail = "no usage";
    }

    let dt = addText(row, detail, Font.regularMonospacedSystemFont(10), COLORS.label);
    dt.lineLimit = 1;
    right.addSpacer(3);
  }

  w.addSpacer();
  addFooter(w, data);
}

// ========================================
// LARGE WIDGET
// ========================================
function buildLarge(w, data) {
  w.setPadding(12, 14, 12, 14);

  // Header
  let hdr = w.addStack();
  hdr.layoutHorizontally();
  hdr.centerAlignContent();
  addText(hdr, "AI Proxy Usage", Font.boldMonospacedSystemFont(15), COLORS.title);
  hdr.addSpacer();
  addText(hdr, fmtPeriod(data.currentPeriod), Font.mediumMonospacedSystemFont(12), COLORS.label);

  w.addSpacer(6);

  // Summary row
  let sum = w.addStack();
  sum.layoutHorizontally();

  let costCol = sum.addStack();
  costCol.layoutVertically();
  let cv = costCol.addText(fmt$(data.totals?.monthlyCostUsd));
  cv.font = Font.boldMonospacedSystemFont(28);
  cv.textColor = COLORS.accent;
  cv.minimumScaleFactor = 0.5;
  addText(costCol, "monthly spend", Font.systemFont(10), COLORS.label);

  sum.addSpacer();

  let statsCol = sum.addStack();
  statsCol.layoutVertically();
  addKVRow(statsCol, "requests", pl(data.totals?.monthlyRequests ?? 0, "req"), COLORS.value);
  statsCol.addSpacer(3);
  addKVRow(statsCol, "all-time", fmt$(data.totals?.totalCostUsd), COLORS.yellow);
  statsCol.addSpacer(3);
  let tReqs = data.totals?.todayRequests ?? 0;
  let todayVal = tReqs > 0
    ? fmt$(data.totals.todayCostUsd) + " / " + pl(tReqs, "req")
    : "no usage";
  addKVRow(statsCol, "today", todayVal, tReqs > 0 ? COLORS.value : COLORS.dim);

  w.addSpacer(8);
  addDivider(w);
  w.addSpacer(6);

  // -- Per-provider sections using routing data --
  let routing = data.routing || {};
  let providers = data.providers || {};

  for (let name of ["zai", "anthropic", "deepseek"]) {
    let prov = providers[name];
    if (!prov) continue;

    // Provider header
    let phdr = w.addStack();
    phdr.layoutHorizontally();
    phdr.centerAlignContent();
    addText(phdr, "●", Font.systemFont(9), provColor(name));
    phdr.addSpacer(5);
    addText(phdr, name.toUpperCase(), Font.boldMonospacedSystemFont(11), COLORS.value);
    phdr.addSpacer();

    // Provider-level info (balance / flat rate / console note)
    if (name === "deepseek" && prov.balance?.totalBalance) {
      addText(phdr, "$" + prov.balance.totalBalance + " balance", Font.regularMonospacedSystemFont(10), COLORS.green);
    } else if (name === "zai") {
      addText(phdr, "$" + (prov.monthlyCostUsd || 10) + "/mo flat", Font.regularMonospacedSystemFont(10), COLORS.purple);
    } else if (name === "anthropic") {
      addText(phdr, "see console", Font.regularMonospacedSystemFont(10), COLORS.label);
    }

    // Show status if error
    if (prov.status === "error") {
      phdr.addSpacer(6);
      addText(phdr, "ERROR", Font.boldMonospacedSystemFont(9), COLORS.red);
    }

    w.addSpacer(3);

    // Find routes that belong to this provider
    let provRoutes = [];
    for (let [tier, route] of Object.entries(routing)) {
      if (tier === "default") continue;
      if (route.provider === name) {
        provRoutes.push({ tier, ...route });
      }
    }

    if (provRoutes.length === 0) {
      let noRow = w.addStack();
      noRow.setPadding(2, 18, 2, 8);
      addText(noRow, "No active routes", Font.regularMonospacedSystemFont(10), COLORS.dim);
    } else {
      for (let route of provRoutes) {
        let modelData = findModelStats(data, name, route.model);
        let mo = modelData?.monthly || {};
        let hasUsage = (mo.requests || 0) > 0 || (mo.inputTokens || 0) > 0 || (mo.outputTokens || 0) > 0;

        let card = w.addStack();
        card.layoutHorizontally();
        card.centerAlignContent();
        card.setPadding(4, 18, 4, 8);
        card.backgroundColor = COLORS.surface;
        card.cornerRadius = 6;

        // Left: tier + model + stats
        let lc = card.addStack();
        lc.layoutVertically();

        // Tier label + actual model
        let tierRow = lc.addStack();
        tierRow.layoutHorizontally();
        tierRow.centerAlignContent();
        addText(tierRow, route.tier, Font.boldMonospacedSystemFont(11), COLORS.value);
        tierRow.addSpacer(6);
        let arrow = addText(tierRow, "→ " + route.model, Font.regularMonospacedSystemFont(10), COLORS.label);
        arrow.lineLimit = 1;
        arrow.minimumScaleFactor = 0.7;

        if (hasUsage) {
          let tokStr = fmtTok(mo.inputTokens) + " in / " + fmtTok(mo.outputTokens) + " out tokens";
          addText(lc, tokStr, Font.regularMonospacedSystemFont(10), COLORS.label);
          addText(lc, pl(mo.requests, "request"), Font.regularMonospacedSystemFont(10), COLORS.dim);
        } else {
          addText(lc, "no usage this period", Font.regularMonospacedSystemFont(10), COLORS.dim);
        }

        card.addSpacer();

        // Right: cost or token count for flat rate
        if (name === "zai") {
          if (hasUsage) {
            addText(card, pl(mo.requests, "req"), Font.mediumMonospacedSystemFont(12), COLORS.purple);
          }
        } else if (hasUsage) {
          addText(card, fmt$(mo.costUsd), Font.boldMonospacedSystemFont(14), provColor(name));
        } else {
          addText(card, "$0.00", Font.regularMonospacedSystemFont(12), COLORS.dim);
        }
      }
    }

    w.addSpacer(5);
  }

  w.addSpacer();
  addFooter(w, data);
}

// ========================================
// Shared UI helpers
// ========================================
function addText(parent, str, font, color) {
  let t = parent.addText(str);
  t.font = font;
  t.textColor = color;
  return t;
}

function addKVRow(stack, label, value, color) {
  let row = stack.addStack();
  row.layoutHorizontally();
  addText(row, label + "  ", Font.systemFont(10), COLORS.label);
  row.addSpacer();
  addText(row, value, Font.boldMonospacedSystemFont(11), color);
}

function addDivider(w) {
  let ctx = new DrawContext();
  ctx.opaque = false;
  ctx.respectScreenScale = true;
  ctx.size = new Size(320, 1);
  ctx.setFillColor(COLORS.border);
  ctx.fillRect(new Rect(0, 0, 320, 1));
  let img = w.addImage(ctx.getImage());
  img.imageSize = new Size(320, 1);
}

function addFooter(w, data) {
  let f = w.addStack();
  f.layoutHorizontally();
  f.addSpacer();
  addText(f, "updated " + fmtTime(data.generatedAt), Font.regularMonospacedSystemFont(9), COLORS.dim);
}

// -- Error --
function buildError(w, msg) {
  w.setPadding(12, 12, 12, 12);
  addText(w, "AI Proxy Usage", Font.boldMonospacedSystemFont(12), COLORS.title);
  w.addSpacer(8);
  let e = addText(w, msg, Font.systemFont(11), COLORS.red);
  e.minimumScaleFactor = 0.7;
}

// -- Main --
async function main() {
  let w = new ListWidget();

  let grad = new LinearGradient();
  grad.locations = [0, 1];
  grad.colors = [COLORS.bg, new Color("#0a0e14")];
  w.backgroundGradient = grad;
  w.url = API_URL;

  let data = await fetchUsage();

  if (!data) {
    buildError(w, "Failed to reach API");
  } else {
    let family = config.widgetFamily;
    if (family === "small") buildSmall(w, data);
    else if (family === "large") buildLarge(w, data);
    else buildMedium(w, data);
  }

  w.refreshAfterDate = new Date(Date.now() + 30 * 60 * 1000);

  if (config.runsInWidget) {
    Script.setWidget(w);
  } else {
    await w.presentLarge();
  }

  Script.complete();
}

await main();