import { KdlStructuredError } from "../../kdl/rpc.js";
import type { MotionSourceMap } from "../../kdl/types.js";
import type { GrlProcedureDeclaration } from "../ast/index.js";
import type {
  ControlExpression,
  ControlFlowInstruction,
  ForInstruction,
  IfInstruction,
  JumpInstruction,
  LabelInstruction,
  ProcedureFlowInstruction,
  RawProcedureStatement,
  SwitchCaseInstruction,
  SwitchInstruction,
  WhileInstruction
} from "../ir/index.js";
import type { GrlToken } from "../lexer/index.js";

type StopKeyword = "elseif" | "else" | "case" | "default" | "end";

interface ParseContext {
  scopePath: string[];
  loopDepth: number;
  switchDepth: number;
}

export function parseProcedureControlFlow(procedure: GrlProcedureDeclaration): ProcedureFlowInstruction[] {
  return parseControlFlowStatements(procedure.bodyTokens);
}

export function parseControlFlowStatements(tokens: GrlToken[]): ProcedureFlowInstruction[] {
  const parser = new ControlFlowParser(tokens);
  const flow = parser.parseRoot();
  validateJumps(flow);
  return flow;
}

class ControlFlowParser {
  private current = 0;

  constructor(private readonly tokens: GrlToken[]) {}

  parseRoot(): ProcedureFlowInstruction[] {
    return this.parseBlock(new Set(), {
      scopePath: [],
      loopDepth: 0,
      switchDepth: 0
    });
  }

  private parseBlock(stopKeywords: Set<StopKeyword>, context: ParseContext): ProcedureFlowInstruction[] {
    const instructions: ProcedureFlowInstruction[] = [];
    while (!this.isAtEnd()) {
      if (this.isStopKeyword(stopKeywords)) {
        break;
      }

      if (this.matchKeyword("if")) {
        instructions.push(this.finishIf(this.previous(), context));
        continue;
      }
      if (this.matchKeyword("while")) {
        instructions.push(this.finishWhile(this.previous(), context));
        continue;
      }
      if (this.matchKeyword("for")) {
        instructions.push(this.finishFor(this.previous(), context));
        continue;
      }
      if (this.matchKeyword("switch")) {
        instructions.push(this.finishSwitch(this.previous(), context));
        continue;
      }
      if (this.matchKeyword("break")) {
        instructions.push(this.finishBreak(this.previous(), context));
        continue;
      }
      if (this.matchKeyword("continue")) {
        instructions.push(this.finishContinue(this.previous(), context));
        continue;
      }
      if (this.matchKeyword("label")) {
        instructions.push(this.finishLabel(this.previous(), context));
        continue;
      }
      if (this.matchKeyword("jump")) {
        instructions.push(this.finishJump(this.previous(), context));
        continue;
      }
      if (this.checkKeyword("end")) {
        throw controlError("GRL_CONTROL_UNEXPECTED_END", "Unexpected end in procedure body", this.peek());
      }

      instructions.push(this.finishRawStatement());
    }
    return instructions;
  }

  private finishIf(start: GrlToken, context: ParseContext): IfInstruction {
    const condition = this.parseBooleanLineExpression(start);
    const branches: IfInstruction["branches"] = [
      {
        branchKind: "if",
        condition,
        body: this.parseBlock(new Set(["elseif", "else", "end"]), {
          ...context,
          scopePath: [...context.scopePath, scopeId(start, "if")]
        }),
        sourceMap: tokenSourceMap(start)
      }
    ];

    while (this.matchKeyword("elseif")) {
      const branchStart = this.previous();
      const branchCondition = this.parseBooleanLineExpression(branchStart);
      branches.push({
        branchKind: "elseif",
        condition: branchCondition,
        body: this.parseBlock(new Set(["elseif", "else", "end"]), {
          ...context,
          scopePath: [...context.scopePath, scopeId(branchStart, `elseif${branches.length}`)]
        }),
        sourceMap: tokenSourceMap(branchStart)
      });
    }

    if (this.matchKeyword("else")) {
      const branchStart = this.previous();
      branches.push({
        branchKind: "else",
        body: this.parseBlock(new Set(["end"]), {
          ...context,
          scopePath: [...context.scopePath, scopeId(branchStart, "else")]
        }),
        sourceMap: tokenSourceMap(branchStart)
      });
    }

    this.consumeKeyword("end", "Expected end after if block");
    return {
      kind: "IF",
      branches,
      sourceMap: tokenSourceMap(start)
    };
  }

