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

public struct SpecialCompilers {
  public struct TransitionStatement {
    static func Compile(
      node: Node, forState state_identifier: Common.Identifier,
      withStatements stmts: [EvaluatableStatement], withContext context: CompilerContext
    ) -> Result<(ParserState, CompilerContext)> {

      #RequireNodeType<Node, (EvaluatableStatement, CompilerContext)>(
        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), context
                ))
            case (_, .none):
              return .Ok(
                (
                  ParserStateDirectTransition(
                    name: state_identifier,
                    withNextStateIdentifier: next_state_id, withStatements: stmts), context
                ))
            }
          } 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,
            ), context
          ))
      case .Error(let e): .Error(e)
      }
    }
  }

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

      node.enumerateNamedChildren { node in
        switch Statement.Compile(
          node: node, withContext: current_context)
        {
        case .Ok((let parsed_statement, let updated_context)):
          current_context = updated_context
          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, current_context))
    }
  }

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

    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, updated_context)):
        parser.states = parser.states.append(state: state)
        current_context = updated_context
      case Result.Error(let e): error = e
      }
    }

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

    return Result.Ok((parser, current_context))
  }

  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 specific_declaration_node = declaration_node.child(at: 0)!

        let declaration_parsers: [CompilableDeclaration.Type] = [
          Declaration.self, P4Lang.Parser.self,
        ]
        var found_parser = false

        for parser in declaration_parsers {
          switch parser.Compile(node: specific_declaration_node, withContext: compilation_context) {
          case .Ok(.none): {}()
          case .Ok(.some((_, let updated_context))):
            found_parser = true
            compilation_context = updated_context
            break
          case .Error(let e):
            found_parser = true
            errors =
              if let errors = errors {
                errors.append(error: e)
              } else {
                e
              }
            break
          }
        }

        // If none of the declaration parsers chose to parse, that's an error, too!
        if !found_parser {

          print("spec: \(declaration_node)")
          // Try to see if it is an instantiation!
          if declaration_node.nodeType == "instantiation" {
            switch Instantiation.Compile(
              node: declaration_node, withContext: compilation_context)
            {
            case .Ok((let x, let updated_context)):
              print("instantiation: \(x)")
              compilation_context = updated_context
            case .Error(let e):
              errors =
                if let errors = errors {
                  errors.append(error: e)
                } else {
                  e
                }
            }
          } else {
            let no_parser_error = ErrorWithLocation(
              sourceLocation: specific_declaration_node.toSourceLocation(),
              withError: "Could not find parser for declaration node"
            )
            errors =
              if let errors = errors {
                errors.append(error: no_parser_error)
              } else {
                no_parser_error
              }
          }
        }
      }

      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)
    }
  }
}