// 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.
import Foundation

extension Wmap {

  struct ParsedMap {
    let entities: [Entity]
    let vertexLabels: Set<String>
  }

  class Parser {
    private let lexer: Lexer
    private var currentToken: Token
    private var vertexLabelMap: [String: String] = [:]

    init(lexer: Lexer) {
      self.lexer = lexer
      self.currentToken = lexer.nextToken()
    }

    func parse() -> ParsedMap {
      var entities: [Entity] = []
      var vertexLabels: Set<String> = Set()

      while currentToken.type != .eof {
        if currentToken.type == .newline {
          advance()
          continue
        }

        do {
          let entity = try parseEntity()
          entities.append(entity)
          switch entity {
          case .vertex(let vertex):
            vertexLabels.insert(vertex.label.lowercased())
          default:
            break
          }

          if currentToken.type == .newline {
            advance()
          } else if currentToken.type != .eof {
            throw ParseError.expectedNewline(actual: currentToken)
          }
        } catch {
          skipToNextLine()
        }
      }

      return ParsedMap(entities: entities, vertexLabels: vertexLabels)
    }

    private func parseEntity() throws -> Entity {
      switch currentToken.type {
      case .vertexLabel:
        return try parseVertexOrEdge()
      case .noteKeyword:
        return .note(try parseNote())
      case .stageNumber:
        return .stage(try parseStage())
      case .groupKeyword:
        return .group(try parseGroup())
      case .inertiaKeyword:
        return .inertia(try parseInertia())
      case .evolutionKeyword:
        return .evolution(try parseEvolution())
      case .error:
        // Error tokens should cause the line to be ignored
        throw ParseError.invalidLine
      default:
        throw ParseError.invalidLine
      }
    }

    private func parseVertexOrEdge() throws -> Entity {
      let startPos = currentToken.position
      let firstLabel = currentToken.value
      let firstLabelRange = makeSourceRange(from: currentToken)

      vertexLabelMap[firstLabel.lowercased()] = firstLabel
      advance()

      if currentToken.type == .edgeUndirected || currentToken.type == .edgeDirected {
        let isDirected = currentToken.type == .edgeDirected
        let edgeTypeRange = makeSourceRange(from: currentToken)
        advance()

        guard currentToken.type == .vertexLabel else {
          throw ParseError.expectedVertexLabel(actual: currentToken)
        }

        let secondLabel = currentToken.value
        let secondLabelRange = makeSourceRange(from: currentToken)
        let endPos = currentToken.position

        vertexLabelMap[secondLabel.lowercased()] = secondLabel
        advance()

        return .edge(
          Edge(
            from: firstLabel,
            to: secondLabel,
            isDirected: isDirected,
            position: startPos,
            range: SourceRange(
              start: startPos, end: endPos,
              stringRange: firstLabelRange.stringRange
                .lowerBound..<secondLabelRange.stringRange.upperBound),
            fromLabelRange: firstLabelRange,
            edgeTypeRange: edgeTypeRange,
            toLabelRange: secondLabelRange
          ))
      } else {
        // Parse vertex
        let (coordinates, coordRanges) = try parsePositionWithRanges()
        let (shape, shapeRange) = try parseOptionalShapeWithRange()

        let endPos = shapeRange?.end ?? coordRanges.rightParen.end

        return .vertex(
          Vertex(
            label: firstLabel,
            position: startPos,
            coordinates: coordinates,
            shape: shape,
            range: SourceRange(
              start: startPos, end: endPos,
              stringRange: firstLabelRange.stringRange
                .lowerBound..<(shapeRange?.stringRange.upperBound
                ?? coordRanges.rightParen.stringRange.upperBound)),
            labelRange: firstLabelRange,
            leftParenRange: coordRanges.leftParen,
            xCoordRange: coordRanges.xCoord,
            commaRange: coordRanges.comma,
            yCoordRange: coordRanges.yCoord,
            rightParenRange: coordRanges.rightParen,
            shapeRange: shapeRange
          ))
      }
    }

    private func parseNote() throws -> Note {
      let startPos = currentToken.position
      let keywordRange = makeSourceRange(from: currentToken)
      advance()

      let (coordinates, coordRanges) = try parsePositionWithRanges()

      guard currentToken.type == .text else {
        throw ParseError.expectedText(actual: currentToken)
      }
      let text = currentToken.value
      let textRange = makeSourceRange(from: currentToken)
      let endPos = currentToken.position
      advance()

      return Note(
        coordinates: coordinates,
        text: text,
        position: startPos,
        range: SourceRange(
          start: startPos, end: endPos,
          stringRange: keywordRange.stringRange.lowerBound..<textRange.stringRange.upperBound),
        keywordRange: keywordRange,
        leftParenRange: coordRanges.leftParen,
        xCoordRange: coordRanges.xCoord,
        commaRange: coordRanges.comma,
        yCoordRange: coordRanges.yCoord,
        rightParenRange: coordRanges.rightParen,
        textRange: textRange
      )
    }

