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

extension P4Value: AST.AnExpression {}

public struct AST {

  public protocol AnExpression {}
  public protocol AnStatement {}
  public protocol AnState {}

  public struct Statement {}
  struct Expression {}

  public struct VariableDeclarationStatement: AnStatement {
    public var initializer: AnExpression?
    public var identifier: AST.Identifier
    public var tipe: AST.Tipe

    public init(
      identifier: Identifier, withType tipe: AST.Tipe, withInitializer initializer: AnExpression?
    ) {
      self.identifier = identifier
      self.initializer = initializer
      self.tipe = tipe
    }
  }

  public struct ConditionalStatement: AnStatement {
    public var condition: AnExpression
    public var thenn: AnStatement
    public var elss: AnStatement?

    public init(condition: AnExpression, withThen thenn: AnStatement) {
      self.condition = condition
      self.thenn = thenn
      self.elss = .none
    }

    public init(
      condition: AnExpression, withThen thenn: AnStatement,
      andElse elss: AnStatement
    ) {
      self.condition = condition
      self.thenn = thenn
      self.elss = elss
    }
  }

  public struct BlockStatement: AnStatement {
    public var statements: [AnStatement]

    public init(_ statements: [AnStatement]) {
      self.statements = statements
    }

  }

  public struct ReturnStatement: AnStatement {
    public let value: AnExpression

    public init(_ value: AnExpression) {
      self.value = value
    }
  }

  public struct ApplyStatement: AnStatement {
    public let body: AST.BlockStatement?

    public init() { self.body = .none }
    public init(_ body: AST.BlockStatement) {
      self.body = body
    }
  }

  public struct Action: CustomStringConvertible {
    public var description: String {
      return "Action: "
        + "\(self.name) with parameters \(self.params) and body \(String(describing: self.body))"
    }

    public var body: AST.BlockStatement?
    public var params: AST.ParameterList
    public var name: AST.Identifier

    public init(
      named name: AST.Identifier = AST.Identifier(Common.Identifier(name: "")),
      withParameters parameters: ParameterList = ParameterList([]),
      withBody body: AST.BlockStatement? = .none
    ) {
      self.name = name
      self.params = parameters
      self.body = body
    }

  }

  public struct Actions: CustomStringConvertible {
    public let actions: [AST.Action]
    public init(withActions actions: [AST.Action]) {
      self.actions = actions
    }

    public var description: String {
      return "Actions: "
        + actions.map { action in
          return "\(action)"
        }.joined(separator: ";")
    }
  }

  public enum TableKeyMatchType {
    case Exact
  }

  public struct TableKeyEntry: CustomStringConvertible {
    public let key: AST.KeysetExpression
    public let match_type: AST.TableKeyMatchType

    public init(_ key: AST.KeysetExpression, _ match: AST.TableKeyMatchType) {
      self.key = key
      self.match_type = match
    }

    public var description: String {
      return "Table Key Entry: " + "\(self.key): \(self.match_type)"
    }
  }

  public struct TableKeys: CustomStringConvertible {
    public let keys: [AST.TableKeyEntry]

    public init(withEntries entries: [AST.TableKeyEntry]) {
      self.keys = entries
    }
    public init() {
      self.keys = []
    }

    public var description: String {
      return "Table Keys: "
        + self.keys.map { key in
          return "\(key)"
        }.joined(separator: ";")
    }
  }

  public struct TableActionsProperty: CustomStringConvertible {
    public let actions: [AST.Identifier]
    public init(_ actions: [AST.Identifier] = []) {
      self.actions = actions
    }

    public var description: String {
      return "Table Actions: "
        + self.actions.map { action in
          return "\(action)"
        }.joined(separator: ";")
    }
  }

  public struct TablePropertyList: CustomStringConvertible {
    public let actions: AST.TableActionsProperty
    public let keys: AST.TableKeys
    public init(withActions actions: AST.TableActionsProperty, withKeys keys: AST.TableKeys) {
      self.actions = actions
      self.keys = keys
    }

