export interface Span {
  file: string;
  start: number;
  end: number;
}

export const enum Tok {
  Eof = "eof",
  Newline = "\n",
  Ident = "ident",
  Int = "int",
  Float = "float",
  StrStart = "strStart",
  StrPart = "strPart",
  StrEnd = "strEnd",
  Char = "char",
  KwLet = "kwLet",
  KwMut = "kwMut",
  KwFun = "kwFun",
  KwAsync = "kwAsync",
  KwAwait = "kwAwait",
  KwIf = "kwIf",
  KwElse = "kwElse",
  KwMatch = "kwMatch",
  KwFor = "kwFor",
  KwIn = "kwIn",
  KwWhile = "kwWhile",
  KwReturn = "kwReturn",
  KwRaise = "kwRaise",
  KwTry = "kwTry",
  KwCatch = "kwCatch",
  KwStruct = "kwStruct",
  KwEnum = "kwEnum",
  KwType = "kwType",
  KwContract = "kwContract",
  KwSchema = "kwSchema",
  KwClass = "kwClass",
  KwImport = "kwImport",
  KwFrom = "kwFrom",
  KwExport = "kwExport",
  KwTest = "kwTest",
  KwIs = "kwIs",
  KwAs = "kwAs",
  KwAnd = "kwAnd",
  KwNot = "kwNot",
  KwTrue = "kwTrue",
  KwFalse = "kwFalse",
  KwUndefined = "kwUndefined",
  KwOk = "kwOk",
  KwErr = "kwErr",
  KwSome = "kwSome",
  KwNone = "kwNone",
  KwWith = "kwWith",
  KwOr = "kwOr",
  Plus = "+",
  Minus = "-",
  Star = "*",
  Slash = "/",
  Percent = "%",
  EqEq = "==",
  NotEq = "!=",
  Lt = "<",
  LtEq = "<=",
  Gt = ">",
  GtEq = ">=",
  AndAnd = "&&",
  OrOr = "||",
  Bang = "!",
  Assign = "=",
  PlusEq = "+=",
  MinusEq = "-=",
  StarEq = "*=",
  SlashEq = "/=",
  PercentEq = "%=",
  Dot = ".",
  DotDot = "..",
  DotDotEq = "..=",
  QDot = "?.",
  Arrow = "->",
  Question = "?",
  InterpEnd = "interpEnd",
  LParen = "(",
  RParen = ")",
  LBracket = "[",
  RBracket = "]",
  LBrace = "{",
  RBrace = "}",
  Comma = ",",
  Colon = ":",
  Pipe = "|",
  At = "@",
}

export interface Token {
  kind: Tok;
  text: string;
  value?: string | number | boolean;
  span: Span;
}

const KEYWORDS: Record<string, Tok> = {
  let: Tok.KwLet,
  mut: Tok.KwMut,
  fun: Tok.KwFun,
  async: Tok.KwAsync,
  await: Tok.KwAwait,
  if: Tok.KwIf,
  else: Tok.KwElse,
  match: Tok.KwMatch,
  for: Tok.KwFor,
  in: Tok.KwIn,
  while: Tok.KwWhile,
  return: Tok.KwReturn,
  raise: Tok.KwRaise,
  try: Tok.KwTry,
  catch: Tok.KwCatch,
  struct: Tok.KwStruct,
  enum: Tok.KwEnum,
  type: Tok.KwType,
  contract: Tok.KwContract,
  schema: Tok.KwSchema,
  class: Tok.KwClass,
  import: Tok.KwImport,
  from: Tok.KwFrom,
  export: Tok.KwExport,
  test: Tok.KwTest,
  is: Tok.KwIs,
  as: Tok.KwAs,
  and: Tok.KwAnd,
  not: Tok.KwNot,
  true: Tok.KwTrue,
  false: Tok.KwFalse,
  undefined: Tok.KwUndefined,
  Ok: Tok.KwOk,
  Err: Tok.KwErr,
  Some: Tok.KwSome,
  None: Tok.KwNone,
  with: Tok.KwWith,
  or: Tok.KwOr,
};

export class LexError extends Error {
  constructor(
    message: string,
    public span: Span,
  ) {
    super(message);
  }
}

const TWO_CHAR_OPS = new Set([
  Tok.EqEq, Tok.NotEq, Tok.LtEq, Tok.GtEq, Tok.AndAnd, Tok.OrOr,
  Tok.PlusEq, Tok.MinusEq, Tok.StarEq, Tok.SlashEq, Tok.PercentEq,
  Tok.DotDot, Tok.QDot, Tok.Arrow,
]);

