// 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 SwiftTreeSitter
import TreeSitterExtensions
import TreeSitterP4

extension AST.TransitionStatement: Compilable {
  public static func Compile(
    node: Node, withContext context: ASTCompilerContext
  ) -> Result<AST.AnState> {

    guard let state_identifier = context.lexical_context_name else {
      return .Error(
        ErrorWithLocation(
          sourceLocation: node.toSourceLocation(),
          withError: "Cannot parse a transition statement without the name of the containing state."
        ))
    }

    guard let stmts = context.lexical_context_statements else {
      return .Error(
        ErrorWithLocation(
          sourceLocation: node.toSourceLocation(),
          withError:
            "Cannot parse a transition statement without statements of the containing state."))
    }

    #RequireNodeType<Node, P4Statement>(
      node: node, type: "parserTransitionStatement", nice_type_name: "parser transition statement"
    )

    guard let tse_node = node.child(at: 1),
      tse_node.nodeType! == "transitionSelectionExpression"
    else {
      return .Error(
        ErrorWithLocation(
          sourceLocation: node.toSourceLocation(),
          withError: "Could not find transition select expression"))
    }

    guard let next_node = tse_node.child(at: 0) else {
      return .Error(
        ErrorWithLocation(
          sourceLocation: node.toSourceLocation(),
          withError: "Could not find the next token in a transition selection expression"))
    }

    // If the next node is an identifier, we have the simple form ...
    if next_node.nodeType == "identifier" {
      let maybe_parsed_next_state_id = AST.Identifier.CompileExpression(
        node: next_node, withContext: context)
      switch maybe_parsed_next_state_id {
      case .Ok(let next_state_id):
        return .Ok(
          AST.ParserStateDirectTransition(
            name: (state_identifier),
            withNextStateIdentifier: next_state_id as! AST.Identifier, withStatements: stmts))
      case .Error(let e):
        return .Error(e)
      }
    }

    // We know that the next node is a select expression.
    return
      switch AST.SelectExpression.CompileExpression(node: next_node, withContext: context)
    {
    case .Ok(let tse):
      .Ok(
        AST.ParserStateSelectTransition(
          name: state_identifier, withTransitionExpression: tse as! AST.SelectExpression,
          withStatements: stmts,
        )
      )
    case .Error(let e): .Error(e)
    }
  }
}

public struct SpecialCompilers {
  public struct Statements {
    static func Compile(
      node: Node, withContext context: ASTCompilerContext
    ) -> Result<[AST.AnStatement]> {
      if node.nodeType != "statements" && node.nodeType != "parserStatements" {
        return Result.Error(
          ErrorWithLocation(
            sourceLocation: node.toSourceLocation(), withError: "Did not find expected statements"))
      }

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

      node.enumerateNamedChildren { node in
        switch AST.Statement.Compile(
          node: node, withContext: context)
        {
        case .Ok(let parsed_statement):
          parsed_s.append(parsed_statement)
        case .Error(let e):
          errors =
            if let errors = errors {
              errors.append(error: e)
            } else {
              e
            }
        }
      }

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

      return Result.Ok(parsed_s)
    }
  }

  static func CompileParserBody(
    withName name: AST.Identifier, withParameters parameters: AST.ParameterList, node: Node,
    withContext context: ASTCompilerContext
  ) -> Result<AST.Parser> {

    var parser = AST.Parser(withName: name, withParameters: parameters)

    // Build a state from each one listed.
    var errors: (any Errorable)? = .none

    /// TODO: Assert that there is only one.
    node.enumerateNamedChildren { parser_state in
      if parser_state.nodeType != "parserState" {
        return
      }

      // Parse a state in a nested scope.
      switch AST.ParserState.Compile(
        node: parser_state,
        withContext: context)
      {
      case Result.Ok(let state):
        // All states are instances inside the parser.
        parser.states = parser.states.append(state: state)
      case Result.Error(let e):
        errors =
          if let errors = errors {
            errors.append(error: e)
          } else {
            e
          }
      }
    }

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

    return .Ok(parser)
  }

  public struct ProgramCompiler {
    public static func Compile(_ source: String) -> Result<AST.Program> {

      let maybe_parser = ConfigureP4Parser()
      guard case .Ok(let p) = maybe_parser else {
        return .Error(maybe_parser.error()!)
      }

      let result = p.parse(source)
      guard let tree = result,
        !tree.isError(lang: p4lang),
        !tree.containsMissing(lang: p4lang)
      else {
        return Result.Error(Error(withMessage: "Could not compile the P4 program"))
      }

      var program = AST.Program()

      // Set up a context for parsing.
      let compilation_context = ASTCompilerContext()

      var errors: (any Errorable)? = .none

      // Try to parse all top-level declarations.
      result?.rootNode?.enumerateNamedChildren { (declaration_node: Node) in
        let declaration_parsers: [String: CompilableStatement.Type] = [
          "declaration": AST.Declaration.self,
          "instantiation": AST.Instantiation.self,
        ]

        if let parser = declaration_parsers[declaration_node.nodeType!] {
          let r = parser.CompileStatement(node: declaration_node, withContext: compilation_context)
          switch r {
          case .Ok(let compiled):
            program.statements = program.statements + [compiled]
          case .Error(let e):
            errors =
              if let errors = errors {
                errors.append(error: e)
              } else {
                e
              }

          }
        } else {
          let e = ErrorWithLocation(
            sourceLocation: declaration_node.toSourceLocation(),
            withError:
              "\(declaration_node.nodeType!) cannot be at a P4 program top level")
          errors =
            if let errors = errors {
              errors.append(error: e)
            } else {
              e
            }
        }
      }

      if let errors = errors {
        return .Error(errors)
      }
      return Result.Ok(program)

    }
  }

}