    public var description: String {
      return "Table Property List: \(self.actions) \(self.keys)"
    }
  }

  public struct Table: CustomStringConvertible {
    public let properties: AST.TablePropertyList
    let name: AST.Identifier
    public let entries: [(P4Value, AST.Identifier)]

    public init(
      withName name: AST.Identifier, withPropertyList property_list: AST.TablePropertyList,
      withEntries entries: [(P4Value, AST.Identifier)] = []
    ) {
      self.name = name
      self.properties = property_list
      self.entries = entries
    }

    public var description: String {
      return "Table named: \(self.name) \(self.properties)"
    }

    /// When the control is evaluated, the value of the x in the table is
    /// compared to the entries and the match is assocated with an action
    /// that is invoked when the match occurs!

    public func update(addEntry entry: (P4Value, AST.Identifier)) -> Table {
      return Table(
        withName: self.name, withPropertyList: self.properties, withEntries: self.entries + [entry])
    }
  }

  public protocol AnDeclaration: AnStatement {}

  public struct Declaration {}

  public struct Control: CustomStringConvertible, AnDeclaration {
    public var description: String {
      return "Control named \(self._name) \(self.parameters) \(self.actions) \(self.table)"
    }

    public let actions: AST.Actions
    public let table: AST.Table
    let _parameters: AST.ParameterList
    let _name: AST.Identifier
    let apply: AST.ApplyStatement

    public var parameters: AST.ParameterList {
      _parameters
    }

    public var name: AST.Identifier {
      _name
    }

    public init(
      named: AST.Identifier, withParameters parameters: AST.ParameterList,
      withTable table: AST.Table,
      withActions actions: AST.Actions, withApply apply: AST.ApplyStatement
    ) {
      self._name = named
      self._parameters = parameters
      self.actions = actions
      self.table = table
      self.apply = apply
    }

    public func updateTable(addEntry entry: (P4Value, AST.Identifier)) -> Control {
      let table = self.table.update(addEntry: entry)

      return Control(
        named: self.name, withParameters: self.parameters, withTable: table,
        withActions: self.actions, withApply: self.apply)
    }

    public func def() -> P4DataValue? {
      return .none
    }

  }

  public struct ExternDeclaration: AnDeclaration {
    public let declaration: AST.AnDeclaration

    public init(_ declaration: AST.AnDeclaration) {
      self.declaration = declaration
    }
  }

  public struct FunctionDeclaration: AnDeclaration {
    public var description: String {
      return "Function named \(self.name) that returns \(self.tipe) with parameters \(self.params)"
    }

    public var body: AST.BlockStatement?
    public var params: AST.ParameterList
    public var name: AST.Identifier
    public var tipe: AST.Tipe

    public init(
      named name: AST.Identifier, ofType type: AST.Tipe,
      withParameters parameters: AST.ParameterList,
      withBody body: AST.BlockStatement?
    ) {
      self.name = name
      self.tipe = type
      self.params = parameters
      self.body = body
    }
  }

  public struct StructDeclaration: AnDeclaration {
    public let fields: [AST.VariableDeclarationStatement]
    public let identifier: AST.Identifier
    public init(_ id: AST.Identifier, _ fields: [AST.VariableDeclarationStatement]) {
      self.identifier = id
      self.fields = fields
    }
  }

  public struct Instantiation: AnStatement {
    public let name: AST.Identifier
    public var tipe: AST.Identifier
    public let arguments: AST.ArgumentList

    public init(
      named name: AST.Identifier, withType tipe: AST.Identifier,
      withArguments arguments: AST.ArgumentList
    ) {
      self.name = name
      self.arguments = arguments
      self.tipe = tipe
    }
  }

  public struct ExpressionStatement: AnStatement {
    public let expression: AnExpression

    public init(_ expr: AnExpression) {
      self.expression = expr
    }
  }

  public struct Identifier: AnExpression {
    public let id: Common.Identifier
    public init(_ id: Common.Identifier) {
      self.id = id
    }
  }

