// Copyright (C) 2024 Rubén Beltrán del Río

// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.

// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU General Public License for more details.

// You should have received a copy of the GNU General Public License
// along with this program. If not, see https://map.tranquil.systems.
extension Wmap {
  class Lexer {
    let input: String
    private var current: String.Index
    private var line = 1
    private var column = 1
    private var expectingText = false

    init(_ input: String) {
      self.input = input
      self.current = input.startIndex
    }

    func nextToken() -> Token {
      skipWhitespace()

      guard current < input.endIndex else {
        return Token(
          type: .eof,
          value: "",
          position: currentPosition(),
          stringRange: current..<current
        )
      }

      let char = input[current].lowercased().first!
      let pos = currentPosition()
      let tokenStart = current

      switch char {
      case "\n", "\r":
        expectingText = false  // Reset expectingText on newline
        return scanLineEnding()

      case "(":
        advance()
        return Token(
          type: .leftParen,
          value: "(",
          position: pos,
          stringRange: tokenStart..<current
        )

      case ")":
        advance()
        return Token(
          type: .rightParen,
          value: ")",
          position: pos,
          stringRange: tokenStart..<current
        )

      case "[":
        return scanBracketedToken()

      case ",":
        advance()
        return Token(
          type: .comma,
          value: ",",
          position: pos,
          stringRange: tokenStart..<current
        )

      case "+":
        advance()
        return Token(
          type: .plus,
          value: "+",
          position: pos,
          stringRange: tokenStart..<current
        )

      case "-":
        if peek() == "-" {
          advance()
          advance()
          return Token(
            type: .edgeUndirected,
            value: "--",
            position: pos,
            stringRange: tokenStart..<current
          )
        } else if peek() == ">" {
          advance()
          advance()
          return Token(
            type: .edgeDirected,
            value: "->",
            position: pos,
            stringRange: tokenStart..<current
          )
        } else {
          advance()
          return Token(
            type: .minus,
            value: "-",
            position: pos,
            stringRange: tokenStart..<current
          )
        }

      case "0"..."9", ".":
        return scanNumber()

      case "i":
        return scanStageOrLabel()

      default:
        // If we're expecting text, scan as text instead of vertex label
        if expectingText {
          expectingText = false  // Reset after scanning text
          return scanText()
        }

        // If we encounter an unexpected character, try to scan it as a vertex label
        // This provides more graceful degradation than failing immediately
        return scanVertexLabel()
      }
    }

    func scanText() -> Token {
      let pos = currentPosition()

      // Skip any leading whitespace
      skipWhitespace()

      let tokenStart = current

      while current < input.endIndex && !isLineEnding(input[current]) {
        advance()
      }

      let originalValue = String(input[tokenStart..<current]).trimmingCharacters(in: .whitespaces)
      return Token(
        type: .text,
        value: originalValue,
        position: pos,
        stringRange: tokenStart..<current
      )
    }

    private func scanLineEnding() -> Token {
      let pos = currentPosition()
      let tokenStart = current

      if input[current] == "\r" {
        advance()
        // Check for CRLF
        if current < input.endIndex && input[current] == "\n" {
          advance()
          let token = Token(
            type: .newline,
            value: "\r\n",
            position: pos,
            stringRange: tokenStart..<current
          )
          line += 1
          column = 1
          return token
        } else {
          // Just CR
          let token = Token(
            type: .newline,
            value: "\r",
            position: pos,
            stringRange: tokenStart..<current
          )
          line += 1
          column = 1
          return token
        }
      } else if input[current] == "\n" {
        advance()
        let token = Token(
          type: .newline,
          value: "\n",
          position: pos,
          stringRange: tokenStart..<current
        )
        line += 1
        column = 1
        return token
      } else {
        fatalError("scanLineEnding called on non-line-ending character")
      }
    }

    private func isLineEnding(_ char: Character) -> Bool {
      return char == "\n" || char == "\r"
    }

