import { resolve } from "node:path";
import ts from "typescript";

export interface TypeError {
  readonly line: number;
  readonly column: number;
  readonly message: string;
}

const compilerOptions: ts.CompilerOptions = {
  target: ts.ScriptTarget.ES2022,
  module: ts.ModuleKind.ESNext,
  moduleResolution: ts.ModuleResolutionKind.NodeNext,
  strict: false,
  noImplicitAny: false,
  strictNullChecks: false,
  strictFunctionTypes: false,
  strictBindCallApply: false,
  alwaysStrict: false,
  strictPropertyInitialization: false,
  noImplicitThis: false,
  useUnknownInCatchVariables: false,
  noEmit: false,
  skipLibCheck: true,
  lib: ["lib.es2022.d.ts"],
};

// Type-correctness codes are filtered out (same set and design as
// pi-fabric's type-checker): the cell runs in a dynamic kernel environment,
// so strict structural complaints (unknown properties, mismatched
// assignments, implicit-any parameters) are noise, not cell errors.
const TYPE_CORRECTNESS_CODES = new Set<number>([
  2339, 2551,
  2322, 2345, 2367,
  2531, 2532, 18047, 18048,
  7006, 7008, 7019, 7031, 7032, 7033, 7034,
]);

const CELL_WRAPPER_PREFIX = "async function __cell() {\n";
const CELL_WRAPPER_SUFFIX = "\n}\n";

const normalizePath = (fileName: string): string =>
  fileName.replaceAll("\\", "/");

/**
 * Type-checks one guest cell. The cell is wrapped in a single async function
 * (same design as pi-fabric's type-checker), so top-level `await` and
 * `return` are valid and each check is a self-contained function body.
 * The virtual guest file lives under the process cwd so module resolution
 * (node_modules, @types) matches the workspace the kernel runs in.
 */
class TypeScriptCellChecker {
  readonly #guestFile: string;
  readonly #baseHost = ts.createCompilerHost(compilerOptions, true);
  readonly #stableFiles = new Map<string, ts.SourceFile>();
  readonly #host: ts.CompilerHost;
  #sourceText = "";
  #sourceFile: ts.SourceFile;
  #program: ts.Program | undefined;
  // Top-level names declared by cells that passed the check. The JS worker
  // persists those bindings on globalThis across cells, so a later cell
  // referencing them is valid at runtime; the checker must not reject it.
  // Only leniency can leak across sessions (names from another session are
  // still valid runtime references when that session's worker is alive).
  readonly #persistedNames = new Set<string>();

