import http from "http";
import crypto from "crypto";
import os from "os";
import fs from "fs";
import path from "path";

const HOST = process.env.MIMO_PROXY_HOST || "127.0.0.1";
const PORT = Number.parseInt(process.env.MIMO_PROXY_PORT || "9100", 10);
const PROXY_KEY =
  process.env.MIMO_PROXY_KEY ||
  `sk-mimo-${crypto.randomBytes(16).toString("hex")}`;
const UPSTREAM_API_KEY =
  process.env.MIMO_API_KEY || process.env.XIAOMI_API_KEY || "";
const PROXY_MODE = normalizeMode(process.env.MIMO_PROXY_MODE || "auto");
const UPSTREAM_MODEL =
  process.env.MIMO_UPSTREAM_MODEL ||
  process.env.MIMO_MODEL ||
  (PROXY_MODE === "api" ? "mimo-v2.5-pro" : "mimo-auto");
const PUBLIC_MODEL = process.env.MIMO_PROXY_MODEL || UPSTREAM_MODEL;

const FREE_BASE_URL = (
  process.env.MIMO_FREE_BASE_URL || "https://api.xiaomimimo.com"
).replace(/\/+$/, "");
const BOOTSTRAP_URL = `${FREE_BASE_URL}/api/free-ai/bootstrap`;
const CHAT_BASE = `${FREE_BASE_URL}/api/free-ai/openai`;
const API_BASE_URL = (
  process.env.MIMO_BASE_URL ||
  inferApiBaseUrl(UPSTREAM_API_KEY) ||
  "https://api.xiaomimimo.com/v1"
).replace(/\/+$/, "");
const CLIENT_FILE =
  process.env.MIMO_PROXY_CLIENT_FILE ||
  path.join(os.homedir(), ".local/share/mimocode/mimo-free-client");
const JWT_REFRESH_MARGIN = 5 * 60 * 1000;
const MIMOCODE_USER_AGENT =
  process.env.MIMO_FREE_USER_AGENT ||
  "mimocode/0.1.1 ai-sdk/provider-utils/4.0.23 runtime/bun/1.3.14";
const MIMOCODE_SYSTEM_PROMPT =
  process.env.MIMO_FREE_SYSTEM_PROMPT ||
  "You are MiMoCode, an interactive CLI tool that helps users with software engineering tasks.";
const FREE_COMPAT_SYSTEM = parseBoolean(
  process.env.MIMO_FREE_COMPAT_SYSTEM,
  true,
);
const FREE_SESSION_AFFINITY =
  process.env.MIMO_FREE_SESSION_AFFINITY ||
  `ses_${crypto.randomBytes(12).toString("hex")}`;
const FREE_MAX_TOKENS = Number.parseInt(
  process.env.MIMO_FREE_MAX_TOKENS || "128000",
  10,
);

let cachedJWT = null;
let jwtExpiry = 0;
let pendingBootstrap = null;

function parseBoolean(value, defaultValue = false) {
  if (value === undefined || value === null || value === "") return defaultValue;
  return /^(1|true|yes|on)$/i.test(String(value));
}

function normalizeMode(value) {
  const mode = String(value || "auto").toLowerCase();
  if (["api", "official", "key"].includes(mode)) return "api";
  if (["free", "legacy"].includes(mode)) return "free";
  return UPSTREAM_API_KEY ? "api" : "free";
}

function inferApiBaseUrl(apiKey) {
  if (!apiKey) return "";
  if (apiKey.startsWith("tp-")) {
    return "https://token-plan-cn.xiaomimimo.com/v1";
  }
  return "https://api.xiaomimimo.com/v1";
}

function getFingerprint() {
  try {
    return fs.readFileSync(CLIENT_FILE, "utf-8").trim();
  } catch {}

  const parts = [
    os.hostname(),
    os.platform(),
    os.arch(),
    os.cpus()[0]?.model || "unknown",
    os.userInfo().username,
  ];
  const fp = crypto.createHash("sha256").update(parts.join("|")).digest("hex");

  try {
    fs.mkdirSync(path.dirname(CLIENT_FILE), { recursive: true });
    fs.writeFileSync(CLIENT_FILE, fp, { mode: 0o600 });
  } catch {}

  return fp;
}

function parseJWTExp(jwt) {
  try {
    const payload = JSON.parse(
      Buffer.from(jwt.split(".")[1], "base64url").toString(),
    );
    return typeof payload.exp === "number"
      ? payload.exp * 1000
      : Date.now() + 3600_000;
  } catch {
    return Date.now() + 3600_000;
  }
}

