import { KdlStructuredError } from "../../kdl/rpc.js";
import type { MotionDiagnostic, MotionSourceMap } from "../../kdl/types.js";
import type {
  GrlDataDeclaration,
  GrlFunctionDeclaration,
  GrlProcedureDeclaration,
  GrlTargetDeclaration,
  GrlTopLevelDeclaration
} from "../ast/index.js";
import type {
  CallInstruction,
  ControlExpression,
  FunctionSignature,
  ProcedureFlowInstruction,
  ProcedureSignature,
  ProcFunctionAnalysis,
  ReturnInstruction,
  RoutineParameter,
  RoutineParameterDirection
} from "../ir/index.js";
import type { GrlToken } from "../lexer/index.js";
import { parseControlFlowStatements } from "./compileControlFlow.js";

type RoutineDeclaration = GrlProcedureDeclaration | GrlFunctionDeclaration;
type RoutineSignature = ProcedureSignature | FunctionSignature;
type RoutineKind = "proc" | "func";
type InferredType = string | "unknown";

interface AnalysisContext {
  routineName: string;
  routineKind: RoutineKind;
  returnType?: string;
  parameters: RoutineParameter[];
  symbols: Map<string, string>;
  outerNames: Set<string>;
  signatures: Map<string, RoutineSignature>;
  diagnostics: MotionDiagnostic[];
  calls: CallInstruction[];
  returns: ReturnInstruction[];
}

interface FlowResult {
  normalExits: Set<string>[];
  returnExits: Set<string>[];
}

export function analyzeProcFunctionSemantics(declarations: GrlTopLevelDeclaration[]): ProcFunctionAnalysis {
  const procedures = declarations.filter(
    (decl): decl is GrlProcedureDeclaration => decl.kind === "ProcedureDeclaration"
  );
  const functions = declarations.filter(
    (decl): decl is GrlFunctionDeclaration => decl.kind === "FunctionDeclaration"
  );
  const globalNames = collectGlobalNames(declarations);
  const diagnostics: MotionDiagnostic[] = [];
  const signatures = new Map<string, RoutineSignature>();
  const procedureSignatures = procedures.map((procedure) => compileProcedureSignature(procedure, globalNames, diagnostics));
  const functionSignatures = functions.map((func) => compileFunctionSignature(func, globalNames, diagnostics));

  for (const signature of [...procedureSignatures, ...functionSignatures]) {
    if (signatures.has(signature.name)) {
      throw routineError("GRL_ROUTINE_DUPLICATE", `Duplicate routine ${signature.name}`, signature.sourceMap);
    }
    signatures.set(signature.name, signature);
  }

  const calls: CallInstruction[] = [];
  const returns: ReturnInstruction[] = [];
  for (const procedure of procedures) {
    analyzeRoutineBody(procedure, "proc", undefined, signatures, globalNames, diagnostics, calls, returns);
  }
  for (const func of functions) {
    analyzeRoutineBody(func, "func", func.returnType, signatures, globalNames, diagnostics, calls, returns);
  }

  return {
    procedures: procedureSignatures,
    functions: functionSignatures,
    calls,
    returns,
    diagnostics
  };
}

export function compileProcedureSignature(
  procedure: GrlProcedureDeclaration,
  globalNames: Set<string> = new Set(),
  diagnostics: MotionDiagnostic[] = []
): ProcedureSignature {
  return {
    kind: "PROC_SIGNATURE",
    name: procedure.name,
    parameters: parseRoutineParameters(procedure.params, procedure.name, globalNames, diagnostics),
    sourceMap: rangeSourceMap(procedure.range.start)
  };
}

export function compileFunctionSignature(
  func: GrlFunctionDeclaration,
  globalNames: Set<string> = new Set(),
  diagnostics: MotionDiagnostic[] = []
): FunctionSignature {
  return {
    kind: "FUNC_SIGNATURE",
    name: func.name,
    returnType: func.returnType,
    parameters: parseRoutineParameters(func.params, func.name, globalNames, diagnostics),
    sourceMap: rangeSourceMap(func.range.start)
  };
}

