import { KdlStructuredError } from "../../kdl/rpc.js";
import type { MotionSourceMap } from "../../kdl/types.js";
import type {
  GrlBinaryExpression,
  GrlCallExpression,
  GrlExpression,
  GrlNumberLiteral,
  GrlUnaryExpression
} from "../ast/index.js";
import { normalizeUnitLiteral, type UnitKind } from "../lexer/units.js";

export type ConstantDimension = UnitKind | "scalar" | "boolean" | "string";

export interface ConstantValue {
  kind: ConstantDimension;
  value: number | boolean | string;
  rawUnit?: string;
}

export interface ConstantEvaluationContext {
  symbols?: Map<string, ConstantValue>;
  defaultUnit?: string;
  expectedKind?: ConstantDimension;
}

const EPSILON = 1e-12;

export function evaluateConstantExpression(
  expression: GrlExpression,
  context: ConstantEvaluationContext = {}
): ConstantValue {
  const value = evaluate(expression, context);
  return applyExpectedContext(value, context, expression);
}

export function evaluateNumberExpression(
  expression: GrlExpression,
  context: ConstantEvaluationContext = {}
): number {
  const value = evaluateConstantExpression(expression, context);
  if (typeof value.value !== "number") {
    throw expressionError("GRL_EXPR_NON_CONSTANT", "Expected numeric constant expression", expression);
  }
  return value.value;
}

export function constantToLiteralValue(expression: GrlExpression): unknown {
  const value = evaluateConstantExpression(expression);
  return value.value;
}

function evaluate(expression: GrlExpression, context: ConstantEvaluationContext): ConstantValue {
  switch (expression.kind) {
    case "NumberLiteral":
      return numberLiteralValue(expression, context);
    case "StringLiteral":
      return { kind: "string", value: expression.value };
    case "BooleanLiteral":
      return { kind: "boolean", value: expression.value };
    case "IdentifierExpression":
      return identifierValue(expression.name, expression, context);
    case "UnaryExpression":
      return unaryValue(expression, context);
    case "BinaryExpression":
      return binaryValue(expression, context);
    case "CallExpression":
      return callValue(expression, context);
    case "UnitExpression": {
      const innerContext: ConstantEvaluationContext = {};
      if (context.symbols) {
        innerContext.symbols = context.symbols;
      }
      const inner = evaluate(expression.expression, innerContext);
      if (typeof inner.value !== "number") {
        throw expressionError("GRL_EXPR_UNIT_MISMATCH", "Unit suffix requires a numeric expression", expression);
      }
      const unit = normalizeUnitLiteral(expression.unit.raw);
      if (inner.kind !== "scalar" && inner.kind !== unit.kind) {
        throw expressionError("GRL_EXPR_UNIT_MISMATCH", `Cannot apply ${unit.literal} to ${inner.kind}`, expression);
      }
      return {
        kind: unit.kind,
        value: inner.value * unit.factor,
        rawUnit: unit.literal
      };
    }
    default:
      throw expressionError("GRL_EXPR_NON_CONSTANT", `${expression.kind} is not a constant expression`, expression);
  }
}

function numberLiteralValue(expression: GrlNumberLiteral, context: ConstantEvaluationContext): ConstantValue {
  if (expression.unit) {
    return {
      kind: expression.unit.kind as UnitKind,
      value: expression.unit.normalizedValue,
      rawUnit: expression.unit.raw
    };
  }
  if (context.defaultUnit) {
    const unit = normalizeUnitLiteral(context.defaultUnit);
    return {
      kind: unit.kind,
      value: expression.value * unit.factor,
      rawUnit: unit.literal
    };
  }
  return { kind: "scalar", value: expression.value };
}

function identifierValue(
  name: string,
  expression: GrlExpression,
  context: ConstantEvaluationContext
): ConstantValue {
  if (name === "pi") {
    return { kind: "scalar", value: Math.PI };
  }
  if (name === "e") {
    return { kind: "scalar", value: Math.E };
  }
  const value = context.symbols?.get(name);
  if (value) {
    return value;
  }
  throw expressionError("GRL_EXPR_UNKNOWN_SYMBOL", `Unknown constant ${name}`, expression);
}

function unaryValue(expression: GrlUnaryExpression, context: ConstantEvaluationContext): ConstantValue {
  const value = evaluate(expression.argument, context);
  if (expression.operator === "not" || expression.operator === "!") {
    return { kind: "boolean", value: !coerceBoolean(value, expression) };
  }
  if (typeof value.value !== "number") {
    throw expressionError("GRL_EXPR_UNIT_MISMATCH", `Unary ${expression.operator} requires a number`, expression);
  }
  return {
    ...value,
    value: expression.operator === "-" ? -value.value : value.value
  };
}