async function bootstrapJWT() {
  const res = await fetch(BOOTSTRAP_URL, {
    method: "POST",
    headers: {
      "Content-Type": "application/json",
      "User-Agent": "mimocode/0.1.1",
      Accept: "*/*",
    },
    body: JSON.stringify({ client: getFingerprint() }),
  });

  if (!res.ok) {
    throw new Error(
      `Bootstrap failed: ${res.status} ${await res.text().catch(() => "")}`,
    );
  }

  const data = await res.json();
  if (!data.jwt) throw new Error("Bootstrap response missing jwt");
  return { jwt: data.jwt, exp: parseJWTExp(data.jwt) };
}

async function getJWT() {
  if (cachedJWT && jwtExpiry - Date.now() > JWT_REFRESH_MARGIN) {
    return cachedJWT;
  }
  if (pendingBootstrap) {
    return (await pendingBootstrap).jwt;
  }

  pendingBootstrap = bootstrapJWT();
  try {
    const result = await pendingBootstrap;
    cachedJWT = result.jwt;
    jwtExpiry = result.exp;
    return result.jwt;
  } finally {
    pendingBootstrap = null;
  }
}

async function mimoFetch(url, options) {
  let jwt = await getJWT();

  const headersFor = (token) => {
    const headers = new Headers(options?.headers);
    headers.set("Authorization", `Bearer ${token}`);
    headers.set("X-Mimo-Source", "mimocode-cli-free");
    headers.set("X-Session-Affinity", FREE_SESSION_AFFINITY);
    headers.set("User-Agent", MIMOCODE_USER_AGENT);
    headers.set("Accept", "*/*");
    return headers;
  };

  let res = await fetch(url, { ...options, headers: headersFor(jwt) });
  if (res.status === 401 || res.status === 403) {
    cachedJWT = null;
    jwt = await getJWT();
    res = await fetch(url, { ...options, headers: headersFor(jwt) });
  }
  return res;
}

async function apiFetch(body) {
  if (!UPSTREAM_API_KEY) {
    throw new Error("MIMO_API_KEY or XIAOMI_API_KEY is required in api mode");
  }

  return fetch(`${API_BASE_URL}/chat/completions`, {
    method: "POST",
    headers: {
      "Content-Type": "application/json",
      Accept: body.stream ? "text/event-stream" : "application/json",
      Authorization: `Bearer ${UPSTREAM_API_KEY}`,
      "User-Agent": "HuggingMes-MiMo-Proxy/1.0",
    },
    body: JSON.stringify(body),
  });
}

function modelForUpstream(model) {
  if (!model || model === "mimo-auto" || model === PUBLIC_MODEL) {
    return UPSTREAM_MODEL;
  }
  return model;
}

function hasMimocodeSystemMessage(messages) {
  if (!Array.isArray(messages) || messages.length === 0) return false;
  const first = messages[0];
  if (first?.role !== "system") return false;
  if (typeof first.content === "string") {
    return first.content.includes("MiMoCode");
  }
  return false;
}

function prepareFreeBody(body) {
  const out = { ...body, model: "mimo-auto" };

  if (
    out.max_tokens === undefined &&
    Number.isFinite(Number(out.max_completion_tokens))
  ) {
    out.max_tokens = Number(out.max_completion_tokens);
    delete out.max_completion_tokens;
  }
  if (out.max_tokens === undefined && Number.isFinite(FREE_MAX_TOKENS)) {
    out.max_tokens = FREE_MAX_TOKENS;
  }
  if (out.temperature === undefined) {
    out.temperature = 1;
  }
  if (out.stream === true && out.stream_options === undefined) {
    out.stream_options = { include_usage: true };
  }
  if (
    Array.isArray(out.tools) &&
    out.tools.length > 0 &&
    out.tool_choice === undefined
  ) {
    out.tool_choice = "auto";
  }

  if (FREE_COMPAT_SYSTEM && Array.isArray(out.messages)) {
    out.messages = hasMimocodeSystemMessage(out.messages)
      ? out.messages
      : [{ role: "system", content: MIMOCODE_SYSTEM_PROMPT }, ...out.messages];
  }

  return out;
}

function prepareApiBody(body) {
  const out = { ...body, model: modelForUpstream(body.model) };

  if (
    out.max_completion_tokens === undefined &&
    Number.isFinite(Number(out.max_tokens))
  ) {
    out.max_completion_tokens = Number(out.max_tokens);
    delete out.max_tokens;
  }

  return out;
}

