// 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 BlockStatement: Compilable {
  public typealias C = BlockStatement
  public static func Compile(
    node: Node, withContext context: CompilerContext
  ) -> Result<BlockStatement> {
    #RequireNodeType<Node, (P4Statement)>(
      node: node, type: "blockStatement", nice_type_name: "block statement")

    var walker = Walker(node: node)
    var current_node: Node? = .none

    #MustOr(
      result: current_node, thing: walker.getNext(),
      or: Result<BlockStatement>.Error(
        ErrorWithLocation(
          sourceLocation: node.toSourceLocation(), withError: "Malformed block statement")))

    if current_node!.nodeType != "{" {
      return Result.Error(
        ErrorWithLocation(
          sourceLocation: current_node!.toSourceLocation(),
          withError: "Missing { on block statement"))
    }

    var statements: [P4Statement] = Array()
    var parse_err: (any Errorable)? = .none

    walker.next()
    #MustOr(
      result: current_node, thing: walker.getNext(),
      or: Result<BlockStatement>.Error(
        ErrorWithLocation(
          sourceLocation: node.toSourceLocation(), withError: "Malformed block statement")))

    if current_node!.nodeType == "statements" {
      switch SpecialCompilers.Statements.Compile(
        node: current_node!, withContext: context)
      {
      case .Ok(let parsed_statements):
        statements = parsed_statements
      case .Error(let error):
        parse_err = error
      }

      walker.next()
    }

    if let err = parse_err {
      return .Error(err)
    }

    #MustOr(
      result: current_node, thing: walker.getNext(),
      or: Result<BlockStatement>.Error(
        ErrorWithLocation(
          sourceLocation: node.toSourceLocation(), withError: "Malformed block statement")))

    if current_node!.nodeType != "}" {
      return Result.Error(
        ErrorWithLocation(
          sourceLocation: current_node!.toSourceLocation(),
          withError: "Missing } on block statement"))
    }

    return .Ok(BlockStatement(statements))
  }
}

@deriveCompilableStatement
extension BlockStatement: CompilableStatement {}

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

    #RequireNodeType<Node, (P4Statement)>(
      node: node, type: "conditionalStatement", nice_type_name: "conditional statement")

    let maybe_condition_expression = node.child(at: 2)
    guard let condition_expression = maybe_condition_expression,
      condition_expression.nodeType == "expression"
    else {
      return Result.Error(
        ErrorWithLocation(
          sourceLocation: node.toSourceLocation(),
          withError: "Did not find condition for conditional statement"))
    }

    let maybe_thens = node.child(at: 4)
    guard let thens = maybe_thens,
      thens.nodeType == "statement"
    else {
      return Result.Error(
        ErrorWithLocation(
          sourceLocation: node.toSourceLocation(),
          withError: "Did not find then statement block for conditional statement"))
    }

    guard
      case .Ok(let condition) = Expression.Compile(
        node: condition_expression, withContext: context)
    else {
      return Result.Error(
        Error(withMessage: "Could not parse a conditional expression in a conditional statement"))
    }

    guard
      case .Ok(let thenns) = Statement.Compile(
        node: thens, withContext: context)
    else {
      return Result.Error(
        Error(
          withMessage:
            "Could not parse the then block in a conditional statement"))
    }

    let optional_elss: Result<any P4Statement>? =
      if let elss = node.child(at: 6) {
        .some(
          Statement.Compile(
            node: elss, withContext: context))
      } else {
        .none
      }

    if let parsed_elss = optional_elss {
      guard
        case .Ok(let elss) = parsed_elss
      else {
        return Result.Error(
          Error(
            withMessage:
              "Could not parse the else block in a conditional statement"))
      }
      return .Ok(
        ConditionalStatement(condition: condition, withThen: thenns, andElse: elss))
    }
    return .Ok(ConditionalStatement(condition: condition, withThen: thenns))
  }
}

