// credit: gist.vyr.my.id
import axios from "axios";
import { readFileSync, writeFileSync, existsSync, unlinkSync } from "fs";
import { homedir } from "os";
import { join } from "path";

const BASE = "https://chatai.org";
const CACHE_PATH = join(homedir(), ".chatai-session.json");
const SESSION_TTL = 2 * 60 * 60 * 1000;
const MAX_RETRIES = 3;
const REFRESH_THRESHOLD = 5;

const UA = "Mozilla/5.0 (Linux; Android 10; K) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/140.0.0.0 Safari/537.36";

const MODELS = {
  chatgpt:   "openai/gpt-4o-mini",
  claude:    "anthropic/claude-haiku-4-5",
  gemini:    "google/gemini-2.0-flash-001",
  grok:      "x-ai/grok-3-mini-beta",
  deepseek:  "deepseek/deepseek-chat-v3-0324",
  qwen:      "qwen/qwen-2.5-72b-instruct",
  kimi:      "moonshotai/moonlight-16k",
  perplexity:"perplexity/sonar",
};

const DEFAULT_MODEL = "chatgpt";

const $$ = (v) => JSON.stringify(v, null, 2);

function resolveModel(input) {
  if (!input) return MODELS[DEFAULT_MODEL];
  const lower = input.toLowerCase();
  if (MODELS[lower]) return MODELS[lower];
  const values = Object.values(MODELS);
  if (values.includes(input)) return input;
  throw new Error(`Model "${input}" tidak dikenal. Gunakan --models untuk melihat daftar.`);
}

function parseCookieHeaders(headers) {
  const raw = headers["set-cookie"] ?? [];
  const jar = {};
  for (const entry of raw) {
    const [pair] = entry.split(";");
    const eq = pair.indexOf("=");
    if (eq === -1) continue;
    jar[pair.slice(0, eq).trim()] = pair.slice(eq + 1).trim();
  }
  return jar;
}

function cookieString(jar) {
  return Object.entries(jar).map(([k, v]) => `${k}=${v}`).join("; ");
}

function extractCsrf(jar) {
  const raw = decodeURIComponent(jar["XSRF-TOKEN"] ?? "");
  try {
    return JSON.parse(Buffer.from(raw.split(".")[0], "base64").toString("utf8"))?.mac ?? raw;
  } catch {
    return raw;
  }
}

function loadCache() {
  try {
    if (!existsSync(CACHE_PATH)) return null;
    const data = JSON.parse(readFileSync(CACHE_PATH, "utf8"));
    if (Date.now() > data.expiresAt) return null;
    return data;
  } catch {
    return null;
  }
}

function saveCache(jar, csrfToken, remaining = null) {
  const prev = loadCache();
  writeFileSync(CACHE_PATH, $$(
    { jar, csrfToken, expiresAt: Date.now() + SESSION_TTL, remaining: remaining ?? prev?.remaining ?? null }
  ), "utf8");
}

function updateCacheRemaining(remaining) {
  try {
    const data = JSON.parse(readFileSync(CACHE_PATH, "utf8"));
    data.remaining = remaining;
    writeFileSync(CACHE_PATH, $$(data), "utf8");
  } catch {}
}

function clearCache() {
  try { if (existsSync(CACHE_PATH)) unlinkSync(CACHE_PATH); } catch {}
}

async function fetchSession(force = false) {
  if (!force) {
    const cached = loadCache();
    if (cached) {
      if (cached.remaining !== null && cached.remaining <= REFRESH_THRESHOLD) {
        process.stderr.write(`[chatai] sisa quota ${cached.remaining} ≤ ${REFRESH_THRESHOLD}, refresh session...\n`);
      } else {
        return cached;
      }
    }
  }

  const res = await axios.get(`${BASE}/chat`, {
    headers: {
      "User-Agent": UA,
      Accept: "text/html,application/xhtml+xml,application/xml;q=0.9,*/*;q=0.8",
      "Accept-Language": "id-ID,id;q=0.9,en-US;q=0.8,en;q=0.7",
      "Cache-Control": "no-cache",
      Pragma: "no-cache",
    },
    maxRedirects: 5,
  });

  const jar = parseCookieHeaders(res.headers);
  const csrfToken = extractCsrf(jar);
  saveCache(jar, csrfToken, null);
  return { jar, csrfToken, remaining: null };
}

