import type { MotionDiagnostic, MotionSourceMap } from "../../kdl/types.js";
import type {
  GrlDataDeclaration,
  GrlFunctionDeclaration,
  GrlOperationDeclaration,
  GrlPathDeclaration,
  GrlProcedureDeclaration,
  GrlProgram,
  GrlRawTopLevelDeclaration,
  GrlTargetDeclaration,
  GrlTopLevelDeclaration
} from "../ast/index.js";
import type {
  AlarmInstruction,
  BreakInstruction,
  CallInstruction,
  ContinueInstruction,
  ControlFlowInstruction,
  ExecutableBranch,
  ExecutableInstruction,
  ExecutableProcedure,
  ExecutableSwitchCase,
  IoFlowInstruction,
  ProcedureFlowInstruction,
  RawProcedureStatement,
  ReturnInstruction,
  SemanticProgramIr,
  SemanticSourceMapEntry,
  SemanticSymbol,
  UnsupportedRuntimeInstruction
} from "../ir/index.js";
import {
  analyzeExceptionSemantics,
  parseExceptionFlowStatements
} from "./compileException.js";
import { analyzeProcFunctionSemantics } from "./compileProcFunction.js";
import {
  buildMotionContext,
  compileMotionToKdlRequest,
  compileOperation,
  compilePathToPlanRequest,
  parseProcedureMotionInstructions,
  parseProcedureRunOperationStatements,
  parseProcedureRunPathStatements
} from "./compileMotion.js";
import {
  parseControlFlowStatements
} from "./compileControlFlow.js";
import { parseIoFlowStatements } from "./compileIo.js";

export interface SemanticCompileOptions {
  startJoints: number[];
  sampleTime: number;
}

const SEMANTIC_CHECKS = [
  "language/module/proc",
  "const/var/persistent symbols",
  "tool/frame/speed/zone",
  "joint_target/pose_target",
  "movej/movel/movec",
  "path/point/event/run_path",
  "operation/run_operation",
  "io/wait/pulse",
  "if/elseif/else/while/for/switch",
  "call/return/break/continue",
  "proc parameter directions",
  "func returns and side effects",
  "alarm/raise/try/catch",
  "source map propagation",
  "KDL motion request bridge",
  "KDL path request bridge",
  "duplicate symbol diagnostics",
  "missing reference diagnostics",
  "P1 unsupported diagnostics",
  "operation action expansion",
  "path event expansion",
  "raw statement preservation"
] as const;

export function compileSemanticProgram(program: GrlProgram, options: SemanticCompileOptions): SemanticProgramIr {
  const declarations = program.module.declarations;
  const diagnostics: MotionDiagnostic[] = [];
  const sourceMap: SemanticSourceMapEntry[] = [];
  const symbols = buildSemanticSymbols(declarations, diagnostics);
  const motionContext = buildMotionContext(
    declarations.filter(
      (decl): decl is GrlDataDeclaration | GrlTargetDeclaration =>
        decl.kind === "DataDeclaration" || decl.kind === "TargetDeclaration"
    )
  );
  const paths = declarations
    .filter((decl): decl is GrlPathDeclaration => decl.kind === "PathDeclaration")
    .map((path) => compilePathToPlanRequest(path, motionContext, options));
  const pathDeclarations = new Map(
    declarations
      .filter((decl): decl is GrlPathDeclaration => decl.kind === "PathDeclaration")
      .map((path) => [path.name, path])
  );
  const operations = declarations
    .filter((decl): decl is GrlOperationDeclaration => decl.kind === "OperationDeclaration")
    .map((operation) => compileOperation(operation, pathDeclarations));
  const procFunction = analyzeProcFunctionSemantics(declarations);
  const exceptionAnalysis = analyzeExceptionSemantics(declarations);
  diagnostics.push(...procFunction.diagnostics, ...exceptionAnalysis.diagnostics);
  const procedures = declarations
    .filter((decl): decl is GrlProcedureDeclaration => decl.kind === "ProcedureDeclaration")
    .map((procedure) => compileExecutableProcedure(procedure, motionContext, diagnostics));

  for (const path of paths) {
    collectPathSourceMaps(path, sourceMap);
  }
  for (const operation of operations) {
    collectOperationSourceMaps(operation, sourceMap);
  }
  for (const procedure of procedures) {
    collectExecutableSourceMaps(procedure.instructions, sourceMap, { procedureId: procedure.name });
  }

  return {
    moduleName: program.module.name,
    symbols,
    semanticChecks: [...SEMANTIC_CHECKS],
    procedures,
    paths,
    operations,
    diagnostics,
    sourceMap,
    kdlBridge: {
      motionRequests: procedures.flatMap((procedure) =>
        procedure.instructions.flatMap((instruction) =>
          instruction.kind === "MOVEJ" || instruction.kind === "MOVEL" || instruction.kind === "MOVEC"
            ? [compileMotionToKdlRequest(instruction, options)]
            : []
        )
      ),
      pathRequests: paths.map((path) => path.request)
    }
  };
}