function writeJson(res, statusCode, payload) {
  res.writeHead(statusCode, { "Content-Type": "application/json" });
  res.end(JSON.stringify(payload));
}

function isAuthorized(req) {
  const authHeader = req.headers.authorization || "";
  return authHeader.startsWith("Bearer ") && authHeader.slice(7) === PROXY_KEY;
}

function modelEntry(id) {
  return {
    id,
    object: "model",
    created: Math.floor(Date.now() / 1000),
    owned_by: "xiaomi",
    permission: [],
    root: id,
    parent: null,
  };
}

function modelsResponse() {
  const ids = new Set([PUBLIC_MODEL, "mimo-auto"]);
  if (PROXY_MODE === "api") ids.add(UPSTREAM_MODEL);
  return {
    object: "list",
    data: [...ids].map(modelEntry),
  };
}

async function handleRequest(req, res) {
  res.setHeader("Access-Control-Allow-Origin", "*");
  res.setHeader("Access-Control-Allow-Methods", "GET, POST, OPTIONS");
  res.setHeader("Access-Control-Allow-Headers", "*");

  if (req.method === "OPTIONS") {
    res.writeHead(204);
    res.end();
    return;
  }

  const url = new URL(req.url, `http://${req.headers.host}`);

  if (!isAuthorized(req)) {
    writeJson(res, 401, {
      error: {
        message: "Invalid API key",
        type: "auth_error",
        code: "401",
      },
    });
    return;
  }

  if (url.pathname === "/health") {
    writeJson(res, 200, {
      status: "ok",
      mode: PROXY_MODE,
      model: PUBLIC_MODEL,
      upstream_model: UPSTREAM_MODEL,
      free: PROXY_MODE === "free",
    });
    return;
  }

  if (url.pathname === "/v1/models" && req.method === "GET") {
    writeJson(res, 200, modelsResponse());
    return;
  }

  if (url.pathname === "/v1/chat/completions" && req.method === "POST") {
    let body = "";
    for await (const chunk of req) body += chunk;

    let parsed;
    try {
      parsed = JSON.parse(body);
    } catch {
      writeJson(res, 400, {
        error: { message: "Invalid JSON body", type: "parse_error" },
      });
      return;
    }

    const upstreamBody =
      PROXY_MODE === "api" ? prepareApiBody(parsed) : prepareFreeBody(parsed);
    const isStreaming = upstreamBody.stream === true;

    try {
      const upstream =
        PROXY_MODE === "api"
          ? await apiFetch(upstreamBody)
          : await mimoFetch(`${CHAT_BASE}/chat`, {
              method: "POST",
              headers: { "Content-Type": "application/json" },
              body: JSON.stringify(upstreamBody),
            });

      if (!upstream.ok) {
        const errBody = await upstream.text().catch(() => "");
        writeJson(res, upstream.status, {
          error: {
            message: `Upstream error: ${upstream.status}`,
            type: "upstream_error",
            mode: PROXY_MODE,
            param: errBody.slice(0, 500),
          },
        });
        return;
      }

      if (isStreaming) {
        res.writeHead(200, {
          "Content-Type": "text/event-stream",
          "Cache-Control": "no-cache",
          Connection: "keep-alive",
        });

        const reader = upstream.body.getReader();
        const decoder = new TextDecoder();
        while (true) {
          const { done, value } = await reader.read();
          if (done) break;
          res.write(decoder.decode(value, { stream: true }));
        }
        res.end();
      } else {
        const data = await upstream.text();
        res.writeHead(200, { "Content-Type": "application/json" });
        res.end(data);
      }
    } catch (err) {
      writeJson(res, 502, {
        error: { message: err.message, type: "proxy_error" },
      });
    }
    return;
  }

  writeJson(res, 404, {
    error: { message: `Not found: ${url.pathname}`, type: "not_found" },
  });
}

const server = http.createServer(handleRequest);
server.keepAliveTimeout = 65_000;
server.listen(PORT, HOST, () => {
  console.log(`MiMo proxy listening on ${HOST}:${PORT}`);
  console.log(`MiMo proxy base URL: http://${HOST}:${PORT}/v1`);
  console.log(`MiMo proxy mode: ${PROXY_MODE}`);
  console.log(`MiMo proxy model: ${PUBLIC_MODEL}`);
  if (PROXY_MODE === "api") {
    console.log(`MiMo proxy upstream: ${API_BASE_URL}`);
  }
});