function buildHeaders(jar, csrfToken) {
  return {
    "Content-Type": "application/json",
    Accept: "*/*",
    Origin: BASE,
    Referer: `${BASE}/chat`,
    "X-CSRF-TOKEN": csrfToken,
    Cookie: cookieString(jar),
    "User-Agent": UA,
    "sec-ch-ua": '"Chromium";v="140", "Not?A_Brand";v="8"',
    "sec-ch-ua-mobile": "?1",
    "sec-ch-ua-platform": '"Android"',
    "Sec-Fetch-Dest": "empty",
    "Sec-Fetch-Mode": "cors",
    "Sec-Fetch-Site": "same-origin",
    "Cache-Control": "no-cache",
    Pragma: "no-cache",
  };
}

function parseSSE(raw) {
  return raw
    .split("\n")
    .filter((l) => l.startsWith("data: ") && l !== "data: [DONE]")
    .map((l) => { try { return JSON.parse(l.slice(6)); } catch { return null; } })
    .filter(Boolean);
}

function extractText(chunks) {
  return chunks.flatMap((c) => c.choices ?? []).map((ch) => ch.delta?.content ?? "").join("");
}

function isQuotaError(status) {
  return [419, 422, 429, 401, 403].includes(status);
}

async function chat({ messages, model, jar, csrfToken }) {
  const res = await axios.post(
    `${BASE}/api/chat`,
    { model, messages },
    { headers: buildHeaders(jar, csrfToken), responseType: "text" }
  );

  const remaining = parseInt(res.headers["x-ratelimit-remaining"] ?? "-1", 10);
  if (remaining >= 0) updateCacheRemaining(remaining);

  const chunks = parseSSE(res.data);
  const last = chunks.at(-1);

  return {
    id: last?.id ?? null,
    model: last?.model ?? model,
    provider: last?.provider ?? null,
    content: extractText(chunks),
    usage: last?.usage ?? null,
    quota: { remaining, limit: parseInt(res.headers["x-ratelimit-limit"] ?? "-1", 10) },
  };
}

async function chatStream({ messages, model, jar, csrfToken }) {
  const res = await axios.post(
    `${BASE}/api/chat`,
    { model, messages },
    { headers: buildHeaders(jar, csrfToken), responseType: "stream" }
  );

  const remaining = parseInt(res.headers["x-ratelimit-remaining"] ?? "-1", 10);
  if (remaining >= 0) updateCacheRemaining(remaining);

  return new Promise((resolve, reject) => {
    let buf = "";
    let fullText = "";
    let meta = {};

    res.data.on("data", (chunk) => {
      buf += chunk.toString();
      const lines = buf.split("\n");
      buf = lines.pop();

      for (const line of lines) {
        if (!line.startsWith("data: ") || line === "data: [DONE]") continue;
        try {
          const parsed = JSON.parse(line.slice(6));
          const delta = parsed.choices?.[0]?.delta?.content ?? "";
          if (delta) {
            fullText += delta;
            process.stderr.write(delta);
          }
          if (parsed.usage) {
            meta = { id: parsed.id, model: parsed.model, provider: parsed.provider, usage: parsed.usage };
          }
        } catch {}
      }
    });

    res.data.on("end", () => {
      process.stderr.write("\n");
      resolve({
        ...meta,
        content: fullText,
        quota: { remaining, limit: parseInt(res.headers?.["x-ratelimit-limit"] ?? "-1", 10) },
      });
    });
    res.data.on("error", reject);
  });
}

