// 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 Foundation
import Macros
import P4Lang
import P4Runtime
import SwiftTreeSitter
import Testing
import TreeSitter
import TreeSitterP4

@testable import P4Compiler

@Test func test_control_single_key() async throws {
  let simple_parser_declaration = """
    control simple(inout int result, bool x) {
      action a() {
        result = 5;
      }
      action b() {
        result = 7;
      }
      table t {
        key = {
          x: exact;
        }
        actions = {
          a;
          b;
        }
      }
      apply {
      }
    };
    """

  let program = try! #UseOkResult(SpecialCompilers.ProgramCompiler.Compile(simple_parser_declaration))

  // Pull the control out of the compiled program.
  let controls = program.TypesWithTypes { (tipe: P4Type) -> Bool in
    return switch tipe {
    case let c as Control: c.name == "simple"
    default: false
    }
  }
  var control = ((controls[0]) as! Control)

  // Add entries to the table.
  control = control.updateTable(
    addEntry: (
      P4Value(P4BooleanValue(withValue: true)),
      TypedIdentifier(name: "a", withType: P4QualifiedType(Action()))
    )
  ).updateTable(
    addEntry: (
      P4Value(P4BooleanValue(withValue: false)),
      TypedIdentifier(name: "b", withType: P4QualifiedType(Action()))
    ))

  // Set a variable in the global scope for the inout first parameter.
  var global_values = VarValueScopes().enter()
  global_values = global_values.declare(
    identifier: Identifier(name: "result_arg"),
    withValue: P4Value(
      P4IntValue(withValue: 0),
      P4QualifiedType(P4Int())))

  let runtime = try #UseOkResult(
    P4Runtime.Runtime<P4TableHitMissValue, Control>.create(
      control: control, withGlobalValues: global_values))

  let (hit_miss, updated_execution) = try #UseOkResult(
    runtime.run(
      withArguments: ArgumentList([
        Argument(
          TypedIdentifier(name: "result_arg", withType: P4QualifiedType(P4Int())), atIndex: 0),
        Argument(P4Value(P4BooleanValue(withValue: true)), atIndex: 1),
      ])))

  // We expect there to be a hit.
  #expect(hit_miss == P4TableHitMissValue.Hit)

  // And that the proper action was invoked.
  let result_arg = try #UseOkResult(
    updated_execution.scopes.lookup(identifier: Identifier(name: "result_arg")))
  #expect(result_arg.eq(P4Value(P4IntValue(withValue: 5))))
}

@Test func test_control_single_key_false() async throws {
  let simple_parser_declaration = """
    control simple(inout int result, bool x) {
      action a() {
        result = 5;
      }
      action b() {
        result = 7;
      }
      table t {
        key = {
          x: exact;
        }
        actions = {
          a;
          b;
        }
      }
      apply {
      }
    };
    """

  let program = try! #UseOkResult(SpecialCompilers.ProgramCompiler.Compile(simple_parser_declaration))

  // Pull the control out of the compiled program.
  let controls = program.TypesWithTypes { (tipe: P4Type) -> Bool in
    switch tipe {
    case let c as Control: c.name == "simple"
    default: false
    }
  }
  var control = ((controls[0]) as! Control)

  // Add entries to the table.
  control = control.updateTable(
    addEntry: (
      P4Value(P4BooleanValue(withValue: true)),
      TypedIdentifier(name: "a", withType: P4QualifiedType(Action()))
    )
  ).updateTable(
    addEntry: (
      P4Value(P4BooleanValue(withValue: false)),
      TypedIdentifier(name: "b", withType: P4QualifiedType(Action()))
    ))

  // Set a variable in the global scope for the inout first parameter.
  var global_values = VarValueScopes().enter()
  global_values = global_values.declare(
    identifier: Identifier(name: "result_arg"),
    withValue: P4Value(
      P4IntValue(withValue: 0),
      P4QualifiedType(P4Int())))

  let runtime = try #UseOkResult(
    P4Runtime.Runtime<P4TableHitMissValue, Control>.create(
      control: control, withGlobalValues: global_values))

  let (hit_miss, updated_execution) = try #UseOkResult(
    runtime.run(
      withArguments: ArgumentList([
        Argument(
          TypedIdentifier(name: "result_arg", withType: P4QualifiedType(P4Int())), atIndex: 0),
        Argument(P4Value(P4BooleanValue(withValue: false)), atIndex: 1),
      ])))

  // We expect there to be a hit.
  #expect(hit_miss == P4TableHitMissValue.Hit)

  // And that the proper action was invoked.
  let result_arg = try #UseOkResult(
    updated_execution.scopes.lookup(identifier: Identifier(name: "result_arg")))
  #expect(result_arg.eq(P4Value(P4IntValue(withValue: 7))))
}