function binaryValue(expression: GrlBinaryExpression, context: ConstantEvaluationContext): ConstantValue {
  if (expression.operator === "and" || expression.operator === "&&") {
    return {
      kind: "boolean",
      value: coerceBoolean(evaluate(expression.left, context), expression.left) &&
        coerceBoolean(evaluate(expression.right, context), expression.right)
    };
  }
  if (expression.operator === "or" || expression.operator === "||") {
    return {
      kind: "boolean",
      value: coerceBoolean(evaluate(expression.left, context), expression.left) ||
        coerceBoolean(evaluate(expression.right, context), expression.right)
    };
  }

  let left = evaluate(expression.left, contextWithoutDefaultUnit(context));
  let right = evaluate(expression.right, contextWithoutDefaultUnit(context));

  if (["==", "!=", "<", "<=", ">", ">="].includes(expression.operator)) {
    [left, right] = normalizeScalarUnits(left, right, context, expression);
    return compareValues(left, right, expression);
  }

  if (typeof left.value !== "number" || typeof right.value !== "number") {
    throw expressionError("GRL_EXPR_UNIT_MISMATCH", `${expression.operator} requires numeric operands`, expression);
  }

  if (expression.operator === "*" || expression.operator === "/") {
    return multiplyDivide(left, right, expression);
  }

  if (expression.operator === "mod") {
    [left, right] = normalizeScalarUnits(left, right, context, expression);
    ensureNumericCompatible(left, right, expression);
    const leftNumber = left.value as number;
    const rightNumber = right.value as number;
    if (Math.abs(rightNumber) < EPSILON) {
      throw expressionError("GRL_EXPR_DIV_ZERO", "Modulo by zero", expression);
    }
    return withRawUnit(left.kind, leftNumber % rightNumber, left.rawUnit);
  }

  [left, right] = normalizeScalarUnits(left, right, context, expression);
  ensureNumericCompatible(left, right, expression);
  const leftNumber = left.value as number;
  const rightNumber = right.value as number;
  return withRawUnit(left.kind, expression.operator === "+" ? leftNumber + rightNumber : leftNumber - rightNumber, left.rawUnit ?? right.rawUnit);
}

function contextWithoutDefaultUnit(context: ConstantEvaluationContext): ConstantEvaluationContext {
  const next: ConstantEvaluationContext = {};
  if (context.symbols) {
    next.symbols = context.symbols;
  }
  return next;
}

function contextWithoutExpectedKind(context: ConstantEvaluationContext): ConstantEvaluationContext {
  const next: ConstantEvaluationContext = {};
  if (context.symbols) {
    next.symbols = context.symbols;
  }
  if (context.defaultUnit) {
    next.defaultUnit = context.defaultUnit;
  }
  return next;
}

function normalizeScalarUnits(
  left: ConstantValue,
  right: ConstantValue,
  context: ConstantEvaluationContext,
  expression: GrlExpression
): [ConstantValue, ConstantValue] {
  if (typeof left.value !== "number" || typeof right.value !== "number") {
    return [left, right];
  }
  if (left.kind === "scalar" && right.kind !== "scalar") {
    return [coerceScalarToReference(left, right, context, expression), right];
  }
  if (right.kind === "scalar" && left.kind !== "scalar") {
    return [left, coerceScalarToReference(right, left, context, expression)];
  }
  return [left, right];
}

function coerceScalarToReference(
  scalarValue: ConstantValue,
  reference: ConstantValue,
  context: ConstantEvaluationContext,
  expression: GrlExpression
): ConstantValue {
  if (typeof scalarValue.value !== "number") {
    return scalarValue;
  }
  const unitLiteral = reference.rawUnit ?? context.defaultUnit;
  if (!unitLiteral) {
    throw expressionError("GRL_EXPR_UNIT_MISMATCH", `Cannot infer unit for scalar ${scalarValue.value}`, expression);
  }
  const unit = normalizeUnitLiteral(unitLiteral);
  if (unit.kind !== reference.kind) {
    throw expressionError("GRL_EXPR_UNIT_MISMATCH", `Cannot use ${unit.literal} with ${reference.kind}`, expression);
  }
  return {
    kind: unit.kind,
    value: scalarValue.value * unit.factor,
    rawUnit: unit.literal
  };
}

