// The small parts of the Python standard library games reach for: re, copy, time, datetime,
// collections.defaultdict and functools.partial. Everything is pure computation.

import { PyDict, PyInstance, pyStr } from "./python";

/** dict that creates missing entries from a factory (collections.defaultdict). */
export class PyDefaultDict extends PyDict {
  constructor(public factory: (() => any) | null, entries?: Iterable<[any, any]>) {
    super(entries);
  }
  missing(k: any) {
    const value = this.factory ? this.factory() : null;
    this.set(k, value);
    return value;
  }
}

class ReMatch {
  constructor(private m: RegExpExecArray) {}
  group(...which: (number | string)[]) {
    const pick = (w: number | string) => {
      const v = typeof w === "string" ? this.m.groups?.[w] : this.m[w];
      return v ?? null;
    };
    if (!which.length) return this.m[0];
    return which.length === 1 ? pick(which[0]) : which.map(pick);
  }
  groups() {
    return this.m.slice(1).map((v) => v ?? null);
  }
  groupdict() {
    return new PyDict(Object.entries(this.m.groups ?? {}));
  }
  start() {
    return this.m.index;
  }
  end() {
    return this.m.index + this.m[0].length;
  }
}

const FLAGS: Record<number, string> = { 2: "i", 8: "m", 16: "s" };