  constructor(cwd: string) {
    this.#guestFile = normalizePath(resolve(cwd, "__pi_codemode_ts_cell__.ts"));
    this.#sourceFile = ts.createSourceFile(
      this.#guestFile,
      "",
      ts.ScriptTarget.ES2022,
      true,
    );
    const isGuestFile = (fileName: string): boolean =>
      this.#baseHost.getCanonicalFileName(normalizePath(fileName)) ===
      this.#baseHost.getCanonicalFileName(this.#guestFile);
    this.#host = {
      ...this.#baseHost,
      fileExists: (fileName) =>
        isGuestFile(fileName) || this.#baseHost.fileExists(fileName),
      readFile: (fileName) =>
        isGuestFile(fileName) ? this.#sourceText : this.#baseHost.readFile(fileName),
      getSourceFile: (fileName, languageVersion, onError, shouldCreateNewSourceFile) => {
        if (isGuestFile(fileName)) {
          return this.#sourceFile;
        }
        const cached = this.#stableFiles.get(fileName);
        if (cached) {
          return cached;
        }
        const source = this.#baseHost.getSourceFile(
          fileName,
          languageVersion,
          onError,
          shouldCreateNewSourceFile,
        );
        if (source) {
          this.#stableFiles.set(fileName, source);
        }
        return source;
      },
    };
  }

  check(code: string): TypeError[] {
    this.#sourceText = `${CELL_WRAPPER_PREFIX}${code}${CELL_WRAPPER_SUFFIX}`;
    this.#sourceFile = ts.createSourceFile(
      this.#guestFile,
      this.#sourceText,
      ts.ScriptTarget.ES2022,
      true,
    );
    const program = ts.createProgram({
      rootNames: [this.#guestFile],
      options: compilerOptions,
      host: this.#host,
      ...(this.#program ? { oldProgram: this.#program } : {}),
    });
    this.#program = program;
    const diagnostics = [
      ...program.getSyntacticDiagnostics(this.#sourceFile),
      ...program
        .getSemanticDiagnostics(this.#sourceFile)
        .filter((diagnostic) => !TYPE_CORRECTNESS_CODES.has(diagnostic.code)),
    ];
    const errors = diagnostics.filter(
      (diagnostic) => !this.#isPersistedNameDiagnostic(diagnostic),
    );
    if (errors.length === 0) {
      this.#rememberDeclaredNames();
    }
    return errors.map((diagnostic) => {
      const message = ts.flattenDiagnosticMessageText(diagnostic.messageText, "\n");
      if (!diagnostic.file || diagnostic.start === undefined) {
        return { line: 0, column: 0, message };
      }
      const position = diagnostic.file.getLineAndCharacterOfPosition(diagnostic.start);
      return {
        line: Math.max(1, position.line),
        column: position.character + 1,
        message,
      };
    });
  }

  /** 2304/2552 diagnostics for a name a previous passing cell declared. */
  #isPersistedNameDiagnostic(diagnostic: ts.Diagnostic): boolean {
    if (diagnostic.code !== 2304 && diagnostic.code !== 2552) {
      return false;
    }
    const text = ts.flattenDiagnosticMessageText(diagnostic.messageText, " ");
    const match = /Cannot find name '([^']+)'/.exec(text);
    return match !== null && this.#persistedNames.has(match[1]);
  }

  /** Records the cell's top-level declarations so later cells can reuse them. */
  #rememberDeclaredNames(): void {
    const source = this.#sourceFile;
    const wrapper = source.statements.find(
      (statement): statement is ts.FunctionDeclaration =>
        ts.isFunctionDeclaration(statement) && statement.name?.text === "__cell",
    );
    const body = wrapper?.body;
    if (!body) {
      return;
    }
    for (const statement of body.statements) {
      if (ts.isVariableStatement(statement)) {
        for (const declaration of statement.declarationList.declarations) {
          collectBindingNames(declaration.name, this.#persistedNames);
        }
      } else if (ts.isFunctionDeclaration(statement)) {
        if (statement.name) {
          this.#persistedNames.add(statement.name.text);
        }
      } else if (ts.isClassDeclaration(statement)) {
        if (statement.name) {
          this.#persistedNames.add(statement.name.text);
        }
      } else if (ts.isEnumDeclaration(statement)) {
        this.#persistedNames.add(statement.name.text);
      } else if (ts.isModuleDeclaration(statement)) {
        if (ts.isIdentifier(statement.name)) {
          this.#persistedNames.add(statement.name.text);
        }
      }
    }
  }

  /** Forgets persisted names; called when the kernel is reset. */
  forgetPersistedNames(): void {
    this.#persistedNames.clear();
  }
}

function collectBindingNames(
  name: ts.BindingName,
  out: Set<string>,
): void {
  if (ts.isIdentifier(name)) {
    out.add(name.text);
    return;
  }
  for (const element of name.elements) {
    if (ts.isBindingElement(element)) {
      collectBindingNames(element.name, out);
    }
  }
}

let checker: TypeScriptCellChecker | undefined;

/**
 * Type-checks a `ts` cell body. Returns one entry per error with the
 * 1-based line and column of the offending position in the cell source.
 * The cell must not run when the returned list is non-empty.
 */
export const typeCheckCell = (code: string): TypeError[] => {
  checker ??= new TypeScriptCellChecker(process.cwd());
  return checker.check(code);
};

/** Forgets the names persisted by earlier cells (kernel reset). */
export const resetTypeScriptCellState = (): void => {
  checker?.forgetPersistedNames();
};

/** Formats checker errors into the message shown when a cell is rejected. */
export const typeErrorsMessage = (errors: readonly TypeError[]): string => {
  const lines = errors.map(
    (error) => `line ${error.line}, column ${error.column}: ${error.message}`,
  );
  return `TypeScript cell check failed:\n${lines.join("\n")}`;
};

/**
 * Transpiles a checked cell body to JavaScript. The body is transpiled as-is
 * (no wrapper): the JS worker wraps the emitted statements in its own async
 * IIFE, which provides last-expression display and globalThis persistence.
 */
export const transpileTypeScriptCell = (code: string): string =>
  ts.transpileModule(code, {
    compilerOptions: {
      target: ts.ScriptTarget.ES2022,
      module: ts.ModuleKind.ESNext,
    },
  }).outputText;