function multiplyDivide(
  left: ConstantValue,
  right: ConstantValue,
  expression: GrlBinaryExpression
): ConstantValue {
  if (expression.operator === "/" && Math.abs(right.value as number) < EPSILON) {
    throw expressionError("GRL_EXPR_DIV_ZERO", "Division by zero", expression);
  }
  if (left.kind === "scalar" && right.kind === "scalar") {
    return {
      kind: "scalar",
      value: expression.operator === "*" ? (left.value as number) * (right.value as number) : (left.value as number) / (right.value as number)
    };
  }
  if (left.kind !== "scalar" && right.kind === "scalar") {
    return withRawUnit(
      left.kind,
      expression.operator === "*" ? (left.value as number) * (right.value as number) : (left.value as number) / (right.value as number),
      left.rawUnit
    );
  }
  if (left.kind === "scalar" && right.kind !== "scalar" && expression.operator === "*") {
    return withRawUnit(right.kind, (left.value as number) * (right.value as number), right.rawUnit);
  }
  throw expressionError("GRL_EXPR_UNIT_MISMATCH", "Compound unit arithmetic is not supported", expression);
}

function compareValues(
  left: ConstantValue,
  right: ConstantValue,
  expression: GrlBinaryExpression
): ConstantValue {
  if (expression.operator === "==" || expression.operator === "!=") {
    if (left.kind !== right.kind) {
      return { kind: "boolean", value: expression.operator === "!=" };
    }
    return {
      kind: "boolean",
      value: expression.operator === "==" ? left.value === right.value : left.value !== right.value
    };
  }

  ensureNumericCompatible(left, right, expression);
  const leftNumber = left.value as number;
  const rightNumber = right.value as number;
  switch (expression.operator) {
    case "<":
      return { kind: "boolean", value: leftNumber < rightNumber };
    case "<=":
      return { kind: "boolean", value: leftNumber <= rightNumber };
    case ">":
      return { kind: "boolean", value: leftNumber > rightNumber };
    case ">=":
      return { kind: "boolean", value: leftNumber >= rightNumber };
    default:
      throw expressionError("GRL_EXPR_UNKNOWN_FUNCTION", `Unsupported comparison ${expression.operator}`, expression);
  }
}

function callValue(expression: GrlCallExpression, context: ConstantEvaluationContext): ConstantValue {
  const args = expression.args.map((arg) => evaluate(arg, context));
  switch (expression.callee) {
    case "sin":
      arity(expression, args, 1);
      return scalar(Math.sin(angleArg(args[0]!, expression)));
    case "cos":
      arity(expression, args, 1);
      return scalar(Math.cos(angleArg(args[0]!, expression)));
    case "tan":
      arity(expression, args, 1);
      return scalar(Math.tan(angleArg(args[0]!, expression)));
    case "asin":
      arity(expression, args, 1);
      return angle(inverseTrig(Math.asin, args[0]!, expression));
    case "acos":
      arity(expression, args, 1);
      return angle(inverseTrig(Math.acos, args[0]!, expression));
    case "atan":
      arity(expression, args, 1);
      return angle(Math.atan(numberArg(args[0]!, expression)));
    case "atan2":
      arity(expression, args, 2);
      return angle(Math.atan2(numberArg(args[0]!, expression), numberArg(args[1]!, expression)));
    case "sqrt":
      arity(expression, args, 1);
      return sameKindUnary(args[0]!, Math.sqrt(numberArg(args[0]!, expression, { min: 0 })));
    case "abs":
      arity(expression, args, 1);
      return sameKindUnary(args[0]!, Math.abs(numberArg(args[0]!, expression)));
    case "pow":
      arity(expression, args, 2);
      if (args[0]!.kind !== "scalar" && args[1]!.kind !== "scalar") {
        throw expressionError("GRL_EXPR_UNIT_MISMATCH", "pow exponent must be scalar", expression);
      }
      return sameKindUnary(args[0]!, Math.pow(numberArg(args[0]!, expression), numberArg(args[1]!, expression)));
    case "min":
      arityAtLeast(expression, args, 1);
      return minMax(args, expression, "min");
    case "max":
      arityAtLeast(expression, args, 1);
      return minMax(args, expression, "max");
    case "clamp":
      arity(expression, args, 3);
      ensureNumericCompatible(args[0]!, args[1]!, expression);
      ensureNumericCompatible(args[0]!, args[2]!, expression);
      return withRawUnit(
        args[0]!.kind,
        Math.min(Math.max(args[0]!.value as number, args[1]!.value as number), args[2]!.value as number),
        args[0]!.rawUnit
      );
    case "floor":
      arity(expression, args, 1);
      return sameKindUnary(args[0]!, Math.floor(numberArg(args[0]!, expression)));
    case "ceil":
      arity(expression, args, 1);
      return sameKindUnary(args[0]!, Math.ceil(numberArg(args[0]!, expression)));
    case "round":
      arity(expression, args, 1);
      return sameKindUnary(args[0]!, Math.round(numberArg(args[0]!, expression)));
    default:
      throw expressionError("GRL_EXPR_UNKNOWN_FUNCTION", `Unknown function ${expression.callee}`, expression);
  }
}