function analyzeRoutineBody(
  declaration: RoutineDeclaration,
  routineKind: RoutineKind,
  returnType: string | undefined,
  signatures: Map<string, RoutineSignature>,
  globalNames: Set<string>,
  diagnostics: MotionDiagnostic[],
  calls: CallInstruction[],
  returns: ReturnInstruction[]
): void {
  const signature = signatures.get(declaration.name);
  if (!signature) {
    throw routineError("GRL_ROUTINE_NOT_FOUND", `Missing routine signature ${declaration.name}`);
  }

  const symbols = new Map(signature.parameters.map((parameter) => [parameter.name, parameter.typeName]));
  const context: AnalysisContext = {
    routineName: declaration.name,
    routineKind,
    ...(returnType ? { returnType } : {}),
    parameters: signature.parameters,
    symbols,
    outerNames: globalNames,
    signatures,
    diagnostics,
    calls,
    returns
  };
  const flow = parseControlFlowStatements(declaration.bodyTokens);
  const result = analyzeFlow(flow, new Set(), context);
  const outParameters = signature.parameters.filter((parameter) => parameter.direction === "out");

  for (const exit of [...result.normalExits, ...result.returnExits]) {
    for (const parameter of outParameters) {
      if (!exit.has(parameter.name)) {
        throw routineError(
          "GRL_OUT_PARAM_NOT_ASSIGNED",
          `out parameter ${parameter.name} is not assigned on all normal return paths`,
          parameter.sourceMap
        );
      }
    }
  }

  if (routineKind === "func" && returnType !== "void" && result.normalExits.length > 0) {
    throw routineError("GRL_FUNC_MISSING_RETURN", `Function ${declaration.name} does not return on all normal paths`, rangeSourceMap(declaration.range.start));
  }
}

function analyzeFlow(flow: ProcedureFlowInstruction[], incoming: Set<string>, context: AnalysisContext): FlowResult {
  let normalStates: Set<string>[] = [new Set(incoming)];
  const returnStates: Set<string>[] = [];

  for (const instruction of flow) {
    const nextNormalStates: Set<string>[] = [];
    for (const state of normalStates) {
      const result = analyzeInstruction(instruction, state, context);
      nextNormalStates.push(...result.normalExits);
      returnStates.push(...result.returnExits);
    }
    normalStates = nextNormalStates;
    if (normalStates.length === 0) {
      break;
    }
  }

  return {
    normalExits: normalStates,
    returnExits: returnStates
  };
}

function analyzeInstruction(
  instruction: ProcedureFlowInstruction,
  incoming: Set<string>,
  context: AnalysisContext
): FlowResult {
  if (instruction.kind === "RAW_STATEMENT") {
    return analyzeRawStatement(instruction.tokens as GrlToken[] | undefined, incoming, context);
  }

  if (instruction.kind === "IF") {
    const normalExits: Set<string>[] = [];
    const returnExits: Set<string>[] = [];
    for (const branch of instruction.branches) {
      const result = analyzeFlow(branch.body, new Set(incoming), context);
      normalExits.push(...result.normalExits);
      returnExits.push(...result.returnExits);
    }
    if (!instruction.branches.some((branch) => branch.branchKind === "else")) {
      normalExits.push(new Set(incoming));
    }
    return { normalExits, returnExits };
  }

  if (instruction.kind === "WHILE" || instruction.kind === "FOR") {
    const body = analyzeFlow(instruction.body, new Set(incoming), context);
    return {
      normalExits: [new Set(incoming), ...body.normalExits],
      returnExits: body.returnExits
    };
  }

  if (instruction.kind === "SWITCH") {
    const normalExits: Set<string>[] = [];
    const returnExits: Set<string>[] = [];
    for (const switchCase of instruction.cases) {
      const result = analyzeFlow(switchCase.body, new Set(incoming), context);
      normalExits.push(...result.normalExits);
      returnExits.push(...result.returnExits);
    }
    if (!instruction.cases.some((switchCase) => switchCase.caseKind === "default")) {
      normalExits.push(new Set(incoming));
    }
    return { normalExits, returnExits };
  }

  return {
    normalExits: [new Set(incoming)],
    returnExits: []
  };
}

