import { Lexer, LexError } from "./tokenizer.ts";

export type FunctionRuntime = "legacy" | "client" | "server" | "shared";
export type StateRuntime = "shared" | "client" | "server";
export type StoreKind = "global" | "page";

export interface FunctionParameterDecl {
  name: string;
  optional: boolean;
  valueType?: string;
  default?: string;
}

export interface RuntimeFunctionDecl {
  name: string;
  runtime: FunctionRuntime;
  async: boolean;
  parameters: FunctionParameterDecl[];
  returnType?: string;
  body: string;
  source: string;
}

export interface OutputDecl {
  name: string;
  payload?: {
    name: string;
    valueType: string;
    optional: boolean;
  };
}

export interface StructuredImportDecl {
  source: string;
  typeOnly: boolean;
  defaultImport?: string;
  namespaceImport?: string;
  namedImports: Array<{ imported: string; local: string; typeOnly: boolean }>;
  raw: string;
}

export interface PersistDecl {
  storage: "memory" | "session" | "local";
  include: string[];
  version: number;
  migrations?: string;
  validation?: string;
}

export interface StoreLifecycleDecl {
  serverInit?: string;
  clientInit?: string;
  hydrate?: string;
  dispose?: string;
}

function splitTopLevel(input: string, separator = ","): string[] {
  const parts: string[] = [];
  let start = 0;
  let quote: string | null = null;
  let angle = 0;
  let square = 0;
  let brace = 0;
  let paren = 0;
  for (let i = 0; i < input.length; i++) {
    const c = input[i]!;
    if (quote) {
      if (c === "\\") i++;
      else if (c === quote) quote = null;
      continue;
    }
    if (c === '"' || c === "'" || c === "`") {
      quote = c;
      continue;
    }
    if (c === "<") angle++;
    else if (c === ">" && angle > 0) angle--;
    else if (c === "[") square++;
    else if (c === "]" && square > 0) square--;
    else if (c === "{") brace++;
    else if (c === "}" && brace > 0) brace--;
    else if (c === "(") paren++;
    else if (c === ")" && paren > 0) paren--;
    else if (c === separator && angle === 0 && square === 0 && brace === 0 && paren === 0) {
      parts.push(input.slice(start, i).trim());
      start = i + 1;
    }
  }
  const tail = input.slice(start).trim();
  if (tail) parts.push(tail);
  return parts;
}

function findTopLevelChar(input: string, wanted: string): number {
  let quote: string | null = null;
  let angle = 0;
  let square = 0;
  let brace = 0;
  let paren = 0;
  for (let i = 0; i < input.length; i++) {
    const c = input[i]!;
    if (quote) {
      if (c === "\\") i++;
      else if (c === quote) quote = null;
      continue;
    }
    if (c === '"' || c === "'" || c === "`") {
      quote = c;
      continue;
    }
    if (c === "<") angle++;
    else if (c === ">" && angle > 0) angle--;
    else if (c === "[") square++;
    else if (c === "]" && square > 0) square--;
    else if (c === "{") brace++;
    else if (c === "}" && brace > 0) brace--;
    else if (c === "(") paren++;
    else if (c === ")" && paren > 0) paren--;
    if (c === wanted && angle === 0 && square === 0 && brace === 0 && paren === 0) return i;
  }
  return -1;
}

function parseParameters(source: string): FunctionParameterDecl[] {
  return splitTopLevel(source)
    .filter(Boolean)
    .map((entry) => {
      const eq = findTopLevelChar(entry, "=");
      const declaration = (eq >= 0 ? entry.slice(0, eq) : entry).trim();
      const defaultValue = eq >= 0 ? entry.slice(eq + 1).trim() : undefined;
      const colon = findTopLevelChar(declaration, ":");
      const rawName = (colon >= 0 ? declaration.slice(0, colon) : declaration).trim();
      const optional = rawName.endsWith("?");
      const name = optional ? rawName.slice(0, -1).trim() : rawName;
      const valueType = colon >= 0 ? declaration.slice(colon + 1).trim() : undefined;
      return {
        name,
        optional,
        ...(valueType ? { valueType } : {}),
        ...(defaultValue ? { default: defaultValue } : {}),
      };
    });
}