    private func scanBracketedToken() -> Token {
      let pos = currentPosition()
      let tokenStart = current
      advance()  // consume '['

      let contentStart = current
      while current < input.endIndex && input[current] != "]" && !isLineEnding(input[current]) {
        advance()
      }

      // If we hit a line ending or EOF before finding the closing bracket, it's an error
      guard current < input.endIndex && input[current] == "]" else {
        // Return error token for unclosed bracket (don't consume past line ending)
        return Token(
          type: .error,
          value: String(input[tokenStart..<current]),
          position: pos,
          stringRange: tokenStart..<current
        )
      }

      let originalContent = String(input[contentStart..<current])
      let lowercaseContent = originalContent.lowercased()
      advance()  // consume ']'

      let tokenEnd = current

      switch lowercaseContent {
      case "note":
        expectingText = true  // Set expectingText when we return a note keyword
        return Token(
          type: .noteKeyword,
          value: "[Note]",
          position: pos,
          stringRange: tokenStart..<tokenEnd
        )
      case "group":
        return Token(
          type: .groupKeyword,
          value: "[Group]",
          position: pos,
          stringRange: tokenStart..<tokenEnd
        )
      case "inertia":
        return Token(
          type: .inertiaKeyword,
          value: "[Inertia]",
          position: pos,
          stringRange: tokenStart..<tokenEnd
        )
      case "evolution":
        return Token(
          type: .evolutionKeyword,
          value: "[Evolution]",
          position: pos,
          stringRange: tokenStart..<tokenEnd
        )
      case "x", "square", "triangle", "circle":
        return Token(
          type: .shapeLabel,
          value: originalContent,  // Use original case, not lowercase
          position: pos,
          stringRange: tokenStart..<tokenEnd
        )
      default:
        if lowercaseContent.allSatisfy({ $0 == "i" }) && lowercaseContent.count <= 4 {
          return Token(
            type: .stageNumber,
            value: lowercaseContent,
            position: pos,
            stringRange: tokenStart..<tokenEnd
          )
        }
        // Return error token for invalid bracketed content
        return Token(
          type: .error,
          value: String(input[tokenStart..<tokenEnd]),
          position: pos,
          stringRange: tokenStart..<tokenEnd
        )
      }
    }

    private func scanNumber() -> Token {
      let pos = currentPosition()
      let tokenStart = current

      // Handle leading decimal point
      if input[current] == "." {
        advance()
      }

      while current < input.endIndex && input[current].isNumber {
        advance()
      }

      // Handle decimal point if we haven't seen one
      if current < input.endIndex && input[current] == "." && input[tokenStart] != "." {
        advance()
        while current < input.endIndex && input[current].isNumber {
          advance()
        }
      }

      let value = String(input[tokenStart..<current])
      return Token(
        type: .realNumber,
        value: value,
        position: pos,
        stringRange: tokenStart..<current
      )
    }

    private func scanStageOrLabel() -> Token {
      let pos = currentPosition()
      let tokenStart = current

      // Check if it's a stage number (i, ii, iii, iv)
      var tempIndex = current
      var iCount = 0

      while tempIndex < input.endIndex && input[tempIndex].lowercased() == "i" && iCount < 4 {
        tempIndex = input.index(after: tempIndex)
        iCount += 1
      }

      // If we have 1-4 i's and next char is whitespace/delimiter, it's a stage
      if iCount > 0 && iCount <= 4 && (tempIndex >= input.endIndex || isDelimiter(input[tempIndex]))
      {
        current = tempIndex
        return Token(
          type: .stageNumber,
          value: String(repeating: "i", count: iCount),
          position: pos,
          stringRange: tokenStart..<current
        )
      }

      // Otherwise, scan as vertex label
      return scanVertexLabel()
    }

    private func scanVertexLabel() -> Token {
      let pos = currentPosition()
      let tokenStart = current

      while current < input.endIndex && !isVertexLabelDelimiter(input[current]) {
        advance()
      }

      let originalValue = String(input[tokenStart..<current]).trimmingCharacters(in: .whitespaces)
      return Token(
        type: .vertexLabel,
        value: originalValue,
        position: pos,
        stringRange: tokenStart..<current
      )
    }

    private func isVertexLabelDelimiter(_ char: Character) -> Bool {
      return char == "-" || char == "+" || char == "," || char == "[" || char == "]" || char == "("
        || char == ")" || isLineEnding(char)
    }

    private func isDelimiter(_ char: Character) -> Bool {
      return char.isWhitespace || isLineEnding(char) || char == "(" || char == ")" || char == "["
        || char == "]" || char == ","
    }

    private func skipWhitespace() {
      while current < input.endIndex && input[current].isWhitespace && !isLineEnding(input[current])
      {
        advance()
      }
    }

    private func advance() {
      if current < input.endIndex {
        current = input.index(after: current)
        column += 1
      }
    }

    private func peek() -> Character? {
      let next = input.index(after: current)
      return next < input.endIndex ? input[next] : nil
    }

    private func currentPosition() -> SourcePosition {
      return SourcePosition(line: line, column: column, stringIndex: current)
    }
  }
}