    private func parseStage() throws -> Stage {
      let startPos = currentToken.position
      let stageToken = currentToken
      let stageValue = stageToken.value

      // For bracketed stage numbers like [ii], we need to parse the components
      // The lexer gives us the full token, but we need to break it down
      let tokenRange = stageToken.stringRange

      // Find the bracket positions within the token
      let leftBracketRange = SourceRange(
        start: stageToken.position,
        end: SourcePosition(
          line: stageToken.position.line,
          column: stageToken.position.column + 1,
          stringIndex: lexer.input.index(after: tokenRange.lowerBound)
        ),
        stringRange: tokenRange.lowerBound..<lexer.input.index(after: tokenRange.lowerBound)
      )

      let stageNumberStart = lexer.input.index(after: tokenRange.lowerBound)
      let stageNumberEnd = lexer.input.index(tokenRange.upperBound, offsetBy: -1)
      let stageNumberRange = SourceRange(
        start: SourcePosition(
          line: stageToken.position.line,
          column: stageToken.position.column + 1,
          stringIndex: stageNumberStart
        ),
        end: SourcePosition(
          line: stageToken.position.line,
          column: stageToken.position.column + 1 + stageValue.count,
          stringIndex: stageNumberEnd
        ),
        stringRange: stageNumberStart..<stageNumberEnd
      )

      let rightBracketRange = SourceRange(
        start: SourcePosition(
          line: stageToken.position.line,
          column: stageToken.position.column + 1 + stageValue.count,
          stringIndex: stageNumberEnd
        ),
        end: SourcePosition(
          line: stageToken.position.line,
          column: stageToken.position.column + 2 + stageValue.count,
          stringIndex: tokenRange.upperBound
        ),
        stringRange: stageNumberEnd..<tokenRange.upperBound
      )

      advance()

      guard currentToken.type == .realNumber else {
        throw ParseError.expectedNumber(actual: currentToken)
      }

      guard let value = Double(currentToken.value) else {
        throw ParseError.invalidNumber(currentToken.value)
      }

      let valueRange = makeSourceRange(from: currentToken)
      let endPos = currentToken.position
      advance()

      return Stage(
        number: stageValue,
        value: value,
        position: startPos,
        range: SourceRange(
          start: startPos,
          end: endPos,
          stringRange: leftBracketRange.stringRange.lowerBound..<valueRange.stringRange.upperBound
        ),
        leftBracketRange: leftBracketRange,
        stageNumberRange: stageNumberRange,
        rightBracketRange: rightBracketRange,
        valueRange: valueRange
      )
    }

    private func parseGroup() throws -> Group {
      let startPos = currentToken.position
      let keywordRange = makeSourceRange(from: currentToken)
      advance()

      var vertices: [String] = []
      var vertexRanges: [SourceRange] = []
      var commaRanges: [SourceRange] = []

      guard currentToken.type == .vertexLabel else {
        throw ParseError.expectedVertexLabel(actual: currentToken)
      }

      vertices.append(currentToken.value)
      vertexRanges.append(makeSourceRange(from: currentToken))
      advance()

      while currentToken.type == .comma {
        commaRanges.append(makeSourceRange(from: currentToken))
        advance()

        guard currentToken.type == .vertexLabel else {
          throw ParseError.expectedVertexLabel(actual: currentToken)
        }

        vertices.append(currentToken.value)
        vertexRanges.append(makeSourceRange(from: currentToken))
        advance()
      }

      let endPos = vertexRanges.last?.end ?? startPos

      return Group(
        vertices: vertices,
        position: startPos,
        range: SourceRange(
          start: startPos, end: endPos,
          stringRange: keywordRange.stringRange
            .lowerBound..<vertexRanges.last!.stringRange.upperBound),
        keywordRange: keywordRange,
        vertexRanges: vertexRanges,
        commaRanges: commaRanges
      )
    }

    private func parseInertia() throws -> Inertia {
      let startPos = currentToken.position
      let keywordRange = makeSourceRange(from: currentToken)
      advance()

      guard currentToken.type == .vertexLabel else {
        throw ParseError.expectedVertexLabel(actual: currentToken)
      }

      let vertex = currentToken.value
      let vertexRange = makeSourceRange(from: currentToken)
      let endPos = currentToken.position
      advance()

      return Inertia(
        vertex: vertex,
        position: startPos,
        range: SourceRange(
          start: startPos, end: endPos,
          stringRange: keywordRange.stringRange.lowerBound..<vertexRange.stringRange.upperBound),
        keywordRange: keywordRange,
        vertexRange: vertexRange
      )
    }