function skipTrivia(source: string, start: number): number {
  let i = start;
  while (i < source.length) {
    if (/\s/.test(source[i]!)) {
      i++;
      continue;
    }
    if (source.startsWith("//", i)) {
      const end = source.indexOf("\n", i + 2);
      i = end < 0 ? source.length : end + 1;
      continue;
    }
    if (source.startsWith("/*", i)) {
      const end = source.indexOf("*/", i + 2);
      i = end < 0 ? source.length : end + 2;
      continue;
    }
    break;
  }
  return i;
}

function readWord(source: string, start: number): { word: string; end: number } | null {
  const match = /^[A-Za-z_$][\w$]*/.exec(source.slice(start));
  return match ? { word: match[0], end: start + match[0].length } : null;
}

function readBalanced(
  source: string,
  start: number,
  open: string,
  close: string,
): { inner: string; end: number } {
  if (source[start] !== open) throw new Error(`Expected '${open}' at offset ${start}`);
  let depth = 0;
  let quote: string | null = null;
  for (let i = start; i < source.length; i++) {
    const c = source[i]!;
    if (quote) {
      if (c === "\\") i++;
      else if (c === quote) quote = null;
      continue;
    }
    if (c === '"' || c === "'" || c === "`") {
      quote = c;
      continue;
    }
    if (c === open) depth++;
    else if (c === close && --depth === 0) return { inner: source.slice(start + 1, i), end: i + 1 };
  }
  throw new Error(`Unbalanced '${open}${close}' starting at offset ${start}`);
}

export function parseRuntimeFunctions(source: string): RuntimeFunctionDecl[] {
  const declarations: RuntimeFunctionDecl[] = [];
  let i = 0;
  while (i < source.length) {
    i = skipTrivia(source, i);
    const start = i;
    let token = readWord(source, i);
    if (!token) {
      i++;
      continue;
    }
    let runtime: FunctionRuntime = "legacy";
    if (["client", "server", "shared"].includes(token.word)) {
      runtime = token.word as FunctionRuntime;
      i = skipTrivia(source, token.end);
      token = readWord(source, i);
      if (!token) continue;
    }
    let isAsync = false;
    if (token.word === "async") {
      isAsync = true;
      i = skipTrivia(source, token.end);
      token = readWord(source, i);
      if (!token) continue;
    }
    if (token.word !== "function") {
      i = token.end;
      continue;
    }
    i = skipTrivia(source, token.end);
    const nameToken = readWord(source, i);
    if (!nameToken) throw new Error(`Expected function name at offset ${i}`);
    const name = nameToken.word;
    i = skipTrivia(source, nameToken.end);
    const params = readBalanced(source, i, "(", ")");
    i = skipTrivia(source, params.end);
    let returnType: string | undefined;
    if (source[i] === ":") {
      i++;
      const typeStart = i;
      let quote: string | null = null;
      let angle = 0;
      let square = 0;
      let paren = 0;
      while (i < source.length) {
        const c = source[i]!;
        if (quote) {
          if (c === "\\") i++;
          else if (c === quote) quote = null;
          i++;
          continue;
        }
        if (c === '"' || c === "'" || c === "`") quote = c;
        else if (c === "<") angle++;
        else if (c === ">" && angle > 0) angle--;
        else if (c === "[") square++;
        else if (c === "]" && square > 0) square--;
        else if (c === "(") paren++;
        else if (c === ")" && paren > 0) paren--;
        else if (c === "{" && angle === 0 && square === 0 && paren === 0) break;
        i++;
      }
      returnType = source.slice(typeStart, i).trim();
    }
    i = skipTrivia(source, i);
    const body = readBalanced(source, i, "{", "}");
    i = body.end;
    declarations.push({
      name,
      runtime,
      async: isAsync,
      parameters: parseParameters(params.inner),
      ...(returnType ? { returnType } : {}),
      body: body.inner,
      source: source.slice(start, body.end).trim(),
    });
  }
  return declarations;
}