@deriveCompilableStatement
extension ConditionalStatement: CompilableStatement {}

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

    #RequireNodeType<Node, (P4Statement)>(
      node: node, type: "variableDeclaration", nice_type_name: "variable declaration statement")

    let maybe_typeref = node.child(at: 0)
    guard let typeref = maybe_typeref,
      typeref.nodeType == "typeRef"
    else {
      return Result.Error(
        ErrorWithLocation(
          sourceLocation: node.toSourceLocation(),
          withError: "Did not find type name for variable declaration statement"))
    }

    let maybe_variablename = node.child(at: 1)
    guard let variablename = maybe_variablename,
      variablename.nodeType == "identifier"
    else {
      return Result.Error(
        ErrorWithLocation(
          sourceLocation: node.toSourceLocation(),
          withError: "Did not find identifier name for variable declaration statement"))
    }

    let maybe_rvalue = node.childCount > 3 ? node.child(at: 3) : .none

    guard
      case .Ok(let parsed_variablename) = Identifier.Compile(
        node: variablename, withContext: context)
    else {
      return Result.Error(
        Error(withMessage: "Could not parse variable name"))
    }

    guard case .Ok(let declaration_p4_type) = Types.CompileType(type: typeref, withContext: context)
    else {
      return Result.Error(
        Error(withMessage: "Could not parse a P4 type from \(typeref.text!)"))
    }

    var initializer: P4Expression? = .none

    // If there is an initializer, it must be an expression.
    if let initializer_expression = maybe_rvalue {
      guard initializer_expression.nodeType == "expression" else {
        return Result.Error(
          ErrorWithLocation(
            sourceLocation: node.toSourceLocation(),
            withError: "initial value for declaration statement is not an expression"))
      }

      let maybe_parsed_rvalue = Expression.Compile(
        node: initializer_expression, withContext: context)
      guard
        case .Ok(let parsed_rvalue) = maybe_parsed_rvalue
      else {
        return .Error(maybe_parsed_rvalue.error()!)
      }

      if parsed_rvalue.type().eq(declaration_p4_type) {
        initializer = parsed_rvalue
      } else {
        return Result.Error(
          Error(
            withMessage:
              "Cannot initialize \(parsed_variablename) (with type \(declaration_p4_type)) from expression with type \(parsed_rvalue.type())"
          ))
      }
    }

    // If there is no initializer, then it must be defaultable.

    if initializer == nil {
      initializer = declaration_p4_type.def()
    }

    guard let initializer = initializer else {
      return Result.Error(
        ErrorWithLocation(
          sourceLocation: node.toSourceLocation(), withError: "No initializer for declaration"))
    }

    return Result.Ok(
      VariableDeclarationStatement(
        identifier: TypedIdentifier(id: parsed_variablename, withType: declaration_p4_type),
        withInitializer: initializer),
    )
  }
}

@deriveCompilableStatement
extension VariableDeclarationStatement: CompilableStatement {}

extension ExpressionStatement: Compilable {
  public typealias C = ExpressionStatement
  public static func Compile(
    node: Node, withContext context: CompilerContext
  ) -> Result<ExpressionStatement> {
    #RequireNodeType<Node, (P4Statement)>(
      node: node, type: "expressionStatement", nice_type_name: "expression statement")

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

    return switch Expression.Compile(node: expression_node, withContext: context) {
    case .Ok(let expression): .Ok(ExpressionStatement(expression))
    case .Error(let e): .Error(e)
    }
  }
}

@deriveCompilableStatement
extension ExpressionStatement: CompilableStatement {}

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

    #RequireNodeType<Node, (P4Statement)>(
      node: node, type: "assignmentStatement", nice_type_name: "assignment statement")

    guard let lvalue_node = node.child(at: 0),
      lvalue_node.nodeType == "expression"
    else {
      return Result.Error(
        ErrorWithLocation(
          sourceLocation: node.toSourceLocation(),
          withError: "Missing lvalue in assignment statement"))
    }

    guard let rvalue_node = node.child(at: 2),
      rvalue_node.nodeType == "expression"
    else {
      return Result.Error(
        ErrorWithLocation(
          sourceLocation: node.toSourceLocation(),
          withError: "Missing rvalue in assignment statement"))
    }

    let maybe_parsed_rvalue = Expression.Compile(
      node: rvalue_node, withContext: context)
    guard case Result.Ok(let rvalue) = maybe_parsed_rvalue else {
      return Result.Error(maybe_parsed_rvalue.error()!)
    }

    let maybe_parsed_lvalue = LValue.Compile(node: lvalue_node, withContext: context)
    guard case .Ok(let lvalue_identifier) = maybe_parsed_lvalue else {
      return Result.Error(maybe_parsed_lvalue.error()!)
    }

    let check_result = lvalue_identifier.check(to: rvalue, inScopes: context.instances)
    guard case .Ok(_) = check_result else {
      return Result.Error(
        ErrorWithLocation(
          sourceLocation: lvalue_node.toSourceLocation(),
          withError: "\(check_result.error()!)"))
    }

    return Result.Ok(
      ParserAssignmentStatement(
        withLValue: lvalue_identifier,
        withValue: rvalue
      ))
  }
}

@deriveCompilableStatement
extension ParserAssignmentStatement: CompilableStatement {}

extension ReturnStatement: Compilable {
  public typealias C = ReturnStatement
  public static func Compile(
    node: SwiftTreeSitter.Node, withContext context: CompilerContext
  ) -> Result<ReturnStatement> {
    #RequireNodeType<Node, (P4Statement)>(
      node: node, type: "return_statement", nice_type_name: "return statement")

    let expression_node = node.child(at: 1)!

    return switch Expression.Compile(node: expression_node, withContext: context) {
    case .Ok(let result):
      if result.type().baseType().eq(rhs: context.expected_type!.baseType()) {
        .Ok(ReturnStatement(result))
      } else {
        .Error(
          ErrorWithLocation(
            sourceLocation: node.toSourceLocation(),
            withError:
              "Type of expression in return statement (\(result.type())) is not compatible with function return type (\(context.expected_type!))"
          ))
      }
    case .Error(let e): .Error(e)
    }
  }
}