function compileExecutableProcedure(
  procedure: GrlProcedureDeclaration,
  motionContext: ReturnType<typeof buildMotionContext>,
  diagnostics: MotionDiagnostic[]
): ExecutableProcedure {
  const motion = parseProcedureMotionInstructions(procedure, cloneMotionContextForSemantic(motionContext));
  const io = parseIoFlowStatements(procedure.bodyTokens);
  const runPaths = parseProcedureRunPathStatements(procedure);
  const runOperations = parseProcedureRunOperationStatements(procedure);
  const controls = parseControlFlowStatements(procedure.bodyTokens);
  const exceptions = parseExceptionFlowStatements(procedure.bodyTokens);
  const nestedControlLines = collectNestedControlLines(controls);
  const topLevelMotion = motion.filter((instruction) => !isNestedInstruction(instruction, nestedControlLines));
  const topLevelIo = io.filter((instruction) => !isNestedInstruction(instruction, nestedControlLines));
  const topLevelRunPaths = runPaths.filter((instruction) => !isNestedInstruction(instruction, nestedControlLines));
  const topLevelRunOperations = runOperations.filter((instruction) => !isNestedInstruction(instruction, nestedControlLines));
  const topLevelExceptions = exceptions.filter(
    (instruction): instruction is AlarmInstruction | UnsupportedRuntimeInstruction =>
      (instruction.kind === "ALARM" || instruction.kind === "UNSUPPORTED_RUNTIME") &&
      !isNestedInstruction(instruction, nestedControlLines)
  );
  const topLevelStructuredLines = new Set([
    ...topLevelMotion,
    ...topLevelIo,
    ...topLevelRunPaths,
    ...topLevelRunOperations,
    ...topLevelExceptions
  ].map((instruction) => instruction.sourceMap?.line).filter((line): line is number => line !== undefined));
  const instructions = mergeExecutableInstructions(
    procedure.bodyTokens,
    [
      ...topLevelMotion,
      ...topLevelIo,
      ...topLevelRunPaths,
      ...topLevelRunOperations,
      ...topLevelExceptions,
      ...flattenControlInstructions(controls, topLevelStructuredLines),
      ...extractRawCallsAndReturns(controls, topLevelStructuredLines)
    ],
    diagnostics
  );

  return {
    name: procedure.name,
    instructions,
    sourceMap: rangeSourceMap(procedure.range.start)
  };
}

function mergeExecutableInstructions(
  tokens: GrlProcedureDeclaration["bodyTokens"],
  instructions: ExecutableInstruction[],
  diagnostics: MotionDiagnostic[]
): ExecutableInstruction[] {
  const sorted = [...instructions].sort((left, right) => sourceOrder(left.sourceMap, right.sourceMap));
  const seen = new Set<string>();
  const merged: ExecutableInstruction[] = [];
  for (const instruction of sorted) {
    const key = instructionKey(instruction);
    if (key && seen.has(key)) {
      continue;
    }
    if (key) {
      seen.add(key);
    }
    merged.push(instruction);
  }
  for (const token of tokens) {
    if (["catch", "finally", "end"].includes(token.raw)) {
      continue;
    }
    if (!merged.some((instruction) => instruction.sourceMap?.line === token.range.start.line)) {
      diagnostics.push(diagnostic("info", "GRL_RAW_STATEMENT_PRESERVED", `Statement ${token.raw} preserved as raw IR`, tokenSourceMap(token)));
    }
  }
  return merged;
}