  public struct Literal: AnExpression {
    public let literal: P4Value
    public init(_ literal: P4Value) {
      self.literal = literal
    }
  }

  public enum KeysetExpression: AnExpression {
    case Default
    case Value(AnExpression)
  }

  public struct SelectCaseExpression: AnExpression {
    public let key: AST.KeysetExpression
    public let next_state_identifier: AST.Identifier

    public init(withKey key: AST.KeysetExpression, withNextState next_state_id: AST.Identifier) {
      self.key = key
      self.next_state_identifier = next_state_id
    }
  }

  public struct SelectExpression: AnExpression {
    public let selector: AnExpression
    public let case_expressions: [AST.SelectCaseExpression]

    public init(
      withSelector selector: AnExpression,
      withSelectCaseExpressions sces: [AST.SelectCaseExpression]
    ) {
      self.selector = selector
      self.case_expressions = sces
    }

    public func append_checked_sce(sce: AST.SelectCaseExpression) -> AST.SelectExpression {
      var new_cses = self.case_expressions
      new_cses.append(sce)
      return SelectExpression(
        withSelector: self.selector, withSelectCaseExpressions: new_cses)
    }
  }

  public enum BinaryOperatorExpressionType {
    case Add
    case Subtract
    case Multiply
    case Divide
    case Lt
    case Lte
    case Gt
    case Gte
    case Eq
    case And
    case Or
  }

  public struct BinaryOperatorExpression: AnExpression {
    public let left: AnExpression
    public let right: AnExpression
    public let type: BinaryOperatorExpressionType

    public init(
      withType tipe: BinaryOperatorExpressionType,
      withLhs lhs: AnExpression,
      withRhs rhs: AnExpression
    ) {
      self.type = tipe
      self.left = lhs
      self.right = rhs
    }
  }

  public struct ArrayAccessExpression: AnExpression {
    public let indexor: AnExpression
    public let name: AnExpression

    public init(
      withName name: AnExpression,
      withIndexor indexor: AnExpression
    ) {
      self.name = name
      self.indexor = indexor
    }
  }

  public struct FieldAccessExpression: AnExpression {
    public let field: Identifier
    public let strct: AnExpression

    public init(withStruct strct: AnExpression, withField field: Identifier) {
      self.strct = strct
      self.field = field
    }
  }

  public struct FunctionCall: AnExpression {
    public let callee: Identifier
    public let arguments: ArgumentList

    public init(_ callee: Identifier, withArguments arguments: ArgumentList) {
      self.callee = callee
      self.arguments = arguments
    }
  }

  public struct LocalElements {}

  public struct LocalElement {}

  public struct ParserAssignmentStatement: AnStatement {
    public let lvalue: AnExpression
    public let value: AnExpression

    public init(
      withLValue lvalue: AnExpression, withValue value: AnExpression
    ) {
      self.lvalue = lvalue
      self.value = value
    }
  }

  /// A P4 Parser State
  public class ParserState {
    let name: Identifier
    public let statements: [AnStatement]

    public var description: String {
      return "Parser State named \(self.name)"
    }

    public func getName() -> Identifier {
      return self.name
    }

    public func getStatements() -> [AnStatement] {
      return self.statements
    }

    /// Construct a ParserState
    public init(_ name: Identifier, _ statements: [AnStatement] = Array()) {
      self.name = name
      self.statements = statements
    }
  }

  /// TransitionStatement
  ///
  /// Only defined to define Compilable extension.
  public struct TransitionStatement {}

  public class ParserStateDirectTransition: ParserState, AnState {
    public let next_state_identifier: Identifier?

    public init(
      name: Identifier, withNextStateIdentifier next_state_id: Identifier,
      withStatements stmts: [AnStatement] = Array(),
    ) {
      self.next_state_identifier = next_state_id
      super.init(name, stmts)
    }
  }

  public class ParserStateNoTransition: ParserState, AnState {
    /// Construct a ParserState
    public init(
      name: Identifier, withStatements stmts: [AnStatement] = Array(),
    ) {
      super.init(name, stmts)
    }
  }