  private finishWhile(start: GrlToken, context: ParseContext): WhileInstruction {
    const condition = this.parseBooleanLineExpression(start);
    const body = this.parseBlock(new Set(["end"]), {
      scopePath: [...context.scopePath, scopeId(start, "while")],
      loopDepth: context.loopDepth + 1,
      switchDepth: context.switchDepth
    });
    this.consumeKeyword("end", "Expected end after while block");
    return {
      kind: "WHILE",
      condition,
      body,
      sourceMap: tokenSourceMap(start)
    };
  }

  private finishFor(start: GrlToken, context: ParseContext): ForInstruction {
    const iterator = this.consumeIdentifier("Expected loop variable after for");
    if (!this.matchOperator("=") && !this.matchOperator(":=")) {
      throw controlError("GRL_FOR_ASSIGNMENT_EXPECTED", "Expected = after for loop variable", this.peek());
    }

    const from = this.parseLineExpressionUntil(start, ["to"]);
    this.consumeKeyword("to", "Expected to in for loop");
    const to = this.parseLineExpressionUntil(start, ["step"]);
    const step = this.matchKeyword("step") ? this.parseLineExpressionUntil(start, []) : undefined;
    const body = this.parseBlock(new Set(["end"]), {
      scopePath: [...context.scopePath, scopeId(start, "for")],
      loopDepth: context.loopDepth + 1,
      switchDepth: context.switchDepth
    });
    this.consumeKeyword("end", "Expected end after for block");

    return {
      kind: "FOR",
      iterator: iterator.raw,
      from,
      to,
      ...(step ? { step } : {}),
      body,
      sourceMap: tokenSourceMap(start)
    };
  }

  private finishSwitch(start: GrlToken, context: ParseContext): SwitchInstruction {
    const expression = this.parseLineExpressionUntil(start, []);
    const cases: SwitchCaseInstruction[] = [];
    const seenCases = new Set<string>();
    let seenDefault = false;

    while (!this.isAtEnd() && !this.checkKeyword("end")) {
      if (this.matchKeyword("case")) {
        const caseStart = this.previous();
        const valueTokens = this.collectLineExpressionTokens(caseStart, []);
        if (valueTokens.length === 0) {
          throw controlError("GRL_SWITCH_CASE_VALUE_MISSING", "Expected case value", caseStart);
        }
        const constant = parseCaseConstant(valueTokens);
        const key = `${typeof constant.value}:${String(constant.value)}`;
        if (seenCases.has(key)) {
          throw controlError("GRL_SWITCH_CASE_DUPLICATE", `Duplicate switch case ${constant.raw}`, caseStart);
        }
        seenCases.add(key);
        cases.push({
          caseKind: "case",
          value: constant.value,
          raw: constant.raw,
          body: this.parseBlock(new Set(["case", "default", "end"]), {
            scopePath: [...context.scopePath, scopeId(caseStart, `case:${constant.raw}`)],
            loopDepth: context.loopDepth,
            switchDepth: context.switchDepth + 1
          }),
          sourceMap: tokenSourceMap(caseStart)
        });
        continue;
      }

      if (this.matchKeyword("default")) {
        const defaultStart = this.previous();
        if (seenDefault) {
          throw controlError("GRL_SWITCH_DEFAULT_DUPLICATE", "Duplicate switch default case", defaultStart);
        }
        seenDefault = true;
        cases.push({
          caseKind: "default",
          body: this.parseBlock(new Set(["case", "default", "end"]), {
            scopePath: [...context.scopePath, scopeId(defaultStart, "default")],
            loopDepth: context.loopDepth,
            switchDepth: context.switchDepth + 1
          }),
          sourceMap: tokenSourceMap(defaultStart)
        });
        continue;
      }

      throw controlError("GRL_SWITCH_CASE_EXPECTED", "Expected case, default, or end in switch", this.peek());
    }

    this.consumeKeyword("end", "Expected end after switch block");
    return {
      kind: "SWITCH",
      expression,
      cases,
      sourceMap: tokenSourceMap(start)
    };
  }

  private finishBreak(start: GrlToken, context: ParseContext): ControlFlowInstruction {
    if (context.loopDepth === 0 && context.switchDepth === 0) {
      throw controlError("GRL_BREAK_OUTSIDE_FLOW", "break is only valid inside loop or switch", start);
    }
    return {
      kind: "BREAK",
      sourceMap: tokenSourceMap(start)
    };
  }

  private finishContinue(start: GrlToken, context: ParseContext): ControlFlowInstruction {
    if (context.loopDepth === 0) {
      throw controlError("GRL_CONTINUE_OUTSIDE_LOOP", "continue is only valid inside loop", start);
    }
    return {
      kind: "CONTINUE",
      sourceMap: tokenSourceMap(start)
    };
  }

