// p4rse, Copyright 2026, Will Hawkins
//
// This file is part of p4rse.
//
// This file 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://www.gnu.org/licenses/>.

import Common
import P4Lang
import P4Runtime
import SwiftTreeSitter
import TreeSitterExtensions
import TreeSitterP4

extension Statement: Compilable {
  public typealias C = P4Statement
  public static func Compile(
    node: SwiftTreeSitter.Node, withContext context: CompilerContext
  ) -> Result<P4Statement> {

    if node.nodeType != "parserStatement" && node.nodeType != "statement" {
      return Result.Error(
        ErrorWithLocation(
          sourceLocation: node.toSourceLocation(), withError: "Missing expected parser statement")
      )
    }

    let statement = node.child(at: 0)!

    let statementParsers: [String: CompilableStatement.Type] = [
      "assignmentStatement": ParserAssignmentStatement.self,
      "expressionStatement": ExpressionStatement.self,
      "variableDeclaration": VariableDeclarationStatement.self,
      "conditionalStatement": ConditionalStatement.self, "blockStatement": BlockStatement.self,
      "return_statement": ReturnStatement.self,
    ]
    guard let parser = statementParsers[statement.nodeType ?? ""] else {
      return Result.Error(
        ErrorWithLocation(
          sourceLocation: statement.toSourceLocation(),
          withError:
            "Unparseable statement type (\(statement.nodeType ?? "Unknown Statement Type"))"))
    }
    switch parser.CompileStatement(node: statement, withContext: context) {
    case Result.Ok(let parsed):
      return .Ok(parsed)
    case Result.Error(let e):
      return .Error(
        ErrorWithLocation(
          sourceLocation: node.toSourceLocation(),
          withError: "Failed to parse a statement element: \(e)"))
    }
  }
}

extension LocalElements: Compilable {
  public typealias C = P4Statement
  public static func Compile(
    node: Node, withContext context: CompilerContext
  ) -> Result<P4Statement> {
    let localElementsParsers: [String: CompilableStatement.Type] = [
      "variableDeclaration": VariableDeclarationStatement.self
    ]

    guard let parser = localElementsParsers[node.nodeType ?? ""] else {
      return Result.Error(
        ErrorWithLocation(
          sourceLocation: node.toSourceLocation(),
          withError: "Unparseable statement type (\(node.nodeType ?? "Unknown Statement Type"))"))
    }

    switch parser.CompileStatement(node: node, withContext: context) {
    case Result.Ok(let parsed):
      return Result.Ok(parsed)
    case Result.Error(let e):
      return Result.Error(Error(withMessage: "Failed to parse local element: \(e)"))
    }
  }
}

extension ParserState: Compilable {
  public typealias C = ParserState
  public static func Compile(
    node: Node, withContext context: CompilerContext
  ) -> Result<ParserState> {
    var walker = Walker(node: node)

    var current_node: Node? = .none

    guard let node_type = node.nodeType,
      node_type == "parserState"
    else {
      return Result.Error(
        ErrorWithLocation(
          sourceLocation: node.toSourceLocation(),
          withError: "Did not find a parser state declaration"))
    }

    #MustOr(
      result: current_node, thing: walker.getNext(),
      or: Result<ParserState>.Error(
        ErrorWithLocation(
          sourceLocation: node.toSourceLocation(),
          withError: "Missing elements in parser state declaration")))

    if current_node!.nodeType == "annotations" {
      return Result.Error(
        ErrorWithLocation(
          sourceLocation: current_node!.toSourceLocation(),
          withError: "Annotations in parser state are not yet handled."))

      // Would increment here.
    }

    // Skip the keyword state
    walker.next()
    #MustOr(
      result: current_node, thing: walker.getNext(),
      or: Result<ParserState>.Error(
        ErrorWithLocation(
          sourceLocation: node.toSourceLocation(),
          withError: "Missing elements in parser state declaration")))

    let maybe_state_identifier = Identifier.Compile(
      node: current_node!, withContext: context)
    guard case Result.Ok(let state_identifier) = maybe_state_identifier else {
      return Result.Error(maybe_state_identifier.error()!)
    }

    walker.next()
    // Skip the '{'
    walker.next()
    #MustOr(
      result: current_node, thing: walker.getNext(),
      or: Result<ParserState>.Error(
        ErrorWithLocation(
          sourceLocation: node.toSourceLocation(), withError: "Missing body of state declaration")
      ))

    var errors: (any Errorable)? = .none
    var parsed_s: [P4Statement] = Array()

    if current_node!.nodeType == "parserStatements" {
      switch SpecialCompilers.Statements.Compile(
        node: current_node!, withContext: context)
      {
      case .Ok(let state_statements):
        parsed_s = state_statements
      case .Error(let error):
        errors =
          if let errors = errors {
            errors.append(error: error)
          } else {
            error
          }
      }
      walker.next()
    }

    if let errors = errors {
      return .Error(errors)
    }

    #MustOr(
      result: current_node, thing: walker.getNext(),
      or: Result<ParserState>.Error(
        ErrorWithLocation(
          sourceLocation: node.toSourceLocation(),
          withError: "Missing transition statement of state declaration")))

    let updated_context = SpecialCompilers.Statements.effect(statements: parsed_s, context: context)
      .update(newLexicalContextName: state_identifier).update(newLexicalContextStatements: parsed_s)

    return TransitionStatement.Compile(node: current_node!, withContext: updated_context)
  }
}