@Test func test_control_single_integer_key_hit() async throws {
  let simple_parser_declaration = """
    control simple(inout int result, int x) {
      action a() {
        result = 5;
      }
      action b() {
        result = 7;
      }
      table t {
        key = {
          x: exact;
        }
        actions = {
          a;
          b;
        }
      }
      apply {
      }
    };
    """

  let program = try! #UseOkResult(SpecialCompilers.ProgramCompiler.Compile(simple_parser_declaration))

  // Pull the control out of the compiled program.
  let controls = program.TypesWithTypes { (tipe: P4Type) -> Bool in
    switch tipe {
    case let c as Control: c.name == "simple"
    default: false
    }
  }
  var control = ((controls[0]) as! Control)

  // Add entries to the table.
  control = control.updateTable(
    addEntry: (
      P4Value(P4IntValue(withValue: 5)),
      TypedIdentifier(name: "a", withType: P4QualifiedType(Action()))
    )
  ).updateTable(
    addEntry: (
      P4Value(P4IntValue(withValue: 2)),
      TypedIdentifier(name: "b", withType: P4QualifiedType(Action()))
    ))

  // Set a variable in the global scope for the inout first parameter.
  var global_values = VarValueScopes().enter()
  global_values = global_values.declare(
    identifier: Identifier(name: "result_arg"),
    withValue: P4Value(
      P4IntValue(withValue: 0),
      P4QualifiedType(P4Int())))

  let runtime = try #UseOkResult(
    P4Runtime.Runtime<P4TableHitMissValue, Control>.create(
      control: control, withGlobalValues: global_values))

  let (hit_miss, updated_execution) = try #UseOkResult(
    runtime.run(
      withArguments: ArgumentList([
        Argument(
          TypedIdentifier(name: "result_arg", withType: P4QualifiedType(P4Int())), atIndex: 0),
        Argument(P4Value(P4IntValue(withValue: 5)), atIndex: 1),
      ])))

  // We expect there to be a hit.
  #expect(hit_miss == P4TableHitMissValue.Hit)

  let result_arg = try #UseOkResult(
    updated_execution.scopes.lookup(identifier: Identifier(name: "result_arg")))
  #expect(result_arg.eq(P4Value(P4IntValue(withValue: 5))))
}

@Test func test_control_single_integer_key_miss() async throws {
  let simple_parser_declaration = """
    control simple(inout int result, int x) {
      action a() {
        result = 5;
      }
      action b() {
        result = 7;
      }
      table t {
        key = {
          x: exact;
        }
        actions = {
          a;
          b;
        }
      }
      apply {
      }
    };
    """

  let program = try! #UseOkResult(SpecialCompilers.ProgramCompiler.Compile(simple_parser_declaration))

  // Pull the control out of the compiled program.
  let controls = program.TypesWithTypes { (tipe: P4Type) -> Bool in
    return switch tipe {
    case let c as Control: c.name == "simple"
    default: false
    }
  }
  var control = ((controls[0]) as! Control)

  // Add entries to the table.
  control = control.updateTable(
    addEntry: (
      P4Value(P4IntValue(withValue: 1)),
      TypedIdentifier(name: "a", withType: P4QualifiedType(Action()))
    )
  ).updateTable(
    addEntry: (
      P4Value(P4IntValue(withValue: 2)),
      TypedIdentifier(name: "b", withType: P4QualifiedType(Action()))
    ))

  // Set a variable in the global scope for the inout first parameter.
  var global_values = VarValueScopes().enter()
  global_values = global_values.declare(
    identifier: Identifier(name: "result_arg"),
    withValue: P4Value(
      P4IntValue(withValue: 0),
      P4QualifiedType(P4Int())))

  let runtime = try #UseOkResult(
    P4Runtime.Runtime<P4TableHitMissValue, Control>.create(
      control: control, withGlobalValues: global_values))

  let (hit_miss, updated_execution) = try #UseOkResult(
    runtime.run(
      withArguments: ArgumentList([
        Argument(
          TypedIdentifier(name: "result_arg", withType: P4QualifiedType(P4Int())), atIndex: 0),
        Argument(P4Value(P4IntValue(withValue: 3)), atIndex: 1),
      ])))

  // We expect there to be a hit.
  #expect(hit_miss == P4TableHitMissValue.Miss)

  let result_arg = try #UseOkResult(
    updated_execution.scopes.lookup(identifier: Identifier(name: "result_arg")))
  #expect(result_arg.eq(P4Value(P4IntValue(withValue: 0))))
}