  private finishLabel(start: GrlToken, context: ParseContext): LabelInstruction {
    const label = this.consumeIdentifier("Expected label name");
    return {
      kind: "LABEL",
      name: label.raw,
      scopePath: [...context.scopePath],
      sourceMap: tokenSourceMap(start)
    };
  }

  private finishJump(start: GrlToken, context: ParseContext): JumpInstruction {
    const label = this.consumeIdentifier("Expected label name after jump");
    return {
      kind: "JUMP",
      label: label.raw,
      scopePath: [...context.scopePath],
      sourceMap: tokenSourceMap(start)
    };
  }

  private finishRawStatement(): RawProcedureStatement {
    const start = this.peek();
    const tokens = this.collectLineExpressionTokens(start, []);
    if (tokens.length === 0) {
      const token = this.advance();
      return {
        kind: "RAW_STATEMENT",
        text: token.raw,
        tokens: [token],
        sourceMap: tokenSourceMap(token)
      };
    }
    return {
      kind: "RAW_STATEMENT",
      text: stringifyTokens(tokens),
      tokens,
      sourceMap: tokenSourceMap(start)
    };
  }

  private parseBooleanLineExpression(start: GrlToken): ControlExpression {
    const expression = this.parseLineExpressionUntil(start, []);
    if (!isBooleanCondition(expression.tokens as GrlToken[])) {
      throw controlError("GRL_CONTROL_CONDITION_NOT_BOOL", "Control condition must be boolean", start);
    }
    return expression;
  }

  private parseLineExpressionUntil(start: GrlToken, stopKeywords: string[]): ControlExpression {
    const tokens = this.collectLineExpressionTokens(start, stopKeywords);
    if (tokens.length === 0) {
      throw controlError("GRL_CONTROL_EXPRESSION_MISSING", "Expected control expression", start);
    }
    return {
      text: stringifyTokens(tokens),
      tokens,
      sourceMap: tokenSourceMap(tokens[0]!)
    };
  }

  private collectLineExpressionTokens(start: GrlToken, stopKeywords: string[]): GrlToken[] {
    const tokens: GrlToken[] = [];
    let parenDepth = 0;
    let bracketDepth = 0;
    let braceDepth = 0;

    while (!this.isAtEnd()) {
      const token = this.peek();
      if (token.range.start.line !== start.range.start.line) {
        break;
      }
      if (
        parenDepth === 0 &&
        bracketDepth === 0 &&
        braceDepth === 0 &&
        (token.kind === "keyword" || token.kind === "identifier") &&
        stopKeywords.includes(token.raw)
      ) {
        break;
      }

      const consumed = this.advance();
      tokens.push(consumed);
      if (consumed.kind === "punctuation") {
        if (consumed.raw === "(") parenDepth += 1;
        if (consumed.raw === ")") parenDepth = Math.max(0, parenDepth - 1);
        if (consumed.raw === "[") bracketDepth += 1;
        if (consumed.raw === "]") bracketDepth = Math.max(0, bracketDepth - 1);
        if (consumed.raw === "{") braceDepth += 1;
        if (consumed.raw === "}") braceDepth = Math.max(0, braceDepth - 1);
      }
    }
    return tokens;
  }

  private isStopKeyword(stopKeywords: Set<StopKeyword>): boolean {
    if (stopKeywords.size === 0) {
      return false;
    }
    const token = this.peek();
    return (token.kind === "keyword" || token.kind === "identifier") && stopKeywords.has(token.raw as StopKeyword);
  }

  private matchKeyword(keyword: string): boolean {
    if (this.checkKeyword(keyword)) {
      this.advance();
      return true;
    }
    return false;
  }

  private checkKeyword(keyword: string): boolean {
    const token = this.peek();
    return (token.kind === "keyword" || token.kind === "identifier") && token.raw === keyword;
  }

  private consumeKeyword(keyword: string, message: string): GrlToken {
    if (this.checkKeyword(keyword)) {
      return this.advance();
    }
    throw controlError("GRL_KEYWORD_EXPECTED", message, this.peek());
  }

  private consumeIdentifier(message: string): GrlToken {
    const token = this.peek();
    if (token.kind === "identifier" || token.kind === "keyword") {
      return this.advance();
    }
    throw controlError("GRL_IDENTIFIER_EXPECTED", message, token);
  }

  private matchOperator(operator: string): boolean {
    const token = this.peek();
    if (token.kind === "operator" && token.raw === operator) {
      this.advance();
      return true;
    }
    return false;
  }

  private advance(): GrlToken {
    this.current += 1;
    return this.previous();
  }

  private previous(): GrlToken {
    return this.tokens[this.current - 1]!;
  }

  private peek(): GrlToken {
    return this.tokens[this.current]!;
  }

  private isAtEnd(): boolean {
    return this.current >= this.tokens.length;
  }
}