@deriveCompilableStatement
extension ReturnStatement: CompilableStatement {}

extension ApplyStatement: Compilable {
  public typealias C = ApplyStatement
  public static func Compile(
    node: SwiftTreeSitter.Node, withContext context: CompilerContext
  ) -> Result<ApplyStatement> {
    #RequireNodeType<Node, (P4Statement)>(
      node: node, type: "apply_statement", nice_type_name: "apply statement")

    let expression_node = node.child(at: 1)!

    return switch BlockStatement.Compile(node: expression_node, withContext: context) {
    case .Ok(let statement):
      .Ok(ApplyStatement(statement))
    case .Error(let e): .Error(e)
    }
  }
}

@deriveCompilableStatement
extension ApplyStatement: CompilableStatement {}

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

    let expression = node
    #RequireNodeType<Node, (Instantiation)>(
      node: expression, type: "instantiation", nice_type_name: "instantiation statement")

    var walker = Walker(node: expression)
    var current_node: Node? = .none

    #MustOr(
      result: current_node, thing: walker.getNext(),
      or: Result<Instantiation>.Error(
        ErrorWithLocation(
          sourceLocation: node.toSourceLocation(), withError: "Missing function call component")))

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

    var maybe_instantiated_type: Result<(P4Lang.Parser?, Declaration?)> =
      switch context.types.lookup(identifier: instantiated_type_name) {
      case .Ok(let looked_up):
        switch looked_up {
        /// TODO: Further filter instantiable things.
        case let instantiated_parser as P4Lang.Parser:
          Result<(P4Lang.Parser?, Declaration?)>.Ok((instantiated_parser, .none))  // What we found is actually a parser declaration
        default:
          Result<(P4Lang.Parser?, Declaration?)>.Error(
            ErrorWithLocation(
              sourceLocation: current_node!.toSourceLocation(),
              withError: "\(instantiated_type_name) cannot be instantiated"))
        }
      case .Error(let e): Result<(P4Lang.Parser?, Declaration?)>.Error(e)
      }

    maybe_instantiated_type =
      if case .Error(let e) = maybe_instantiated_type {
        switch context.externs.lookup(identifier: instantiated_type_name) {
        case .Ok(let callee as Declaration):
          // Now, make sure that it is a function declaration!
          switch callee.identifier.type.baseType() {
          case is P4Lang.Parser: Result.Ok((.none, callee))
          default:
            .Error(
              ErrorWithLocation(
                sourceLocation: current_node!.toSourceLocation(),
                withError: "\(instantiated_type_name) cannot be instantiated"))
          }
        default: .Error(e)
        }
      } else {
        maybe_instantiated_type
      }

    guard case .Ok(let callee) = maybe_instantiated_type else {
      return .Error(maybe_instantiated_type.error()!)
    }

    walker.next()

    #MustOr(
      result: current_node, thing: walker.getNext(),
      or: Result<Instantiation>.Error(
        ErrorWithLocation(
          sourceLocation: node.toSourceLocation(), withError: "Missing instantiation component")))

    let maybe_argument_list = ArgumentList.Compile(node: current_node!, withContext: context)

    guard case .Ok(let arguments) = maybe_argument_list else {
      return .Error(maybe_argument_list.error()!)
    }

    // Now, compare the arguments with the parameters:

    let params =
      switch callee {
      case (.some(let callee), .none): Optional<ParameterList>.some(callee.parameters)
      case (.none, .some(let callee)):
        Optional<ParameterList>.some((callee.ffi!.type().baseType() as! P4Lang.Parser).parameters)
      default: Optional<ParameterList>.none
      }

    guard case .some(let params) = params else {
      return Result.Error(
        ErrorWithLocation(
          sourceLocation: node.toSourceLocation(),
          withError:
            "Could not lookup the parameters for the instantiated parser (\(instantiated_type_name))"
        ))
    }

    if case .Error(let e) = arguments.compatible(params) {
      return .Error(e)
    }

    walker.next()

    #MustOr(
      result: current_node, thing: walker.getNext(),
      or: Result<Instantiation>.Error(
        ErrorWithLocation(
          sourceLocation: node.toSourceLocation(), withError: "Missing instantiation name")))

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

    let inst: Result<Instantiation> =
      switch callee {
      case (.some(let callee), .none):
        .Ok(Instantiation(named: name, ofType: callee, withArguments: arguments))
      case (.none, .some(let callee)):
        .Ok(Instantiation(named: name, ofType: callee, withArguments: arguments))
      default:
        Result.Error(
          ErrorWithLocation(
            sourceLocation: node.toSourceLocation(),
            withError:
              "Unexpected error occurred calling function named (\(instantiated_type_name))"
          ))
      }
    return inst
  }
}

@deriveCompilableStatement
extension Instantiation: CompilableStatement {}