  public class ParserStateSelectTransition: ParserState, AnState {

    public let te: SelectExpression

    public init(
      name: Identifier, withTransitionExpression te: SelectExpression,
      withStatements stmts: [AnStatement] = Array()
    ) {
      self.te = te
      super.init(name, stmts)
    }
  }

  public struct ParserStates {
    public var states: [AnState] = Array()

    public func count() -> Int {
      return states.count
    }

    public init(_ states: [AnState] = Array()) {
      self.states = states
    }

    public func append(state: AnState) -> ParserStates {
      var new_states = self.states
      new_states.append(state)
      return ParserStates(new_states)
    }
  }

  /// A P4 Parser
  ///
  /// Note: A Parser is a type
  public struct Parser: AnDeclaration {
    public var states: ParserStates

    public var name: Identifier
    public var parameters: ParameterList

    public init(withName name: Identifier) {
      self.states = ParserStates()
      self.parameters = ParameterList()
      self.name = name
    }

    public init(withName name: Identifier, withParameters parameters: ParameterList) {
      self.states = ParserStates()
      self.parameters = parameters
      self.name = name
    }

    public var description: String {
      return "Parser \(self.name) with parameters: \(parameters) and states: \(self.states)"
    }
  }

  public struct Types {}

  public struct Tipe {
    public let tipe: P4QualifiedType
    public init(_ tipe: P4QualifiedType) {
      self.tipe = tipe
    }
  }

  public struct Parameter: CustomStringConvertible {
    public var name: AST.Identifier
    public var type: AST.Tipe

    public init(
      identifier: Identifier, withType type: AST.Tipe
    ) {
      self.name = identifier
      self.type = type
    }

    public var description: String {
      return "Parameter: \(self.name) with type \(self.type)"
    }

  }

  public struct ParameterList: CustomStringConvertible {
    public var parameters: [AST.Parameter]

    public init() {
      self.parameters = Array()
    }

    public init(_ parameters: [AST.Parameter]) {
      self.parameters = parameters
    }

    public func addParameter(_ parameter: AST.Parameter) -> AST.ParameterList {
      return AST.ParameterList(self.parameters + [parameter])
    }

    public var description: String {
      let parameters = self.parameters.map { parameter in
        parameter.description
      }.joined(separator: ";")
      return "Parameter list: \(parameters)"
    }
  }

  public struct ArgumentList {
    public let arguments: [AST.Argument]

    public init(_ arguments: [AST.Argument] = []) {
      self.arguments = arguments
    }

    public func addArgument(_ argument: AST.Argument) -> AST.ArgumentList {
      return ArgumentList(self.arguments + [argument])
    }
  }

  public struct Argument {
    public let index: Int
    public let argument: AST.AnExpression

    public init(_ argument: AST.AnExpression, atIndex index: Int) {
      self.argument = argument
      self.index = index
    }
  }

  public struct Program {
    public var statements: [AnStatement]
    public init(_ stmts: [AnStatement] = Array()) {
      self.statements = stmts
    }
  }
}

public struct ASTCompilerContext {
  public let lexical_context_name: AST.Identifier?
  public let lexical_context_statements: [AST.AnStatement]?
  public let extern_context: Bool

  public init(
    _ name: AST.Identifier? = .none, _ stmts: [AST.AnStatement]? = .none, _ extern: Bool = false
  ) {
    self.lexical_context_name = name
    self.lexical_context_statements = stmts
    self.extern_context = extern
  }

  public func update(withContextName cn: AST.Identifier?) -> ASTCompilerContext {
    return ASTCompilerContext(cn, self.lexical_context_statements, self.extern_context)
  }

  public func update(withContextStatements stmts: [AST.AnStatement]?) -> ASTCompilerContext {
    return ASTCompilerContext(self.lexical_context_name, stmts, self.extern_context)
  }

  public func update(withExtern extern: Bool) -> ASTCompilerContext {
    return ASTCompilerContext(self.lexical_context_name, self.lexical_context_statements, extern)
  }
}