function analyzeRawStatement(
  tokens: GrlToken[] | undefined,
  incoming: Set<string>,
  context: AnalysisContext
): FlowResult {
  if (!tokens || tokens.length === 0) {
    return { normalExits: [new Set(incoming)], returnExits: [] };
  }

  checkFunctionSideEffects(tokens, context);
  const assigned = new Set(incoming);
  const declaration = parseLocalDeclaration(tokens);
  if (declaration) {
    if (context.symbols.has(declaration.name) || context.outerNames.has(declaration.name)) {
      context.diagnostics.push(diagnostic("warning", "GRL_NAME_SHADOWS_OUTER_SCOPE", `Local ${declaration.name} shadows an outer name`, declaration.sourceMap));
    }
    context.symbols.set(declaration.name, declaration.typeName);
    assigned.add(declaration.name);
    return { normalExits: [assigned], returnExits: [] };
  }

  const assignment = parseAssignment(tokens);
  if (assignment) {
    assigned.add(assignment.name);
  }

  const call = parseCallStatement(tokens);
  if (call) {
    validateCall(call, context);
    applyCallAssignments(call, assigned, context);
    context.calls.push(call);
    if (call.target === context.routineName) {
      context.diagnostics.push(diagnostic("warning", "GRL_RECURSIVE_CALL", `Routine ${context.routineName} calls itself`, call.sourceMap));
    }
    return { normalExits: [assigned], returnExits: [] };
  }

  const returnInstruction = parseReturnStatement(tokens);
  if (returnInstruction) {
    validateReturn(returnInstruction, context);
    context.returns.push(returnInstruction);
    return { normalExits: [], returnExits: [assigned] };
  }

  return { normalExits: [assigned], returnExits: [] };
}

function parseRoutineParameters(
  tokens: GrlToken[],
  routineName: string,
  globalNames: Set<string>,
  diagnostics: MotionDiagnostic[]
): RoutineParameter[] {
  const parameters: RoutineParameter[] = [];
  const seen = new Set<string>();
  for (const group of splitTopLevel(tokens, ",")) {
    if (group.length === 0) {
      continue;
    }
    let offset = 0;
    let direction: RoutineParameterDirection = "in";
    const first = group[0]!;
    if (isDirection(first)) {
      direction = first.raw as RoutineParameterDirection;
      offset = 1;
    }
    const typeName = group[offset];
    const name = group[offset + 1];
    if (!typeName || !name || !isIdentifierLike(typeName) || !isIdentifierLike(name)) {
      throw routineError("GRL_PARAMETER_INVALID", `Invalid parameter list for ${routineName}`, tokenSourceMap(first));
    }
    if (seen.has(name.raw)) {
      throw routineError("GRL_PARAMETER_DUPLICATE", `Duplicate parameter ${name.raw}`, tokenSourceMap(name));
    }
    seen.add(name.raw);
    if (globalNames.has(name.raw)) {
      diagnostics.push(diagnostic("warning", "GRL_NAME_SHADOWS_OUTER_SCOPE", `Parameter ${name.raw} shadows an outer name`, tokenSourceMap(name)));
    }
    parameters.push({
      name: name.raw,
      typeName: typeName.raw,
      direction,
      sourceMap: tokenSourceMap(name)
    });
  }
  return parameters;
}

