// 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

protocol AnyCompilable {
  static func Compile(
    node: SwiftTreeSitter.Node, withContext context: CompilerContext
  ) -> Common.Result<(Any, CompilerContext)>
}

extension TransitionStatement: Compilable {
  public static func Compile(
    node: Node, withContext context: CompilerContext
  ) -> Result<ParserState> {

    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 = Identifier.Compile(
        node: next_node, withContext: context)
      if case .Ok(let next_state_id) = maybe_parsed_next_state_id {
        if case .Ok(let next_state) = context.instances.lookup(identifier: next_state_id) {
          switch next_state {
          case (_, .some(let instance)):
            return .Ok(
              ParserStateDirectTransition(
                name: state_identifier,
                withNextState: instance.dataValue() as! InstantiatedParserState,
                withStatements: stmts,
              )
            )
          case (_, .none):
            return .Ok(
              ParserStateDirectTransition(
                name: state_identifier,
                withNextStateIdentifier: next_state_id, withStatements: stmts)
            )
          }
        } else {

          return .Error(
            Error(
              withMessage:
                "\(next_state_id) does not name a parser state in scope"
            ))
        }
      } else {
        return .Error(
          Error(
            withMessage:
              "Could not parse the next state in a transition statement: \(maybe_parsed_next_state_id.error()!)"
          ))
      }
    }

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

public struct SpecialCompilers {
  public struct Statements {
    static func Compile(
      node: Node, withContext context: CompilerContext
    ) -> Result<[P4Statement]> {
      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 current_context = context
      var parsed_s: [P4Statement] = Array()

      node.enumerateNamedChildren { node in
        switch Statement.Compile(
          node: node, withContext: current_context)
        {
        case .Ok(let parsed_statement):
          parsed_s.append(parsed_statement)
          current_context = parsed_statement.effect(context: current_context)
        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 effect(statements: [P4Statement], context: CompilerContext) -> CompilerContext {
      var current = context
      for s in statements {
        current = s.effect(context: current)
      }
      return current
    }

  }

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

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

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

    var current_context = context
    /// 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 ParserState.Compile(
        node: parser_state,
        withContext: context.update(newInstances: current_context.instances.enter()))
      {
      case Result.Ok(let state):
        let statement = state as P4Statement
        current_context = statement.effect(context: current_context)
        parser.states = parser.states.append(state: state)
      case Result.Error(let e): error = e
      }
    }

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

    return Result.Ok(parser)
  }

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

      // Certain names are always in scope during compilation.
      var globals = StaticVarValueScopes().enter()
      globals = globals.declare(
        identifier: accept.state().getName(),
        withValue: (P4QualifiedType(accept.type()), P4Value(accept))
      )
      .declare(
        identifier: reject.state.getName(),
        withValue: (P4QualifiedType(reject.type()), P4Value(reject)))

      return ProgramCompiler.Compile(
        source, withGlobalInstances: globals, withGlobalTypes: .none, withFFIs: [])
    }

    public static func Compile(
      _ source: String, withGlobalInstances globalInstances: StaticVarValueScopes
    ) -> Result<P4Lang.Program> {
      return ProgramCompiler.Compile(
        source, withGlobalInstances: globalInstances, withGlobalTypes: .none, withFFIs: [])
    }

    public static func Compile(
      _ source: String, withGlobalInstances globalInstances: StaticVarValueScopes?,
      withGlobalTypes globalTypes: TypeTypeScopes?, withFFIs ffis: [P4FFI] = Array()
    ) -> Result<P4Lang.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 = P4Lang.Program()

      // Set up a context for parsing.
      var compilation_context = CompilerContext()

      // Add our FFIs
      compilation_context = compilation_context.update(newFFIs: ffis)

      var errors: (any Errorable)? = .none

      // If the caller gave any global instances, add them here.
      if let globalInstances = globalInstances {
        compilation_context = compilation_context.update(newInstances: globalInstances)
      }

      // If the caller gave any global types, add them here.
      if let globalTypes = globalTypes {
        compilation_context = compilation_context.update(newTypes: globalTypes)
      }

      // Try to parse all top-level declarations.
      result?.rootNode?.enumerateNamedChildren { (declaration_node: Node) in
        let declaration_parsers: [String: CompilableStatement.Type] = [
          "declaration": Declaration.self,
          "instantiation": 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):
            compilation_context = compiled.effect(context: compilation_context)
          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)
      }

      // Any of the instances that are in the top-level scope should go into the program!
      program.instances = Array(
        compilation_context.instances.filter { (_, v) in
          v.1 != nil
        }.map { (_, v) in
          v.1!
        })

      // Any of the types that are in the top-level scope should go into the program!
      program.types = Array(
        compilation_context.types.map { (_, v) in
          v
        })

      // Any of the extern types that are in the top-level scope should go into the program!
      program.externs = Array(
        compilation_context.externs.map { (_, v) in
          v
        })

      return Result.Ok(program)
    }
  }

}