async function runWithRetry({ messages, model, stream, attempt = 0 }) {
  const forceRefresh = attempt > 0;
  if (forceRefresh) {
    clearCache();
    process.stderr.write(`[chatai] session refresh (attempt ${attempt + 1}/${MAX_RETRIES})\n`);
  }

  const session = await fetchSession(forceRefresh);
  const payload = { messages, model, jar: session.jar, csrfToken: session.csrfToken };

  try {
    return stream ? await chatStream(payload) : await chat(payload);
  } catch (err) {
    const status = err.response?.status;
    if (isQuotaError(status) && attempt < MAX_RETRIES - 1) {
      return runWithRetry({ messages, model, stream, attempt: attempt + 1 });
    }
    throw err;
  }
}

function parseArgs(argv) {
  const args = argv.slice(2);
  const out = {
    message: null,
    model: null,
    stream: false,
    refresh: false,
    showSession: false,
    showModels: false,
    help: false,
  };

  for (let i = 0; i < args.length; i++) {
    const a = args[i];
    if (a === "-m" || a === "--message") out.message = args[++i];
    else if (a === "--model") out.model = args[++i];
    else if (a === "--stream") out.stream = true;
    else if (a === "--refresh") out.refresh = true;
    else if (a === "--session") out.showSession = true;
    else if (a === "--models") out.showModels = true;
    else if (a === "-h" || a === "--help") out.help = true;
    else if (!out.message && !a.startsWith("-")) out.message = a;
  }

  return out;
}

function printHelp() {
  const aliases = Object.entries(MODELS).map(([k, v]) => `  ${k.padEnd(12)} → ${v}`).join("\n");
  console.log(`
chatai — CLI wrapper untuk chatai.org

USAGE
  node chatai.mjs "<pesan>"
  node chatai.mjs -m "<pesan>" [options]

OPTIONS
  -m, --message <text>   Pesan ke AI
      --model <alias>    Alias atau full model ID (default: ${DEFAULT_MODEL})
      --stream           Streaming ke stderr, metadata JSON ke stdout
      --refresh          Paksa refresh session sekarang
      --session          Info session & sisa quota
      --models           Daftar semua model yang tersedia
  -h, --help             Tampilkan bantuan ini

MODEL ALIASES
${aliases}

AUTO RETRY
  Jika sisa quota ≤ ${REFRESH_THRESHOLD} atau dapat error 4xx, otomatis
  generate session baru dan retry sampai ${MAX_RETRIES}x.
`.trim());
}

async function main() {
  const flags = parseArgs(process.argv);

  if (flags.help) { printHelp(); process.exit(0); }

  if (flags.showModels) {
    const rows = Object.entries(MODELS).map(([alias, id]) => ({ alias, id }));
    console.log($$(rows));
    process.exit(0);
  }

  if (flags.showSession) {
    const cached = loadCache();
    if (!cached) { console.error("Tidak ada session yang ter-cache."); process.exit(1); }
    console.log($$(
      {
        createdAt: new Date(cached.expiresAt - SESSION_TTL).toISOString(),
        expiresAt: new Date(cached.expiresAt).toISOString(),
        quota: { remaining: cached.remaining ?? "unknown", refreshThreshold: REFRESH_THRESHOLD },
        csrfToken: cached.csrfToken,
        cookies: Object.keys(cached.jar),
      }
    ));
    process.exit(0);
  }

  if (flags.refresh) {
    clearCache();
    const s = await fetchSession(true);
    console.log($$({ ok: true, csrfToken: s.csrfToken, cookies: Object.keys(s.jar) }));
    process.exit(0);
  }

  if (!flags.message) { printHelp(); process.exit(1); }

  let model;
  try {
    model = resolveModel(flags.model);
  } catch (err) {
    console.error(err.message);
    process.exit(1);
  }

  try {
    const result = await runWithRetry({
      messages: [{ role: "user", content: flags.message }],
      model,
      stream: flags.stream,
    });
    console.log($$(result));
  } catch (err) {
    const status = err.response?.status;
    console.error(`Error${status ? ` (${status})` : ""}: ${err.message}`);
    process.exit(1);
  }
}

main();