function compile(pattern: string, flags = 0, global = false): RegExp {
  const js = pattern.replace(/\(\?P<(\w+)>/g, "(?<$1>").replace(/\(\?P=(\w+)\)/g, "\\k<$1>").replace(/\\Z/g, "$");
  const letters = Object.entries(FLAGS).filter(([bit]) => flags & Number(bit)).map(([, f]) => f).join("");
  return new RegExp(js, letters + (global ? "g" : "") + "u");
}

function expand(replacement: string): string {
  // \1 and \g<name> are the group references of Python templates
  return replacement.replace(/\\(\d+)/g, "$$$1").replace(/\\g<(\w+)>/g, "$$<$1>");
}

function deepcopy(v: any, seen = new Map<any, any>()): any {
  if (v === null || typeof v !== "object") return v;
  if (seen.has(v)) return seen.get(v);
  if (Array.isArray(v)) {
    const out: any[] = [];
    seen.set(v, out);
    v.forEach((x) => out.push(deepcopy(x, seen)));
    return out;
  }
  if (v instanceof PyDefaultDict) {
    const out = new PyDefaultDict(v.factory);
    seen.set(v, out);
    v.map.forEach((x, k) => out.map.set(k, deepcopy(x, seen)));
    return out;
  }
  if (v instanceof PyDict) {
    const out = new PyDict();
    seen.set(v, out);
    v.map.forEach((x, k) => out.map.set(k, deepcopy(x, seen)));
    return out;
  }
  if (v instanceof PyInstance) {
    const attrs: Record<string, any> = Object.create(null);
    const out = new PyInstance(v.cls, attrs);
    seen.set(v, out);
    for (const [k, x] of Object.entries(v.attrs)) attrs[k] = deepcopy(x, seen);
    return out;
  }
  return v;
}

const pad = (n: number, width = 2) => String(n).padStart(width, "0");

class DateTime {
  constructor(public d: Date) {}
  get year() { return this.d.getFullYear(); }
  get month() { return this.d.getMonth() + 1; }
  get day() { return this.d.getDate(); }
  get hour() { return this.d.getHours(); }
  get minute() { return this.d.getMinutes(); }
  get second() { return this.d.getSeconds(); }
  weekday() { return (this.d.getDay() + 6) % 7; }
  strftime(format: string) { return strftime(format, this.d); }
  isoformat() { return `${this.year}-${pad(this.month)}-${pad(this.day)}T${pad(this.hour)}:${pad(this.minute)}:${pad(this.second)}`; }
}

export function strftime(format: string, d: Date): string {
  const days = ["Sunday", "Monday", "Tuesday", "Wednesday", "Thursday", "Friday", "Saturday"];
  const months = ["January", "February", "March", "April", "May", "June", "July", "August", "September", "October", "November", "December"];
  return format.replace(/%([a-zA-Z%])/g, (_m, c: string) => {
    switch (c) {
      case "Y": return String(d.getFullYear());
      case "y": return pad(d.getFullYear() % 100);
      case "m": return pad(d.getMonth() + 1);
      case "d": return pad(d.getDate());
      case "H": return pad(d.getHours());
      case "I": return pad(d.getHours() % 12 || 12);
      case "M": return pad(d.getMinutes());
      case "S": return pad(d.getSeconds());
      case "p": return d.getHours() < 12 ? "AM" : "PM";
      case "A": return days[d.getDay()];
      case "a": return days[d.getDay()].slice(0, 3);
      case "B": return months[d.getMonth()];
      case "b": return months[d.getMonth()].slice(0, 3);
      case "%": return "%";
      default: return "";
    }
  });
}

const now = () => new DateTime(new Date());
const DAY_MS = 86_400_000;

/** `datetime.timedelta`: kept in milliseconds. */
class TimeDelta {
  constructor(public ms: number) {}
  get days() { return Math.floor(this.ms / DAY_MS); }
  get seconds() { return Math.floor((this.ms - this.days * DAY_MS) / 1000); }
  total_seconds() { return this.ms / 1000; }
  __binop__(op: string, other: any) {
    if (op === "+" && other instanceof TimeDelta) return new TimeDelta(this.ms + other.ms);
    if (op === "-" && other instanceof TimeDelta) return new TimeDelta(this.ms - other.ms);
    if (op === "*" && typeof other === "number") return new TimeDelta(this.ms * other);
    throw new TypeError(`timedelta ${op}`);
  }
}

function timedelta(...a: any[]) {
  const last = a[a.length - 1];
  const kw = last && typeof last === "object" && last.__kw ? last.__kw : {};
  const unit: Record<string, number> = { days: DAY_MS, seconds: 1000, minutes: 60_000, hours: 3_600_000,
                                         weeks: 7 * DAY_MS, milliseconds: 1 };
  let ms = typeof a[0] === "number" ? a[0] * DAY_MS : 0;
  for (const [name, factor] of Object.entries(unit)) if (typeof kw[name] === "number") ms += kw[name] * factor;
  return new TimeDelta(ms);
}

/** `date(y, m, d)` and `datetime(y, m, d, ...)`: a DateTime at a given moment, with date arithmetic. */
const fromParts = (...a: any[]) => {
  const [y, m, d, h = 0, min = 0, s = 0] = a.filter((x) => typeof x === "number");
  const date = new Date(y, m - 1, d, h, min, s);
  if (date.getMonth() !== m - 1 || date.getDate() !== d) throw new Error("ValueError: day is out of range for month");
  return new DateTime(date);
};
const datetimeClass = { now, today: now, utcnow: now };
const dateClass = Object.assign(fromParts, datetimeClass);
Object.assign(DateTime.prototype, {
  __binop__(this: DateTime, op: string, other: any) {
    if (other instanceof TimeDelta && (op === "+" || op === "-")) {
      const moved = new Date(this.d.getTime());
      moved.setTime(moved.getTime() + (op === "+" ? other.ms : -other.ms));
      // whole days keep the wall-clock time across daylight-saving changes
      if (other.ms % DAY_MS === 0) {
        const shifted = new Date(this.d.getTime());
        shifted.setDate(shifted.getDate() + (op === "+" ? 1 : -1) * (other.ms / DAY_MS));
        return new DateTime(shifted);
      }
      return new DateTime(moved);
    }
    if (other instanceof DateTime && op === "-") return new TimeDelta(this.d.getTime() - other.d.getTime());
    throw new TypeError(`datetime ${op}`);
  },
});

function md5(bytes: Uint8Array): string {
  const K = Array.from({ length: 64 }, (_, i) => Math.floor(Math.abs(Math.sin(i + 1)) * 2 ** 32) >>> 0);
  const S = [7, 12, 17, 22, 5, 9, 14, 20, 4, 11, 16, 23, 6, 10, 15, 21];
  let a0 = 0x67452301, b0 = 0xefcdab89, c0 = 0x98badcfe, d0 = 0x10325476;
  const padded = new Uint8Array(((bytes.length + 8) >> 6 << 6) + 64);
  padded.set(bytes);
  padded[bytes.length] = 0x80;
  const view = new DataView(padded.buffer);
  view.setUint32(padded.length - 8, (bytes.length * 8) >>> 0, true);
  view.setUint32(padded.length - 4, Math.floor(bytes.length * 8 / 2 ** 32), true);
  for (let off = 0; off < padded.length; off += 64) {
    const m = Array.from({ length: 16 }, (_, i) => view.getUint32(off + i * 4, true));
    let a = a0, b = b0, c = c0, d = d0;
    for (let i = 0; i < 64; i++) {
      let f: number, g: number;
      if (i < 16) { f = (b & c) | (~b & d); g = i; }
      else if (i < 32) { f = (d & b) | (~d & c); g = (5 * i + 1) % 16; }
      else if (i < 48) { f = b ^ c ^ d; g = (3 * i + 5) % 16; }
      else { f = c ^ (b | ~d); g = (7 * i) % 16; }
      f = (f + a + K[i] + m[g]) >>> 0;
      a = d; d = c; c = b;
      const s = S[(i >> 4) * 4 + (i % 4)];
      b = (b + ((f << s) | (f >>> (32 - s)))) >>> 0;
    }
    a0 = (a0 + a) >>> 0; b0 = (b0 + b) >>> 0; c0 = (c0 + c) >>> 0; d0 = (d0 + d) >>> 0;
  }
  return [a0, b0, c0, d0].map((v) => [0, 8, 16, 24].map((sh) => ((v >>> sh) & 255).toString(16).padStart(2, "0")).join("")).join("");
}

function sha256(bytes: Uint8Array): string {
  const K = Array.from({ length: 64 }, (_, i) => {
    let n = 2, count = 0;
    for (; count <= i; n++) if (Array.from({ length: n - 2 }, (_, d) => d + 2).every((d) => n % d)) count++;
    return Math.floor((Math.cbrt(n - 1) % 1) * 2 ** 32) >>> 0;
  });
  const h = [2, 3, 5, 7, 11, 13, 17, 19].map((p) => Math.floor((Math.sqrt(p) % 1) * 2 ** 32) >>> 0);
  const padded = new Uint8Array(((bytes.length + 8) >> 6 << 6) + 64);
  padded.set(bytes);
  padded[bytes.length] = 0x80;
  const view = new DataView(padded.buffer);
  view.setUint32(padded.length - 8, Math.floor(bytes.length * 8 / 2 ** 32));
  view.setUint32(padded.length - 4, (bytes.length * 8) >>> 0);
  const rotr = (x: number, n: number) => (x >>> n) | (x << (32 - n));
  for (let off = 0; off < padded.length; off += 64) {
    const w = Array.from({ length: 64 }, (_, i) => (i < 16 ? view.getUint32(off + i * 4) : 0));
    for (let i = 16; i < 64; i++) {
      const s0 = rotr(w[i - 15], 7) ^ rotr(w[i - 15], 18) ^ (w[i - 15] >>> 3);
      const s1 = rotr(w[i - 2], 17) ^ rotr(w[i - 2], 19) ^ (w[i - 2] >>> 10);
      w[i] = (w[i - 16] + s0 + w[i - 7] + s1) >>> 0;
    }
    let [a, b, c, d, e, f, g, hh] = h;
    for (let i = 0; i < 64; i++) {
      const t1 = (hh + (rotr(e, 6) ^ rotr(e, 11) ^ rotr(e, 25)) + ((e & f) ^ (~e & g)) + K[i] + w[i]) >>> 0;
      const t2 = ((rotr(a, 2) ^ rotr(a, 13) ^ rotr(a, 22)) + ((a & b) ^ (a & c) ^ (b & c))) >>> 0;
      hh = g; g = f; f = e; e = (d + t1) >>> 0; d = c; c = b; b = a; a = (t1 + t2) >>> 0;
    }
    [a, b, c, d, e, f, g, hh].forEach((v, i) => { h[i] = (h[i] + v) >>> 0; });
  }
  return h.map((v) => v.toString(16).padStart(8, "0")).join("");
}

/** `hashlib.md5(data)`: the data may come later through `update`. */
class Hash {
  private parts: Uint8Array[] = [];
  constructor(private digest: (bytes: Uint8Array) => string, data?: any) { if (data !== undefined) this.update(data); }
  update(data: any) {
    this.parts.push(typeof data === "string" ? new TextEncoder().encode(data) : Uint8Array.from(data ?? []));
    return null;
  }
  hexdigest() {
    const all = new Uint8Array(this.parts.reduce((n, p) => n + p.length, 0));
    let at = 0;
    for (const p of this.parts) { all.set(p, at); at += p.length; }
    return this.digest(all);
  }
}

/** Namespaces added to the sandbox's builtins. */
export function pythonModules(): Record<string, any> {
  const re = {
    I: 2, IGNORECASE: 2, M: 8, MULTILINE: 8, S: 16, DOTALL: 16,
    escape: (s: string) => String(s).replace(/[.*+?^${}()|[\]\\]/g, "\\$&"),
    compile: (pattern: string, flags = 0) => ({
      search: (s: string) => re.search(pattern, s, flags),
      match: (s: string) => re.match(pattern, s, flags),
      sub: (r: string, s: string) => re.sub(pattern, r, s, 0, flags),
      findall: (s: string) => re.findall(pattern, s, flags),
      fullmatch: (s: string) => re.fullmatch(pattern, s, flags),
      split: (s: string) => re.split(pattern, s, 0, flags),
    }),
    search: (pattern: string, s: string, flags = 0) => {
      const m = compile(pattern, flags).exec(String(s));
      return m ? new ReMatch(m) : null;
    },
    match: (pattern: string, s: string, flags = 0) => {
      const m = compile(`^(?:${pattern})`, flags).exec(String(s));
      return m ? new ReMatch(m) : null;
    },
    fullmatch: (pattern: string, s: string, flags = 0) => {
      const m = compile(`^(?:${pattern})$`, flags).exec(String(s));
      return m ? new ReMatch(m) : null;
    },
    findall: (pattern: string, s: string, flags = 0) =>
      [...String(s).matchAll(compile(pattern, flags, true))].map((m) => (m.length > 2 ? m.slice(1) : m[1] ?? m[0])),
    sub: (pattern: string, replacement: any, s: string, count = 0, flags = 0) => {
      let left = count || Infinity;
      const regex = compile(pattern, flags, true);
      return String(s).replace(regex, (...args: any[]) => {
        if (left-- <= 0) return args[0];
        if (typeof replacement === "function") return pyStr(replacement(new ReMatch(Object.assign(args.slice(0, -2).filter((x) => typeof x !== "object"), { index: args[args.length - 2], groups: undefined, input: s }) as any)));
        return String(replacement).replace(/\\(\d+)/g, (_m: string, n: string) => args[Number(n)] ?? "");
      });
    },
    split: (pattern: string, s: string, maxsplit = 0, flags = 0) => {
      const parts = String(s).split(compile(pattern, flags, true));
      return maxsplit ? [...parts.slice(0, maxsplit), parts.slice(maxsplit).join("")] : parts;
    },
  };
  void expand;
  const partial = (f: any, ...bound: any[]) => (...args: any[]) => f(...bound, ...args);
  const functionModules = { partial };
  return {
    re,
    copy: { copy: (v: any) => (Array.isArray(v) ? [...v] : v instanceof PyDict ? new PyDict(v.map.entries()) : v), deepcopy: (v: any) => deepcopy(v) },
    time: {
      time: () => Date.now() / 1000,
      monotonic: () => performance.now() / 1000,
      sleep: () => null,
      strftime: (format: string, t?: any) => strftime(format, t instanceof DateTime ? t.d : new Date()),
    },
    datetime: Object.assign({ datetime: dateClass, date: dateClass, timedelta }, datetimeClass),
    collections: {
      defaultdict: (factory?: any) => new PyDefaultDict(typeof factory === "function" ? () => factory() : null),
      OrderedDict: () => new PyDict(),
    },
    functools: functionModules,
    colorsys: {
      rgb_to_hls: (r: number, g: number, b: number) => {
        const max = Math.max(r, g, b), min = Math.min(r, g, b), l = (max + min) / 2, d = max - min;
        if (d === 0) return [0, l, 0];
        const s = l > 0.5 ? d / (2 - max - min) : d / (max + min);
        const h = max === r ? ((g - b) / d) % 6 : max === g ? (b - r) / d + 2 : (r - g) / d + 4;
        return [((h / 6) % 1 + 1) % 1, l, s];
      },
      hls_to_rgb: (h: number, l: number, s: number) => {
        if (s === 0) return [l, l, l];
        const q = l < 0.5 ? l * (1 + s) : l + s - l * s, p = 2 * l - q;
        const channel = (t: number) => {
          t = ((t % 1) + 1) % 1;
          return t < 1 / 6 ? p + (q - p) * 6 * t : t < 1 / 2 ? q : t < 2 / 3 ? p + (q - p) * (2 / 3 - t) * 6 : p;
        };
        return [channel(h + 1 / 3), channel(h), channel(h - 1 / 3)];
      },
      rgb_to_hsv: (r: number, g: number, b: number) => {
        const max = Math.max(r, g, b), min = Math.min(r, g, b), d = max - min;
        const h = d === 0 ? 0 : max === r ? ((g - b) / d) % 6 : max === g ? (b - r) / d + 2 : (r - g) / d + 4;
        return [((h / 6) % 1 + 1) % 1, max === 0 ? 0 : d / max, max];
      },
      hsv_to_rgb: (h: number, s: number, v: number) => {
        const i = Math.floor(h * 6), f = h * 6 - i, p = v * (1 - s), q = v * (1 - f * s), t = v * (1 - (1 - f) * s);
        return [[v, q, p, p, t, v][i % 6], [t, v, v, q, p, p][i % 6], [p, p, t, v, v, q][i % 6]];
      },
    },
    fnmatch: {
      fnmatch: (name: any, pattern: any) => new RegExp("^" + String(pattern).replace(/[.+^${}()|\\]/g, "\\$&")
        .replace(/\*/g, ".*").replace(/\?/g, ".") + "$", "i").test(String(name)),
    },
    traceback: { format_exc: () => "", print_exc: () => null },
    hashlib: {
      md5: (data?: any) => new Hash(md5, data),
      sha256: (data?: any) => new Hash(sha256, data),
    },
    uuid: {
      uuid4: () => ({ hex: crypto.randomUUID().replaceAll("-", ""), toString: () => crypto.randomUUID() }),
    },
    // `from collections import defaultdict` and friends leave the bare names
    defaultdict: (factory?: any) => new PyDefaultDict(typeof factory === "function" ? () => factory() : null),
    OrderedDict: () => new PyDict(),
    deepcopy: (v: any) => deepcopy(v),
    partial,
    timedelta,
  };
}