function flattenControlInstructions(instructions: ProcedureFlowInstruction[], excludedRawLines = new Set<number>()): ExecutableInstruction[] {
  return instructions.flatMap((instruction): ExecutableInstruction[] => {
    if (instruction.kind === "RAW_STATEMENT") {
      return rawStatementToExecutable(instruction, excludedRawLines);
    }
    if (instruction.kind === "IF") {
      return [
        {
          kind: "EXEC_IF",
          branches: instruction.branches.map((branch): ExecutableBranch => ({
            branchKind: branch.branchKind,
            ...(branch.condition ? { condition: branch.condition } : {}),
            body: flattenProcedureFlow(branch.body),
            ...(branch.sourceMap ? { sourceMap: branch.sourceMap } : {})
          })),
          ...(instruction.sourceMap ? { sourceMap: instruction.sourceMap } : {})
        }
      ];
    }
    if (instruction.kind === "WHILE") {
      return [
        {
          kind: "EXEC_WHILE",
          condition: instruction.condition,
          body: flattenProcedureFlow(instruction.body),
          ...(instruction.sourceMap ? { sourceMap: instruction.sourceMap } : {})
        }
      ];
    }
    if (instruction.kind === "FOR") {
      return [
        {
          kind: "EXEC_FOR",
          iterator: instruction.iterator,
          from: instruction.from,
          to: instruction.to,
          ...(instruction.step ? { step: instruction.step } : {}),
          body: flattenProcedureFlow(instruction.body),
          ...(instruction.sourceMap ? { sourceMap: instruction.sourceMap } : {})
        }
      ];
    }
    if (instruction.kind === "SWITCH") {
      return [
        {
          kind: "EXEC_SWITCH",
          expression: instruction.expression,
          cases: instruction.cases.map((switchCase): ExecutableSwitchCase => ({
            caseKind: switchCase.caseKind,
            ...(switchCase.value !== undefined ? { value: switchCase.value } : {}),
            ...(switchCase.raw ? { raw: switchCase.raw } : {}),
            body: flattenProcedureFlow(switchCase.body),
            ...(switchCase.sourceMap ? { sourceMap: switchCase.sourceMap } : {})
          })),
          ...(instruction.sourceMap ? { sourceMap: instruction.sourceMap } : {})
        }
      ];
    }
    if (instruction.kind === "BREAK" || instruction.kind === "CONTINUE") {
      return [instruction as BreakInstruction | ContinueInstruction];
    }
    return [];
  });
}

function flattenProcedureFlow(instructions: ProcedureFlowInstruction[]): ExecutableInstruction[] {
  const controls = instructions.filter((instruction): instruction is ControlFlowInstruction => instruction.kind !== "RAW_STATEMENT");
  const raw = instructions.filter((instruction): instruction is RawProcedureStatement => instruction.kind === "RAW_STATEMENT");
  return [...flattenControlInstructions(controls), ...extractRawCallsAndReturns(raw)];
}

function extractRawCallsAndReturns(
  instructions: Array<ControlFlowInstruction | RawProcedureStatement>,
  excludedRawLines = new Set<number>()
): ExecutableInstruction[] {
  const extracted: ExecutableInstruction[] = [];
  for (const instruction of instructions) {
    if (instruction.kind !== "RAW_STATEMENT") {
      continue;
    }
    extracted.push(...rawStatementToExecutable(instruction, excludedRawLines));
  }
  return extracted;
}