function parseLocalDeclaration(tokens: GrlToken[]): { name: string; typeName: string; sourceMap: MotionSourceMap } | undefined {
  const storage = tokens[0];
  if (!storage || !["var", "const", "persistent"].includes(storage.raw)) {
    return undefined;
  }
  const typeName = tokens[1];
  const name = tokens[2];
  if (!typeName || !name || !isIdentifierLike(typeName) || !isIdentifierLike(name)) {
    return undefined;
  }
  return {
    name: name.raw,
    typeName: typeName.raw,
    sourceMap: tokenSourceMap(name)
  };
}

function parseAssignment(tokens: GrlToken[]): { name: string; sourceMap: MotionSourceMap } | undefined {
  const name = tokens[0];
  const operator = tokens[1];
  if (!name || !operator || !isIdentifierLike(name) || operator.kind !== "operator" || (operator.raw !== "=" && operator.raw !== ":=")) {
    return undefined;
  }
  return {
    name: name.raw,
    sourceMap: tokenSourceMap(name)
  };
}

function parseCallStatement(tokens: GrlToken[]): CallInstruction | undefined {
  const start = tokens[0];
  const target = tokens[1];
  if (!start || start.raw !== "call" || !target || !isIdentifierLike(target)) {
    return undefined;
  }
  const argTokens = tokens.slice(2);
  const args = parseCallArgs(argTokens);
  return {
    kind: "CALL",
    target: target.raw,
    args,
    sourceMap: tokenSourceMap(start)
  };
}

function parseReturnStatement(tokens: GrlToken[]): ReturnInstruction | undefined {
  const start = tokens[0];
  if (!start || start.raw !== "return") {
    return undefined;
  }
  const valueTokens = tokens.slice(1);
  return {
    kind: "RETURN",
    ...(valueTokens.length > 0 ? { value: expressionFromTokens(valueTokens) } : {}),
    sourceMap: tokenSourceMap(start)
  };
}

function parseCallArgs(tokens: GrlToken[]): ControlExpression[] {
  if (tokens[0]?.raw === "(" && tokens.at(-1)?.raw === ")") {
    return splitTopLevel(tokens.slice(1, -1), ",").filter((group) => group.length > 0).map(expressionFromTokens);
  }
  return splitTopLevel(tokens, ",").filter((group) => group.length > 0).map(expressionFromTokens);
}

function validateCall(call: CallInstruction, context: AnalysisContext): void {
  const signature = context.signatures.get(call.target);
  if (!signature) {
    throw routineError("GRL_CALL_TARGET_NOT_FOUND", `Unknown call target ${call.target}`, call.sourceMap);
  }
  if (call.args.length !== signature.parameters.length) {
    throw routineError("GRL_CALL_ARITY_MISMATCH", `Call ${call.target} expects ${signature.parameters.length} arguments`, call.sourceMap);
  }
  if (context.routineKind === "func" && signature.kind === "PROC_SIGNATURE") {
    throw routineError("GRL_FUNC_SIDE_EFFECT", `Function ${context.routineName} cannot call procedure ${call.target}`, call.sourceMap);
  }

  for (let index = 0; index < signature.parameters.length; index += 1) {
    const parameter = signature.parameters[index]!;
    const arg = call.args[index]!;
    const argTokens = arg.tokens as GrlToken[] | undefined;
    if ((parameter.direction === "out" || parameter.direction === "inout") && (!argTokens || !isLValueExpression(argTokens))) {
      throw routineError("GRL_ARGUMENT_NOT_LVALUE", `${parameter.direction} argument ${parameter.name} must be a writable lvalue`, arg.sourceMap);
    }
    const actualType = inferExpressionType(argTokens ?? [], context);
    if (!isTypeCompatible(parameter.typeName, actualType)) {
      throw routineError("GRL_CALL_ARGUMENT_TYPE", `Argument ${index + 1} for ${call.target} is not compatible with ${parameter.typeName}`, arg.sourceMap);
    }
  }
}