export function stripRuntimeFunctionModifiers(source: string, include: FunctionRuntime[]): string {
  const allowed = new Set(include);
  return parseRuntimeFunctions(source)
    .filter((entry) => allowed.has(entry.runtime))
    .map((entry) => {
      const params = entry.parameters
        .map(
          (param) =>
            `${param.name}${param.optional ? "?" : ""}${param.valueType ? `: ${param.valueType}` : ""}${param.default ? ` = ${param.default}` : ""}`,
        )
        .join(", ");
      return `${entry.async ? "async " : ""}function ${entry.name}(${params})${entry.returnType ? `: ${entry.returnType}` : ""} {${entry.body}}`;
    })
    .join("\n\n");
}

export function parseOutputs(source: string): OutputDecl[] {
  const out: OutputDecl[] = [];
  let i = 0;
  while (i < source.length) {
    i = skipTrivia(source, i);
    if (i >= source.length) break;
    const nameToken = readWord(source, i);
    if (!nameToken) throw new Error(`Expected output name at offset ${i}`);
    i = skipTrivia(source, nameToken.end);
    const args = readBalanced(source, i, "(", ")");
    i = args.end;
    const parameters = parseParameters(args.inner);
    if (parameters.length > 1)
      throw new Error(`Output '${nameToken.word}' accepts zero or one payload`);
    const payload = parameters[0];
    if (payload && !payload.valueType)
      throw new Error(`Output '${nameToken.word}' payload requires a type`);
    out.push({
      name: nameToken.word,
      ...(payload
        ? {
            payload: {
              name: payload.name,
              valueType: payload.valueType!,
              optional: payload.optional,
            },
          }
        : {}),
    });
  }
  return out;
}

export function parseStructuredImports(imports: string[]): StructuredImportDecl[] {
  return imports.map((raw) => {
    const sourceMatch = /\sfrom\s+["']([^"']+)["']|^import\s+["']([^"']+)["']/.exec(raw);
    const source = sourceMatch?.[1] ?? sourceMatch?.[2] ?? "";
    const typeOnly = /^import\s+type\b/.test(raw);
    const clause = raw
      .replace(/^import\s+(?:type\s+)?/, "")
      .replace(/\s+from\s+["'][^"']+["']\s*;?$/, "")
      .trim();
    const declaration: StructuredImportDecl = { source, typeOnly, namedImports: [], raw };
    if (!clause || clause.startsWith('"') || clause.startsWith("'")) return declaration;
    if (clause.startsWith("*")) {
      declaration.namespaceImport = /\*\s+as\s+([A-Za-z_$][\w$]*)/.exec(clause)?.[1];
      return declaration;
    }
    let rest = clause;
    if (!rest.startsWith("{")) {
      const comma = findTopLevelChar(rest, ",");
      declaration.defaultImport = (comma < 0 ? rest : rest.slice(0, comma)).trim();
      rest = comma < 0 ? "" : rest.slice(comma + 1).trim();
    }
    const named = /^\{([\s\S]*)\}$/.exec(rest)?.[1];
    if (named !== undefined) {
      declaration.namedImports = splitTopLevel(named).map((item) => {
        const localTypeOnly = /^type\s+/.test(item);
        const cleaned = item.replace(/^type\s+/, "").trim();
        const [imported, local] = cleaned.split(/\s+as\s+/);
        return {
          imported: imported!.trim(),
          local: (local ?? imported)!.trim(),
          typeOnly: typeOnly || localTypeOnly,
        };
      });
    }
    return declaration;
  });
}