@Test func test_control_multiple_keys() async throws {
  let simple_parser_declaration = """
    control simple(inout int result, bool x, int f) {
      action a() {
        result = 5;
      }
      action b() {
        result = 7;
      }
      table t {
        key = {
          x: exact;
          f: exact;
        }
        actions = {
          a;
          b;
        }
      }
      apply {
      }
    };
    """

  let program = try! #UseOkResult(SpecialCompilers.ProgramCompiler.Compile(simple_parser_declaration))

  // Pull the control out of the compiled program.
  let controls = program.TypesWithTypes { (tipe: P4Type) -> Bool in
    return switch tipe {
    case let c as Control: c.name == "simple"
    default: false
    }
  }
  var control = ((controls[0]) as! Control)

  // Add entries to the table.
  control = control.updateTable(
    addEntry: (
      P4Value(P4BooleanValue(withValue: true)),
      TypedIdentifier(name: "a", withType: P4QualifiedType(Action()))
    )
  ).updateTable(
    addEntry: (
      P4Value(P4IntValue(withValue: 5)),
      TypedIdentifier(name: "b", withType: P4QualifiedType(Action()))
    ))

  // Set a variable in the global scope for the inout first parameter.
  var global_values = VarValueScopes().enter()
  global_values = global_values.declare(
    identifier: Identifier(name: "result_arg"),
    withValue: P4Value(
      P4IntValue(withValue: 0),
      P4QualifiedType(P4Int())))

  let runtime = try #UseOkResult(
    P4Runtime.Runtime<P4TableHitMissValue, Control>.create(
      control: control, withGlobalValues: global_values))

  let (hit_miss, updated_execution) = try #UseOkResult(
    runtime.run(
      withArguments: ArgumentList([
        Argument(
          TypedIdentifier(name: "result_arg", withType: P4QualifiedType(P4Int())), atIndex: 0),
        Argument(P4Value(P4BooleanValue(withValue: false)), atIndex: 1),  // false will make the x key miss.
        Argument(P4Value(P4IntValue(withValue: 5)), atIndex: 2),
      ])))

  // We expect there to be a hit -- but from the second key!
  #expect(hit_miss == P4TableHitMissValue.Hit)

  // And that the proper action was invoked.
  let result_arg = try #UseOkResult(
    updated_execution.scopes.lookup(identifier: Identifier(name: "result_arg")))
  #expect(result_arg.eq(P4Value(P4IntValue(withValue: 7))))
}

@Test func test_control_key_from_struct() async throws {
  let simple_parser_declaration = """
    struct K {
      int i;
      int j;
    };
    control simple(inout int result, K k) {
      action a() {
        result = 5;
      }
      action b() {
        result = 7;
      }
      table t {
        key = {
          k.i: exact;
        }
        actions = {
          a;
          b;
        }
      }
      apply {
      }
    };
    """

  let program = try! #UseOkResult(SpecialCompilers.ProgramCompiler.Compile(simple_parser_declaration))

  // Pull the control out of the compiled program.
  let controls = program.TypesWithTypes { (tipe: P4Type) -> Bool in
    return switch tipe {
    case let c as Control: c.name == "simple"
    default: false
    }
  }
  var control = ((controls[0]) as! Control)

  let k_fields = P4StructFields([
    P4StructFieldIdentifier(name: "i", withType: P4QualifiedType(P4Int())),
    P4StructFieldIdentifier(name: "j", withType: P4QualifiedType(P4Int())),
  ])
  let k_type = P4Struct(withName: Identifier(name: "K"), andFields: k_fields)

  let k_instance = P4StructValue(
    withType: k_type,
    andInitializers: [
      P4Value(P4IntValue(withValue: 5)),
      P4Value(P4IntValue(withValue: 1)),
    ])

  // Add entries to the table.
  control = control.updateTable(
    addEntry: (
      P4Value(P4IntValue(withValue: 5)),
      TypedIdentifier(name: "a", withType: P4QualifiedType(Action()))
    )
  ).updateTable(
    addEntry: (
      P4Value(P4IntValue(withValue: 4)),
      TypedIdentifier(name: "b", withType: P4QualifiedType(Action()))
    ))

  // Set a variable in the global scope for the inout first parameter.
  var global_values = VarValueScopes().enter()
  global_values = global_values.declare(
    identifier: Identifier(name: "result_arg"),
    withValue: P4Value(
      P4IntValue(withValue: 0),
      P4QualifiedType(P4Int())))

  let runtime = try #UseOkResult(
    P4Runtime.Runtime<P4TableHitMissValue, Control>.create(
      control: control, withGlobalValues: global_values))

  let (hit_miss, updated_execution) = try #UseOkResult(
    runtime.run(
      withArguments: ArgumentList([
        Argument(
          TypedIdentifier(name: "result_arg", withType: P4QualifiedType(P4Int())), atIndex: 0),
        Argument(P4Value(k_instance), atIndex: 1),
      ])))

  // We expect there to be a hit.
  #expect(hit_miss == P4TableHitMissValue.Hit)

  // And that the proper action was invoked.
  let result_arg = try #UseOkResult(
    updated_execution.scopes.lookup(identifier: Identifier(name: "result_arg")))
  #expect(result_arg.eq(P4Value(P4IntValue(withValue: 5))))

  // Now, check whether the b action can be invoked.
  let k_instance2 = P4StructValue(
    withType: k_type,
    andInitializers: [
      P4Value(P4IntValue(withValue: 4)),
      P4Value(P4IntValue(withValue: 1)),
    ])

  // Set a variable in the global scope for the inout first parameter.
  var next_global_values = VarValueScopes().enter()
  next_global_values = global_values.declare(
    identifier: Identifier(name: "result_arg"),
    withValue: P4Value(
      P4IntValue(withValue: 0),
      P4QualifiedType(P4Int())))

  let runtime2 = try #UseOkResult(
    P4Runtime.Runtime<P4TableHitMissValue, Control>.create(
      control: control, withGlobalValues: next_global_values))

  let (hit_miss2, updated_execution2) = try #UseOkResult(
    runtime2.run(
      withArguments: ArgumentList([
        Argument(
          TypedIdentifier(name: "result_arg", withType: P4QualifiedType(P4Int())), atIndex: 0),
        Argument(P4Value(k_instance2), atIndex: 1),
      ])))

  // We expect there to be a hit.
  #expect(hit_miss2 == P4TableHitMissValue.Hit)

  // And that the proper action was invoked.
  let result_arg2 = try #UseOkResult(
    updated_execution2.scopes.lookup(identifier: Identifier(name: "result_arg")))
  #expect(result_arg2.eq(P4Value(P4IntValue(withValue: 7))))
}