function applyCallAssignments(call: CallInstruction, assigned: Set<string>, context: AnalysisContext): void {
  const signature = context.signatures.get(call.target);
  if (!signature) {
    return;
  }
  for (let index = 0; index < signature.parameters.length; index += 1) {
    const parameter = signature.parameters[index]!;
    if (parameter.direction !== "out" && parameter.direction !== "inout") {
      continue;
    }
    const argTokens = call.args[index]?.tokens as GrlToken[] | undefined;
    const target = argTokens?.[0];
    if (target && isIdentifierLike(target)) {
      assigned.add(target.raw);
    }
  }
}

function validateReturn(returnInstruction: ReturnInstruction, context: AnalysisContext): void {
  if (context.routineKind === "proc") {
    if (returnInstruction.value) {
      throw routineError("GRL_RETURN_VALUE_IN_PROC", "proc return cannot include a value", returnInstruction.sourceMap);
    }
    return;
  }

  if (context.returnType === "void") {
    if (returnInstruction.value) {
      throw routineError("GRL_RETURN_TYPE_MISMATCH", "void function cannot return a value", returnInstruction.sourceMap);
    }
    return;
  }

  if (!returnInstruction.value) {
    throw routineError("GRL_RETURN_VALUE_MISSING", `Function ${context.routineName} must return ${context.returnType}`, returnInstruction.sourceMap);
  }

  const actualType = inferExpressionType(returnInstruction.value.tokens as GrlToken[] | undefined ?? [], context);
  if (!isTypeCompatible(context.returnType ?? "unknown", actualType)) {
    throw routineError("GRL_RETURN_TYPE_MISMATCH", `Return value is not compatible with ${context.returnType}`, returnInstruction.value.sourceMap);
  }
}

function checkFunctionSideEffects(tokens: GrlToken[], context: AnalysisContext): void {
  if (context.routineKind !== "func") {
    return;
  }
  const first = tokens[0];
  if (!first) {
    return;
  }
  if (["movej", "movel", "movec", "wait", "pulse", "run_path", "run_operation"].includes(first.raw)) {
    throw routineError("GRL_FUNC_SIDE_EFFECT", `Function ${context.routineName} cannot execute ${first.raw}`, tokenSourceMap(first));
  }
}

function inferExpressionType(tokens: GrlToken[], context: AnalysisContext): InferredType {
  if (tokens.length === 0) {
    return "unknown";
  }
  if (tokens.some((token) => token.kind === "operator" && ["==", "!=", "<", ">", "<=", ">=", "&&", "||", "!"].includes(token.raw))) {
    return "bool";
  }
  if (tokens.length === 1) {
    return inferSingleTokenType(tokens[0]!, context);
  }
  if (tokens[0] && isIdentifierLike(tokens[0]) && tokens[1]?.raw === "(") {
    const signature = context.signatures.get(tokens[0].raw);
    if (signature?.kind === "FUNC_SIGNATURE") {
      return signature.returnType;
    }
  }
  const operandTypes = tokens
    .filter((token) => token.kind !== "operator" && token.kind !== "punctuation")
    .map((token) => inferSingleTokenType(token, context))
    .filter((type) => type !== "unknown");
  if (operandTypes.length > 0 && operandTypes.every((type) => ["int", "real"].includes(type))) {
    return operandTypes.includes("real") ? "real" : "int";
  }
  return "unknown";
}

function inferSingleTokenType(token: GrlToken, context: AnalysisContext): InferredType {
  if (token.kind === "number") {
    if (token.unit?.kind === "time") return "time";
    if (token.unit?.kind === "length") return "length";
    if (token.unit?.kind === "angle") return "angle";
    if (token.unit?.kind === "percent") return "percent";
    return Number.isInteger(token.value) ? "int" : "real";
  }
  if (token.kind === "string") {
    return "string";
  }
  if (token.kind === "keyword" && (token.raw === "true" || token.raw === "false")) {
    return "bool";
  }
  if (isIdentifierLike(token)) {
    return context.symbols.get(token.raw) ?? "unknown";
  }
  return "unknown";
}