function rawStatementToExecutable(
  instruction: RawProcedureStatement,
  excludedRawLines = new Set<number>()
): ExecutableInstruction[] {
    if (instruction.sourceMap?.line && excludedRawLines.has(instruction.sourceMap.line)) {
      return [];
    }
    const tokens = instruction.tokens as { raw: string }[] | undefined;
    const first = tokens?.[0]?.raw;
    const second = tokens?.[1]?.raw;
    if (first && ["set_tool", "set_frame", "set_speed", "set_zone"].includes(first)) {
      return [];
    }
    const typedTokens = instruction.tokens as GrlProcedureDeclaration["bodyTokens"] | undefined;
    if (typedTokens && (first === "io" || first === "wait" || first === "pulse")) {
      return parseIoFlowStatements(typedTokens);
    }
    if (typedTokens && (first === "alarm" || first === "raise" || first === "enable" || first === "disable")) {
      return parseExceptionFlowStatements(typedTokens).filter(
        (item): item is AlarmInstruction | UnsupportedRuntimeInstruction => item.kind === "ALARM" || item.kind === "UNSUPPORTED_RUNTIME"
      );
    }
    if (first === "call" && second) {
      return [{
        kind: "CALL" as const,
        target: second,
        args: [],
        ...(instruction.sourceMap ? { sourceMap: instruction.sourceMap } : {})
      }];
    }
    if (first === "return") {
      return [{
        kind: "RETURN" as const,
        ...(instruction.sourceMap ? { sourceMap: instruction.sourceMap } : {})
      }];
    }
    return [instruction];
}

function collectNestedControlLines(instructions: ProcedureFlowInstruction[]): Set<number> {
  const lines = new Set<number>();
  for (const instruction of instructions) {
    if (instruction.kind === "RAW_STATEMENT") {
      continue;
    }
    if (instruction.kind === "IF") {
      for (const branch of instruction.branches) {
        collectFlowLines(branch.body, lines);
      }
    } else if (instruction.kind === "WHILE" || instruction.kind === "FOR") {
      collectFlowLines(instruction.body, lines);
    } else if (instruction.kind === "SWITCH") {
      for (const switchCase of instruction.cases) {
        collectFlowLines(switchCase.body, lines);
      }
    }
  }
  return lines;
}

function collectFlowLines(instructions: ProcedureFlowInstruction[], lines: Set<number>): void {
  for (const instruction of instructions) {
    if (instruction.sourceMap?.line) {
      lines.add(instruction.sourceMap.line);
    }
    if (instruction.kind === "IF") {
      for (const branch of instruction.branches) {
        collectFlowLines(branch.body, lines);
      }
    } else if (instruction.kind === "WHILE" || instruction.kind === "FOR") {
      collectFlowLines(instruction.body, lines);
    } else if (instruction.kind === "SWITCH") {
      for (const switchCase of instruction.cases) {
        collectFlowLines(switchCase.body, lines);
      }
    }
  }
}

function isNestedInstruction(instruction: { sourceMap?: MotionSourceMap }, nestedLines: Set<number>): boolean {
  return Boolean(instruction.sourceMap?.line && nestedLines.has(instruction.sourceMap.line));
}

function buildSemanticSymbols(declarations: GrlTopLevelDeclaration[], diagnostics: MotionDiagnostic[]): SemanticSymbol[] {
  const symbols: SemanticSymbol[] = [];
  const seen = new Map<string, SemanticSymbol>();
  for (const declaration of declarations) {
    const symbol = symbolFromDeclaration(declaration);
    if (!symbol) {
      continue;
    }
    const previous = seen.get(symbol.name);
    if (previous) {
      diagnostics.push(diagnostic("error", "GRL_SYMBOL_DUPLICATE", `Duplicate symbol ${symbol.name}`, symbol.sourceMap));
    }
    seen.set(symbol.name, symbol);
    symbols.push(symbol);
  }
  return symbols;
}

function symbolFromDeclaration(declaration: GrlTopLevelDeclaration): SemanticSymbol | undefined {
  if (declaration.kind === "DataDeclaration") {
    return { kind: "data", name: declaration.name, typeName: declaration.typeName, sourceMap: rangeSourceMap(declaration.range.start) };
  }
  if (declaration.kind === "TargetDeclaration") {
    return { kind: "target", name: declaration.name, sourceMap: rangeSourceMap(declaration.range.start) };
  }
  if (declaration.kind === "PathDeclaration") {
    return { kind: "path", name: declaration.name, sourceMap: rangeSourceMap(declaration.range.start) };
  }
  if (declaration.kind === "OperationDeclaration") {
    return { kind: "operation", name: declaration.name, sourceMap: rangeSourceMap(declaration.range.start) };
  }
  if (declaration.kind === "ProcedureDeclaration") {
    return { kind: "procedure", name: declaration.name, sourceMap: rangeSourceMap(declaration.range.start) };
  }
  if (declaration.kind === "FunctionDeclaration") {
    return { kind: "function", name: declaration.name, typeName: declaration.returnType, sourceMap: rangeSourceMap(declaration.range.start) };
  }
  if (declaration.kind === "RawTopLevelDeclaration") {
    const name = declaration.tokens[1]?.raw ?? declaration.declarationType;
    return { kind: "raw", name, typeName: declaration.declarationType, sourceMap: rangeSourceMap(declaration.range.start) };
  }
  return undefined;
}