function validateJumps(flow: ProcedureFlowInstruction[]): void {
  const labels = new Map<string, LabelInstruction>();
  const jumps: JumpInstruction[] = [];

  visitFlow(flow, (instruction) => {
    if (instruction.kind === "LABEL") {
      if (labels.has(instruction.name)) {
        throw controlError("GRL_LABEL_DUPLICATE", `Duplicate label ${instruction.name}`, undefined, instruction.sourceMap);
      }
      labels.set(instruction.name, instruction);
    } else if (instruction.kind === "JUMP") {
      jumps.push(instruction);
    }
  });

  for (const jump of jumps) {
    const label = labels.get(jump.label);
    if (!label) {
      throw controlError("GRL_LABEL_NOT_FOUND", `Unknown label ${jump.label}`, undefined, jump.sourceMap);
    }
    if (!isPrefix(label.scopePath, jump.scopePath)) {
      throw controlError(
        "GRL_JUMP_INTO_BLOCK",
        `jump ${jump.label} cannot enter a nested or sibling block`,
        undefined,
        jump.sourceMap
      );
    }
  }
}

function visitFlow(flow: ProcedureFlowInstruction[], visit: (instruction: ProcedureFlowInstruction) => void): void {
  for (const instruction of flow) {
    visit(instruction);
    if (instruction.kind === "IF") {
      for (const branch of instruction.branches) {
        visitFlow(branch.body, visit);
      }
    } else if (instruction.kind === "WHILE" || instruction.kind === "FOR") {
      visitFlow(instruction.body, visit);
    } else if (instruction.kind === "SWITCH") {
      for (const switchCase of instruction.cases) {
        visitFlow(switchCase.body, visit);
      }
    }
  }
}

function isPrefix(prefix: string[], value: string[]): boolean {
  return prefix.length <= value.length && prefix.every((part, index) => value[index] === part);
}

function isBooleanCondition(tokens: GrlToken[]): boolean {
  if (tokens.length === 0) {
    return false;
  }
  if (tokens.length === 1) {
    const [token] = tokens;
    if (!token) return false;
    if (token.kind === "number" || token.kind === "string") {
      return false;
    }
    return token.kind === "identifier" || token.kind === "keyword";
  }
  if (tokens.length === 2 && tokens[0]?.raw === "-" && tokens[1]?.kind === "number") {
    return false;
  }
  if (tokens.some((token) => token.kind === "operator" && ["==", "!=", "<", ">", "<=", ">=", "&&", "||", "!"].includes(token.raw))) {
    return true;
  }
  if (tokens.some((token) => token.kind === "keyword" && (token.raw === "true" || token.raw === "false"))) {
    return true;
  }
  const first = tokens[0];
  return Boolean(
    first &&
      (first.kind === "keyword" || first.kind === "identifier") &&
      ["all", "any", "rising", "falling", "changed"].includes(first.raw)
  );
}

function parseCaseConstant(tokens: GrlToken[]): { value: string | number | boolean; raw: string } {
  if (tokens.length === 2 && tokens[0]?.kind === "operator" && tokens[0].raw === "-" && tokens[1]?.kind === "number") {
    const value = -(tokens[1].unit?.normalizedValue ?? tokens[1].value);
    return { value, raw: stringifyTokens(tokens) };
  }

  if (tokens.length !== 1) {
    throw controlError("GRL_SWITCH_CASE_NOT_CONSTANT", "switch case must be a constant expression", tokens[0]);
  }

  const token = tokens[0]!;
  if (token.kind === "number") {
    return { value: token.unit?.normalizedValue ?? token.value, raw: token.raw };
  }
  if (token.kind === "string") {
    return { value: token.value, raw: token.raw };
  }
  if (token.kind === "keyword" && (token.raw === "true" || token.raw === "false")) {
    return { value: token.raw === "true", raw: token.raw };
  }
  if (token.kind === "identifier" || token.kind === "keyword") {
    return { value: token.raw, raw: token.raw };
  }

  throw controlError("GRL_SWITCH_CASE_NOT_CONSTANT", "switch case must be a constant expression", token);
}

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

function scopeId(token: GrlToken, kind: string): string {
  return `${kind}@${token.range.start.line}:${token.range.start.column}`;
}

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

function controlError(
  code: string,
  message: string,
  token?: GrlToken,
  sourceMap?: MotionSourceMap
): KdlStructuredError {
  const map = sourceMap ?? (token ? tokenSourceMap(token) : undefined);
  return new KdlStructuredError(
    code,
    message,
    map
      ? [
          {
            severity: "error",
            code,
            message,
            sourceMap: map
          }
        ]
      : undefined
  );
}