@Test func test_control_key_from_struct_miss() async throws {
  let simple_parser_declaration = """
    struct K {
      int i;
      int j;
    };
    control simple(inout int result, K k) {
      action a() {
        result = 5;
      }
      action b() {
        result = 7;
      }
      table t {
        key = {
          k.i: exact;
        }
        actions = {
          a;
          b;
        }
      }
      apply {
      }
    };
    """

  let program = try! #UseOkResult(SpecialCompilers.ProgramCompiler.Compile(simple_parser_declaration))

  // Pull the control out of the compiled program.
  let controls = program.TypesWithTypes { (tipe: P4Type) -> Bool in
    return switch tipe {
    case let c as Control: c.name == "simple"
    default: false
    }
  }
  var control = ((controls[0]) as! Control)

  let k_fields = P4StructFields([
    P4StructFieldIdentifier(name: "i", withType: P4QualifiedType(P4Int())),
    P4StructFieldIdentifier(name: "j", withType: P4QualifiedType(P4Int())),
  ])
  let k_type = P4Struct(withName: Identifier(name: "K"), andFields: k_fields)

  let k_instance = P4StructValue(
    withType: k_type,
    andInitializers: [
      P4Value(P4IntValue(withValue: 8)),
      P4Value(P4IntValue(withValue: 1)),
    ])

  // Add entries to the table.
  control = control.updateTable(
    addEntry: (
      P4Value(P4IntValue(withValue: 5)),
      TypedIdentifier(name: "a", withType: P4QualifiedType(Action()))
    )
  ).updateTable(
    addEntry: (
      P4Value(P4IntValue(withValue: 4)),
      TypedIdentifier(name: "b", withType: P4QualifiedType(Action()))
    ))

  // Set a variable in the global scope for the inout first parameter.
  var global_values = VarValueScopes().enter()
  global_values = global_values.declare(
    identifier: Identifier(name: "result_arg"),
    withValue: P4Value(
      P4IntValue(withValue: 0),
      P4QualifiedType(P4Int())))

  let runtime = try #UseOkResult(
    P4Runtime.Runtime<P4TableHitMissValue, Control>.create(
      control: control, withGlobalValues: global_values))

  let (hit_miss, updated_execution) = try #UseOkResult(
    runtime.run(
      withArguments: ArgumentList([
        Argument(
          TypedIdentifier(name: "result_arg", withType: P4QualifiedType(P4Int())), atIndex: 0),
        Argument(P4Value(k_instance), atIndex: 1),
      ])))

  // We expect there to be a hit.
  #expect(hit_miss == P4TableHitMissValue.Miss)

  // And that the proper action was invoked.
  let result_arg = try #UseOkResult(
    updated_execution.scopes.lookup(identifier: Identifier(name: "result_arg")))
  #expect(result_arg.eq(P4Value(P4IntValue(withValue: 0))))
}