function collectPathSourceMaps(path: SemanticProgramIr["paths"][number], sourceMap: SemanticSourceMapEntry[]): void {
  for (const segment of path.request.segments) {
    if (segment.sourceMap) {
      sourceMap.push({
        kind: "path_point",
        id: segment.id ?? `${path.pathId}:${segment.motion}`,
        pathId: path.pathId,
        ...(segment.id ? { pointId: segment.id } : {}),
        sourceMap: segment.sourceMap
      });
    }
  }
  for (const event of path.events) {
    if (event.sourceMap) {
      sourceMap.push({
        kind: "path_event",
        id: event.id ?? `${path.pathId}:event`,
        pathId: path.pathId,
        pointId: event.pointId,
        sourceMap: event.sourceMap
      });
    }
  }
}

function collectOperationSourceMaps(operation: SemanticProgramIr["operations"][number], sourceMap: SemanticSourceMapEntry[]): void {
  for (const action of [...operation.startActions, ...operation.endActions]) {
    if (action.sourceMap) {
      sourceMap.push({
        kind: "operation_action",
        id: `${operation.operationId}:${action.actionKind}`,
        operationId: operation.operationId,
        sourceMap: action.sourceMap
      });
    }
  }
}

function collectExecutableSourceMaps(
  instructions: ExecutableInstruction[],
  sourceMap: SemanticSourceMapEntry[],
  context: { procedureId: string }
): void {
  for (const instruction of instructions) {
    if (instruction.sourceMap) {
      sourceMap.push({
        kind: instruction.kind,
        id: `${context.procedureId}:${instruction.kind}:${instruction.sourceMap.line ?? 0}:${instruction.sourceMap.column ?? 0}`,
        procedureId: context.procedureId,
        sourceMap: instruction.sourceMap
      });
    }
    if (instruction.kind === "EXEC_IF") {
      for (const branch of instruction.branches) {
        collectExecutableSourceMaps(branch.body, sourceMap, context);
      }
    } else if (instruction.kind === "EXEC_WHILE" || instruction.kind === "EXEC_FOR") {
      collectExecutableSourceMaps(instruction.body, sourceMap, context);
    } else if (instruction.kind === "EXEC_SWITCH") {
      for (const switchCase of instruction.cases) {
        collectExecutableSourceMaps(switchCase.body, sourceMap, context);
      }
    }
  }
}

function cloneMotionContextForSemantic(context: ReturnType<typeof buildMotionContext>): ReturnType<typeof buildMotionContext> {
  return {
    targets: context.targets,
    speeds: context.speeds,
    zones: context.zones,
    tools: context.tools,
    frames: context.frames,
    ...(context.currentSpeed ? { currentSpeed: context.currentSpeed } : {}),
    ...(context.currentZone ? { currentZone: context.currentZone } : {}),
    ...(context.currentTool ? { currentTool: context.currentTool } : {}),
    ...(context.currentFrame ? { currentFrame: context.currentFrame } : {})
  };
}

function sourceOrder(left: MotionSourceMap | undefined, right: MotionSourceMap | undefined): number {
  return (left?.line ?? Number.MAX_SAFE_INTEGER) - (right?.line ?? Number.MAX_SAFE_INTEGER) ||
    (left?.column ?? Number.MAX_SAFE_INTEGER) - (right?.column ?? Number.MAX_SAFE_INTEGER);
}

function instructionKey(instruction: ExecutableInstruction): string | undefined {
  const line = instruction.sourceMap?.line;
  const column = instruction.sourceMap?.column;
  return line ? `${instruction.kind}:${line}:${column ?? 0}` : undefined;
}

function tokenSourceMap(token: GrlProcedureDeclaration["bodyTokens"][number]): 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 } : {})
  };
}