export function lex(source: string, file: string): Token[] {
  const tokens: Token[] = [];
  let i = 0;

  const span = (start: number): Span => ({ file, start, end: i });
  const push = (kind: Tok, text: string, start: number, value?: string | number | boolean) => {
    tokens.push({ kind, text, value, span: span(start) });
  };
  function err(msg: string, at: number): never {
    throw new LexError(msg, { file, start: at, end: at });
  }
  const peek = (offset = 0) => source[i + offset] ?? "";

  // String state machine: when strCtx is active we lex string content.
  // Nested strings (inside interpolations) push onto strStack.
  let strBuf: string | null = null;
  let strBufStart = 0;
  let interpDepth = 0;
  interface StrCtx {
    depth: number;
    buf: string;
    bufStart: number;
  }
  const strStack: StrCtx[] = [];

  const enterString = (triple: boolean, start: number) => {
    push(Tok.StrStart, triple ? '"""' : '"', start);
    strBuf = "";
    strBufStart = i;
  };

  const lexStringContent = (): boolean => {
    // returns false when the string closed (and we should re-run main loop)
    while (i < source.length) {
      const ch = source[i]!;
      if (ch === "\\" && strTriple === false) {
        const n = source[i + 1];
        if (n === undefined) err("unterminated escape sequence", i);
        if (n === "n") strBuf += "\n";
        else if (n === "t") strBuf += "\t";
        else if (n === "r") strBuf += "\r";
        else if (n === '"') strBuf += '"';
        else if (n === "\\") strBuf += "\\";
        else if (n === "{") strBuf += "{";
        else if (n === "}") strBuf += "}";
        else if (n === "u") {
          if (source[i + 2] !== "{") err("expected \\u{XXXX} escape", i);
          const hexStart = i + 3;
          let k = hexStart;
          while (k < source.length && /[0-9a-fA-F]/.test(source[k]!)) k++;
          if (source[k] !== "}") err("expected \\u{XXXX} escape", i);
          const code = parseInt(source.slice(hexStart, k), 16);
          if (isNaN(code)) err("invalid unicode escape", i);
          strBuf += String.fromCodePoint(code);
          i = k + 1;
          continue;
        } else err(`unknown escape \\${n}`, i);
        i += 2;
        continue;
      }
      if (!strTriple && ch === "{") {
        if (strBuf!.length > 0) {
          push(Tok.StrPart, strBuf!, strBufStart, strBuf!);
          strBuf = "";
        }
        interpDepth = 1;
        i++;
        strBuf = null;
        return true; // switch to code mode (inside interpolation)
      }
      const closing = strTriple
        ? ch === '"' && peek(1) === '"' && peek(2) === '"'
        : ch === '"';
      if (closing) {
        if (strBuf!.length > 0) {
          push(Tok.StrPart, strBuf!, strBufStart, strBuf!);
          strBuf = "";
        }
        i += strTriple ? 3 : 1;
        push(Tok.StrEnd, strTriple ? '"""' : '"', i - (strTriple ? 3 : 1));
        // restore parent context if any: the parent was inside an interpolation
        const parent = strStack.pop();
        if (parent) {
          interpDepth = parent.depth;
          strBuf = null;
          strTriple = false;
          return true; // continue in code mode of the parent interpolation
        }
        strBuf = null;
        return false;
      }
      strBuf += ch;
      i++;
    }
    err("unterminated string", strBufStart);
  };

  let strTriple = false;

  while (i < source.length) {
    const c = source[i]!;
    const start = i;

    // 1. string content mode
    if (strBuf !== null) {
      const wentInterp = lexStringContent();
      if (wentInterp) continue; // now in code mode with interpDepth > 0
      continue;
    }

    // 2. code mode inside an interpolation
    if (interpDepth > 0) {
      if (c === "{") {
        interpDepth++;
        i++;
        push(Tok.LBrace, "{", start);
        continue;
      }
      if (c === "}") {
        interpDepth--;
        i++;
        if (interpDepth === 0) {
          // back inside the string: resume string content
          push(Tok.InterpEnd, "}", start);
          strBuf = "";
          strBufStart = i;
          strTriple = false;
        } else {
          push(Tok.RBrace, "}", start);
        }
        continue;
      }
      if (c === '"') {
        if (peek(1) === '"' && peek(2) === '"') {
          err("triple-quoted strings cannot appear inside interpolations", i);
        }
        // nested string: save current context
        strStack.push({ depth: interpDepth, buf: "", bufStart: i });
        i++;
        strTriple = false;
        enterString(false, start);
        continue;
      }
      // otherwise fall through to normal token lexing
    }

    // 3. whitespace / comments
    if (c === " " || c === "\t" || c === "\r") {
      i++;
      continue;
    }
    if (c === "\n") {
      i++;
      push(Tok.Newline, "\n", start);
      continue;
    }
    if (c === "/" && peek(1) === "/") {
      while (i < source.length && source[i] !== "\n") i++;
      continue;
    }
    if (c === "/" && peek(1) === "*") {
      i += 2;
      let depth = 1;
      while (i < source.length && depth > 0) {
        if (source[i] === "/" && peek(1) === "*") {
          depth++;
          i++;
        } else if (source[i] === "*" && peek(1) === "/") {
          depth--;
          i++;
        }
        i++;
      }
      if (depth > 0) err("unterminated block comment", start);
      continue;
    }

    // 4. identifiers / keywords
    if (/[A-Za-z_]/.test(c)) {
      let j = i + 1;
      while (j < source.length && /[A-Za-z0-9_]/.test(source[j]!)) j++;
      const text = source.slice(i, j);
      i = j;
      const kw = KEYWORDS[text];
      if (kw) push(kw, text, start);
      else push(Tok.Ident, text, start);
      continue;
    }

    // 5. numbers
    if (/[0-9]/.test(c)) {
      let j = i;
      while (j < source.length && /[0-9_]/.test(source[j]!)) j++;
      let isFloat = false;
      if (source[j] === "." && /[0-9]/.test(peek(j - i + 1))) {
        isFloat = true;
        j++;
        while (j < source.length && /[0-9_]/.test(source[j]!)) j++;
      }
      if (!isFloat && (source[j] === "e" || source[j] === "E")) {
        let k = j + 1;
        if (source[k] === "+" || source[k] === "-") k++;
        if (/[0-9]/.test(source[k] ?? "")) {
          isFloat = true;
          j = k;
          while (j < source.length && /[0-9_]/.test(source[j]!)) j++;
        }
      }
      const text = source.slice(i, j).replace(/_/g, "");
      i = j;
      if (isFloat) push(Tok.Float, text, start, parseFloat(text));
      else push(Tok.Int, text, start, parseInt(text, 10));
      continue;
    }

    // 6. strings
    if (c === '"') {
      strTriple = peek(1) === '"' && peek(2) === '"';
      if (strTriple) i += 3;
      else i++;
      enterString(strTriple, start);
      continue;
    }

    // 7. chars
    if (c === "'") {
      i++;
      let value: string;
      if (source[i] === "\\") {
        const n = source[i + 1];
        if (n === "n") value = "\n";
        else if (n === "t") value = "\t";
        else if (n === "r") value = "\r";
        else if (n === "'") value = "'";
        else if (n === "\\") value = "\\";
        else if (n === "u") {
          if (source[i + 2] !== "{") err("expected \\u{XXXX} escape", i);
          const hexStart = i + 3;
          let k = hexStart;
          while (k < source.length && /[0-9a-fA-F]/.test(source[k]!)) k++;
          if (source[k] !== "}") err("expected \\u{XXXX} escape", i);
          const code = parseInt(source.slice(hexStart, k), 16);
          if (isNaN(code)) err("invalid unicode escape", i);
          value = String.fromCodePoint(code);
          i = k + 1;
        } else err(`unknown escape \\${n}`, i);
        i += 2;
      } else {
        if (i >= source.length) err("unterminated char literal", start);
        value = source[i]!;
        i++;
      }
      if (source[i] !== "'") err("char literal must contain exactly one character", start);
      i++;
      push(Tok.Char, `'${value}'`, start, value);
      continue;
    }

    // 8. operators
    const three = c + peek(1) + peek(2);
    const two = c + peek(1);
    let kind: Tok | null = null;
    if (three === "..=") kind = Tok.DotDotEq;
    else if (two === "==") kind = Tok.EqEq;
    else if (two === "!=") kind = Tok.NotEq;
    else if (two === "<=") kind = Tok.LtEq;
    else if (two === ">=") kind = Tok.GtEq;
    else if (two === "&&") kind = Tok.AndAnd;
    else if (two === "||") kind = Tok.OrOr;
    else if (two === "+=") kind = Tok.PlusEq;
    else if (two === "-=") kind = Tok.MinusEq;
    else if (two === "*=") kind = Tok.StarEq;
    else if (two === "/=") kind = Tok.SlashEq;
    else if (two === "%=") kind = Tok.PercentEq;
    else if (two === "..") kind = Tok.DotDot;
    else if (two === "?.") kind = Tok.QDot;
    else if (two === "->") kind = Tok.Arrow;
    else if (c === "+") kind = Tok.Plus;
    else if (c === "-") kind = Tok.Minus;
    else if (c === "*") kind = Tok.Star;
    else if (c === "/") kind = Tok.Slash;
    else if (c === "%") kind = Tok.Percent;
    else if (c === "<") kind = Tok.Lt;
    else if (c === ">") kind = Tok.Gt;
    else if (c === "!") kind = Tok.Bang;
    else if (c === "=") kind = Tok.Assign;
    else if (c === ".") kind = Tok.Dot;
    else if (c === "?") kind = Tok.Question;
    else if (c === "(") kind = Tok.LParen;
    else if (c === ")") kind = Tok.RParen;
    else if (c === "[") kind = Tok.LBracket;
    else if (c === "]") kind = Tok.RBracket;
    else if (c === "{") kind = Tok.LBrace;
    else if (c === "}") kind = Tok.RBrace;
    else if (c === ",") kind = Tok.Comma;
    else if (c === ":") kind = Tok.Colon;
    else if (c === "|") kind = Tok.Pipe;
    else if (c === "@") kind = Tok.At;
    if (kind) {
      const len = kind === Tok.DotDotEq ? 3 : TWO_CHAR_OPS.has(kind) ? 2 : 1;
      i += len;
      push(kind, source.slice(start, i), start);
      continue;
    }

    err(`unexpected character '${c}'`, i);
  }

  if (strBuf !== null) err("unterminated string", strBufStart);
  tokens.push({ kind: Tok.Eof, text: "", span: { file, start: i, end: i } });
  return tokens;
}