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

// Witness the parts of a P4 Program that are visitable.
extension Program: Visitable {}
extension Parser: Visitable {}
extension ParserState: Visitable {}
extension VariableDeclarationStatement: Visitable {}
extension ConditionalStatement: Visitable {}
extension BlockStatement: Visitable {}
extension ReturnStatement: Visitable {}
extension ApplyStatement: Visitable {}
extension KeysetExpression: Visitable {}
extension SelectCaseExpression: Visitable {}
extension SelectExpression: Visitable {}
extension ArrayAccessExpression: Visitable {}
extension FieldAccessExpression: Visitable {}
extension FunctionCall: Visitable {}
extension BinaryOperatorExpression: Visitable {}
extension Declaration: Visitable {}
extension ExternDeclaration: Visitable {}
extension FunctionDeclaration: Visitable {}
extension Action: Visitable {}
extension TableKeyEntry: Visitable {}
extension TablePropertyList: Visitable {}
extension Table: Visitable {}
extension Control: Visitable {}

/// Context for the visiting process.
public struct VisitorContext<UserContext> {
  let visitor: VisitorDriver<UserContext>
  let uc: UserContext
  public init(_ v: VisitorDriver<UserContext>, _ uc: UserContext) {
    self.visitor = v
    self.uc = uc
  }

  public func getUserContext() -> UserContext {
    return self.uc
  }

  public func getVisitorDriver() -> VisitorDriver<UserContext> {
    return self.visitor
  }

  public func next(uc: UserContext) -> VisitorContext<UserContext> {
    return VisitorContext<UserContext>(self.visitor, uc)
  }
}

/// A driver for visiting the components of a parsed P4 program.
public struct VisitorDriver<UserContext> {
  let visitor: any LanguageVisitor<UserContext>

  public init(_ visitor: any LanguageVisitor<UserContext>) {
    self.visitor = visitor
  }

  public func generateContext(uc: UserContext) -> VisitorContext<UserContext> {
    return VisitorContext(self, uc)
  }

  /// Visit a `P4Type`.
  public func visit(
    _ t: P4Type, context: VisitorContext<UserContext>
  ) -> Result<VisitorContext<UserContext>> {
    return switch t {
    case let vv as Control: visitor.visit(vv, context)
    case let vv as Parser: visitor.visit(vv, context)
    default: .Error(Error(withMessage: "Could not visit type \(t)"))
    }
  }

  /// Visit a part of a P4 program.
  public func visit(
    _ v: Visitable, context: VisitorContext<UserContext>
  ) -> Result<VisitorContext<UserContext>> {
    return switch v {
    case let vv as Program: visitor.visit(vv, context)
    case let vv as Parser: visitor.visit(vv, context)
    case let vv as InstantiatedParserState: visitor.visit(vv, context)
    case let vv as VariableDeclarationStatement: visitor.visit(vv, context)
    case let vv as ConditionalStatement: visitor.visit(vv, context)
    case let vv as BlockStatement: visitor.visit(vv, context)
    case let vv as ReturnStatement: visitor.visit(vv, context)
    case let vv as ApplyStatement: visitor.visit(vv, context)
    case let vv as KeysetExpression: visitor.visit(vv, context)
    case let vv as SelectCaseExpression: visitor.visit(vv, context)
    case let vv as SelectExpression: visitor.visit(vv, context)
    case let vv as ArrayAccessExpression: visitor.visit(vv, context)
    case let vv as FieldAccessExpression: visitor.visit(vv, context)
    case let vv as FunctionCall: visitor.visit(vv, context)
    case let vv as BinaryOperatorExpression: visitor.visit(vv, context)
    case let vv as Declaration: visitor.visit(vv, context)
    case let vv as ExternDeclaration: visitor.visit(vv, context)
    case let vv as FunctionDeclaration: visitor.visit(vv, context)
    case let vv as Action: visitor.visit(vv, context)
    case let vv as TableKeyEntry: visitor.visit(vv, context)
    case let vv as TablePropertyList: visitor.visit(vv, context)
    case let vv as Table: visitor.visit(vv, context)
    case let vv as Control: visitor.visit(vv, context)
    default: .Error(Error(withMessage: "Could not visit \(v)"))
    }
  }

  /// Start the process of visiting a P4 program.
  public func start(_ v: Visitable, context: UserContext) -> Result<UserContext> {
    let visit_result = self.visit(v, context: VisitorContext(self, context))
    return switch visit_result {
    case .Ok(let vc): .Ok(vc.getUserContext())
    case .Error(let e): .Error(e)
    }
  }
}