function minMax(args: ConstantValue[], expression: GrlExpression, mode: "min" | "max"): ConstantValue {
  const first = args[0]!;
  for (const arg of args.slice(1)) {
    ensureNumericCompatible(first, arg, expression);
  }
  const values = args.map((arg) => numberArg(arg, expression));
  return withRawUnit(first.kind, mode === "min" ? Math.min(...values) : Math.max(...values), first.rawUnit);
}

function applyExpectedContext(
  value: ConstantValue,
  context: ConstantEvaluationContext,
  expression: GrlExpression
): ConstantValue {
  if (!context.expectedKind || context.expectedKind === value.kind) {
    return value;
  }
  if (context.expectedKind !== "scalar" && value.kind === "scalar" && context.defaultUnit) {
    const unit = normalizeUnitLiteral(context.defaultUnit);
    if (unit.kind !== context.expectedKind) {
      throw expressionError("GRL_EXPR_UNIT_MISMATCH", `Default unit ${unit.literal} does not match ${context.expectedKind}`, expression);
    }
    return {
      kind: unit.kind,
      value: (value.value as number) * unit.factor,
      rawUnit: unit.literal
    };
  }
  throw expressionError("GRL_EXPR_UNIT_MISMATCH", `Expected ${context.expectedKind}, got ${value.kind}`, expression);
}

function ensureNumericCompatible(left: ConstantValue, right: ConstantValue, expression: GrlExpression): void {
  if (typeof left.value !== "number" || typeof right.value !== "number") {
    throw expressionError("GRL_EXPR_UNIT_MISMATCH", "Expected numeric operands", expression);
  }
  if (left.kind === right.kind) {
    return;
  }
  if (left.kind === "scalar" || right.kind === "scalar") {
    return;
  }
  throw expressionError("GRL_EXPR_UNIT_MISMATCH", `Cannot combine ${left.kind} and ${right.kind}`, expression);
}

function numberArg(
  value: ConstantValue,
  expression: GrlExpression,
  options: { min?: number; max?: number } = {}
): number {
  if (typeof value.value !== "number") {
    throw expressionError("GRL_EXPR_UNIT_MISMATCH", "Expected numeric argument", expression);
  }
  if (options.min !== undefined && value.value < options.min) {
    throw expressionError("GRL_EXPR_DOMAIN", "Function argument is below domain", expression);
  }
  if (options.max !== undefined && value.value > options.max) {
    throw expressionError("GRL_EXPR_DOMAIN", "Function argument is above domain", expression);
  }
  return value.value;
}

function angleArg(value: ConstantValue, expression: GrlExpression): number {
  if (typeof value.value !== "number") {
    throw expressionError("GRL_EXPR_UNIT_MISMATCH", "Expected angle argument", expression);
  }
  if (value.kind !== "angle" && value.kind !== "scalar") {
    throw expressionError("GRL_EXPR_UNIT_MISMATCH", `Expected angle, got ${value.kind}`, expression);
  }
  return value.value;
}

function inverseTrig(fn: (value: number) => number, value: ConstantValue, expression: GrlExpression): number {
  return fn(numberArg(value, expression, { min: -1, max: 1 }));
}

function coerceBoolean(value: ConstantValue, expression: GrlExpression): boolean {
  if (typeof value.value !== "boolean") {
    throw expressionError("GRL_EXPR_UNIT_MISMATCH", "Expected boolean expression", expression);
  }
  return value.value;
}

function arity(expression: GrlCallExpression, args: ConstantValue[], expected: number): void {
  if (args.length !== expected) {
    throw expressionError("GRL_EXPR_ARITY", `${expression.callee} expects ${expected} argument(s)`, expression);
  }
}

function arityAtLeast(expression: GrlCallExpression, args: ConstantValue[], minimum: number): void {
  if (args.length < minimum) {
    throw expressionError("GRL_EXPR_ARITY", `${expression.callee} expects at least ${minimum} argument(s)`, expression);
  }
}

function scalar(value: number): ConstantValue {
  return { kind: "scalar", value };
}

function angle(value: number): ConstantValue {
  return { kind: "angle", value, rawUnit: "rad" };
}

function sameKindUnary(base: ConstantValue, value: number): ConstantValue {
  return withRawUnit(base.kind, value, base.rawUnit);
}

function withRawUnit(kind: ConstantDimension, value: number, rawUnit?: string): ConstantValue {
  return rawUnit ? { kind, value, rawUnit } : { kind, value };
}

function expressionError(code: string, message: string, expression: GrlExpression): KdlStructuredError {
  const sourceMap: MotionSourceMap = {
    line: expression.range.start.line,
    column: expression.range.start.column
  };
  return new KdlStructuredError(code, message, [
    {
      severity: "error",
      code,
      message,
      sourceMap
    }
  ]);
}