    private func parseEvolution() throws -> Evolution {
      let startPos = currentToken.position
      let keywordRange = makeSourceRange(from: currentToken)
      advance()

      guard currentToken.type == .vertexLabel else {
        throw ParseError.expectedVertexLabel(actual: currentToken)
      }

      let vertex = currentToken.value
      let vertexRange = makeSourceRange(from: currentToken)
      advance()

      let isPositive: Bool
      let signRange: SourceRange
      if currentToken.type == .plus {
        isPositive = true
        signRange = makeSourceRange(from: currentToken)
        advance()
      } else if currentToken.type == .minus {
        isPositive = false
        signRange = makeSourceRange(from: currentToken)
        advance()
      } else {
        throw ParseError.expectedSign(actual: currentToken)
      }

      guard currentToken.type == .realNumber else {
        throw ParseError.expectedNumber(actual: currentToken)
      }

      guard let value = Double(currentToken.value) else {
        throw ParseError.invalidNumber(currentToken.value)
      }

      let valueRange = makeSourceRange(from: currentToken)
      let endPos = currentToken.position
      advance()

      return Evolution(
        vertex: vertex,
        isPositive: isPositive,
        value: value,
        position: startPos,
        range: SourceRange(
          start: startPos, end: endPos,
          stringRange: keywordRange.stringRange.lowerBound..<valueRange.stringRange.upperBound),
        keywordRange: keywordRange,
        vertexRange: vertexRange,
        signRange: signRange,
        valueRange: valueRange
      )
    }

    // Helper methods
    private func parsePositionWithRanges() throws -> (
      (Double, Double),
      (
        leftParen: SourceRange, xCoord: SourceRange, comma: SourceRange, yCoord: SourceRange,
        rightParen: SourceRange
      )
    ) {
      guard currentToken.type == .leftParen else {
        throw ParseError.expectedLeftParen(actual: currentToken)
      }
      let leftParenRange = makeSourceRange(from: currentToken)
      advance()

      guard currentToken.type == .realNumber else {
        throw ParseError.expectedNumber(actual: currentToken)
      }
      guard let x = Double(currentToken.value) else {
        throw ParseError.invalidNumber(currentToken.value)
      }
      let xCoordRange = makeSourceRange(from: currentToken)
      advance()

      guard currentToken.type == .comma else {
        throw ParseError.expectedComma(actual: currentToken)
      }
      let commaRange = makeSourceRange(from: currentToken)
      advance()

      guard currentToken.type == .realNumber else {
        throw ParseError.expectedNumber(actual: currentToken)
      }
      guard let y = Double(currentToken.value) else {
        throw ParseError.invalidNumber(currentToken.value)
      }
      let yCoordRange = makeSourceRange(from: currentToken)
      advance()

      guard currentToken.type == .rightParen else {
        throw ParseError.expectedRightParen(actual: currentToken)
      }
      let rightParenRange = makeSourceRange(from: currentToken)
      advance()

      return ((x, y), (leftParenRange, xCoordRange, commaRange, yCoordRange, rightParenRange))
    }

    private func parseOptionalShapeWithRange() throws -> (String?, SourceRange?) {
      guard currentToken.type == .shapeLabel else {
        return (nil, nil)
      }

      let shape = currentToken.value
      let shapeRange = makeSourceRange(from: currentToken)
      advance()
      return (shape, shapeRange)
    }

    private func makeSourceRange(from token: Token) -> SourceRange {
      return SourceRange(
        start: token.position,
        end: SourcePosition(
          line: token.position.line,
          column: token.position.column + token.value.count,
          stringIndex: token.stringRange.upperBound
        ),
        stringRange: token.stringRange
      )
    }

    private func advance() {
      currentToken = lexer.nextToken()
    }

    private func skipToNextLine() {
      while currentToken.type != .newline && currentToken.type != .eof {
        advance()
      }
      if currentToken.type == .newline {
        advance()
      }
    }
  }

  enum ParseError: Error, LocalizedError {
    case unexpectedToken(expected: String, actual: Token)
    case expectedVertexLabel(actual: Token)
    case expectedNumber(actual: Token)
    case expectedText(actual: Token)
    case expectedSign(actual: Token)
    case expectedComma(actual: Token)
    case expectedLeftParen(actual: Token)
    case expectedRightParen(actual: Token)
    case expectedNewline(actual: Token)
    case invalidNumber(String)
    case invalidLine

    var errorDescription: String? {
      switch self {
      case .unexpectedToken(let expected, let actual):
        return
          "Expected \(expected), got \(actual.type) '\(actual.value)' at line \(actual.position.line)"
      case .expectedVertexLabel(let actual):
        return "Expected vertex label, got \(actual.type) at line \(actual.position.line)"
      case .invalidNumber(let value):
        return "Invalid number format: \(value)"
      // ... other cases
      default:
        return "Parse error"
      }
    }
  }
}