export function parseStateDeclarations(
  source: string,
  runtime: StateRuntime,
): Array<{ name: string; valueType?: string; expr: string; runtime: StateRuntime }> {
  const lx = new Lexer(source);
  const out: Array<{ name: string; valueType?: string; expr: string; runtime: StateRuntime }> = [];
  while (lx.peek().type !== "eof") {
    const nameToken = lx.next();
    if (nameToken.type !== "ident")
      throw new LexError(`Expected a state name at offset ${nameToken.pos}`);
    let valueType: string | undefined;
    if (lx.peek().type === "colon") {
      lx.next();
      const annotation = lx.readTypeAnnotation();
      valueType = annotation.type;
      if (!annotation.hasDefault)
        throw new LexError(`State '${nameToken.value}' requires an initializer`);
    } else {
      const eq = lx.next();
      if (eq.type !== "eq")
        throw new LexError(`Expected '=' after state '${nameToken.value}' at offset ${eq.pos}`);
    }
    out.push({
      name: nameToken.value,
      ...(valueType ? { valueType } : {}),
      expr: lx.readPropInitializer(),
      runtime,
    });
  }
  return out;
}

export function parseComputedDeclarations(
  source: string,
): Array<{ name: string; valueType?: string; expr: string }> {
  const lx = new Lexer(source);
  const out: Array<{ name: string; valueType?: string; expr: string }> = [];
  while (lx.peek().type !== "eof") {
    const nameToken = lx.next();
    if (nameToken.type !== "ident")
      throw new LexError(`Expected a computed name at offset ${nameToken.pos}`);
    let valueType: string | undefined;
    if (lx.peek().type === "colon") {
      lx.next();
      const annotation = lx.readTypeAnnotation();
      valueType = annotation.type;
      if (!annotation.hasDefault)
        throw new LexError(`Computed '${nameToken.value}' requires an expression`);
    } else {
      const eq = lx.next();
      if (eq.type !== "eq")
        throw new LexError(`Expected '=' after computed '${nameToken.value}' at offset ${eq.pos}`);
    }
    out.push({
      name: nameToken.value,
      ...(valueType ? { valueType } : {}),
      expr: lx.readPropInitializer(),
    });
  }
  return out;
}

function nestedBlock(source: string, name: string): string | undefined {
  const match = new RegExp(`\\b${name}\\s*\\{`).exec(source);
  if (!match) return undefined;
  const brace = source.indexOf("{", match.index);
  return readBalanced(source, brace, "{", "}").inner.trim() || undefined;
}

export function parsePersist(source: string): PersistDecl {
  const storage = /\bstorage\s*=\s*["'](memory|session|local)["']/.exec(source)?.[1] as
    PersistDecl["storage"] | undefined;
  const includeRaw = /\binclude\s*=\s*\[([\s\S]*?)\]/.exec(source)?.[1] ?? "";
  const include = Array.from(includeRaw.matchAll(/["']([^"']+)["']/g), (match) => match[1]!);
  const version = Number(/\bversion\s*=\s*(\d+)/.exec(source)?.[1] ?? "1");
  const migrations = nestedBlock(source, "migrations");
  const validation = nestedBlock(source, "validate");
  return {
    storage: storage ?? "memory",
    include,
    version,
    ...(migrations ? { migrations } : {}),
    ...(validation ? { validation } : {}),
  };
}

export function parseStoreLifecycle(source: string): StoreLifecycleDecl {
  const out: StoreLifecycleDecl = {};
  for (const hook of ["serverInit", "clientInit", "hydrate", "dispose"] as const) {
    const start = new RegExp(`\\b${hook}\\s*\\{`).exec(source);
    if (!start) continue;
    const brace = source.indexOf("{", start.index);
    out[hook] = readBalanced(source, brace, "{", "}").inner;
  }
  return out;
}