function isTypeCompatible(expected: string, actual: InferredType): boolean {
  if (expected === "unknown" || actual === "unknown") {
    return true;
  }
  if (expected === actual) {
    return true;
  }
  return expected === "real" && actual === "int";
}

function isLValueExpression(tokens: GrlToken[]): boolean {
  const first = tokens[0];
  if (!first || !isIdentifierLike(first) || ["true", "false"].includes(first.raw)) {
    return false;
  }
  return !tokens.some((token) => token.kind === "operator");
}

function expressionFromTokens(tokens: GrlToken[]): ControlExpression {
  return {
    text: stringifyTokens(tokens),
    tokens,
    ...(tokens[0] ? { sourceMap: tokenSourceMap(tokens[0]) } : {})
  };
}

function splitTopLevel(tokens: GrlToken[], separator: string): GrlToken[][] {
  const groups: GrlToken[][] = [];
  let current: GrlToken[] = [];
  let parenDepth = 0;
  let bracketDepth = 0;
  let braceDepth = 0;

  for (const token of tokens) {
    if (
      token.kind === "punctuation" &&
      token.raw === separator &&
      parenDepth === 0 &&
      bracketDepth === 0 &&
      braceDepth === 0
    ) {
      groups.push(current);
      current = [];
      continue;
    }
    current.push(token);
    if (token.kind === "punctuation") {
      if (token.raw === "(") parenDepth += 1;
      if (token.raw === ")") parenDepth = Math.max(0, parenDepth - 1);
      if (token.raw === "[") bracketDepth += 1;
      if (token.raw === "]") bracketDepth = Math.max(0, bracketDepth - 1);
      if (token.raw === "{") braceDepth += 1;
      if (token.raw === "}") braceDepth = Math.max(0, braceDepth - 1);
    }
  }
  groups.push(current);
  return groups;
}

function collectGlobalNames(declarations: GrlTopLevelDeclaration[]): Set<string> {
  const names = new Set<string>();
  for (const declaration of declarations) {
    if (
      declaration.kind === "DataDeclaration" ||
      declaration.kind === "TargetDeclaration" ||
      declaration.kind === "PathDeclaration" ||
      declaration.kind === "OperationDeclaration" ||
      declaration.kind === "ProcedureDeclaration" ||
      declaration.kind === "FunctionDeclaration"
    ) {
      names.add(declaration.name);
    }
  }
  return names;
}

function isDirection(token: GrlToken | undefined): boolean {
  return Boolean(token && (token.raw === "in" || token.raw === "out" || token.raw === "inout"));
}

function isIdentifierLike(token: GrlToken): boolean {
  return token.kind === "identifier" || token.kind === "keyword";
}

function stringifyTokens(tokens: GrlToken[]): string {
  return tokens.map((token) => token.raw).join(" ");
}

function tokenSourceMap(token: GrlToken): MotionSourceMap {
  return {
    line: token.range.start.line,
    column: token.range.start.column
  };
}

function rangeSourceMap(position: { line: number; column: number }): MotionSourceMap {
  return {
    line: position.line,
    column: position.column
  };
}

function diagnostic(
  severity: MotionDiagnostic["severity"],
  code: string,
  message: string,
  sourceMap?: MotionSourceMap
): MotionDiagnostic {
  return {
    severity,
    code,
    message,
    ...(sourceMap ? { sourceMap } : {})
  };
}

function routineError(code: string, message: string, sourceMap?: MotionSourceMap): KdlStructuredError {
  return new KdlStructuredError(
    code,
    message,
    sourceMap
      ? [
          {
            severity: "error",
            code,
            message,
            sourceMap
          }
        ]
      : undefined
  );
}
