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

public struct Generated {
  let gen: String
  public init(_ base: String = "") {
    self.gen = base
  }

  public func getGeneratedCode() -> String {
    return self.gen
  }

  public func append(_ a: String) -> Generated {
    return Generated(self.gen + a)
  }
}

/// Generate code for a P4 program.
///
/// See the [P4 behavioral model](https://github.com/p4lang/behavioral-model) for
/// the format.
public struct CodeGenerator: LanguageVisitor {

  public init() {}

  /// Generate code.
  public func codeGen(_ node: any Visitable) -> Result<Generated> {
    let visitor = VisitorDriver<Generated>(self)
    let generated = Generated()
    return visitor.start(node, context: generated)
  }

  public typealias Context = Generated

  public func visit(
    _ v: Program, _ c: VisitorContext<Generated>
  ) -> Result<VisitorContext<Generated>> {

    var result: Result<VisitorContext<Generated>> = Fold(
      input: v.types, initial: .Ok(c.next(uc: c.getUserContext().append("[")))
    ) { (current, acc) in
      return switch acc {
      case .Ok(let acc): acc.getVisitorDriver().visit(current, context: acc)
      case .Error(let e): .Error(e)
      }
    }

    result = Fold(
      input: v.externs, initial: result
    ) { (current, acc) in
      return switch acc {
      case .Ok(let acc): acc.getVisitorDriver().visit(current, context: acc)
      case .Error(let e): .Error(e)
      }
    }

    /// TODO: Handle instances.
    /*
    result = Fold(
      input: v.instances, initial: result
    ) { (current, acc) in
      return switch acc {
      case .Ok(let acc): acc.getVisitorDriver().visit(current, context: acc)
      case .Error(let e): .Error(e)
      }
    }*/

    result = result.map {
      .Ok($0.next(uc: $0.getUserContext().append("]")))
    }

    return result
  }

  public func visit(
    _ v: Parser, _ c: VisitorContext<Generated>
  ) -> Result<VisitorContext<Generated>> {
    var initial = "{"
    initial += "name: \"\(v.name)\","
    initial += "init_state: \"start\","
    initial += "parse_states: ["

    let result: Result<VisitorContext<Generated>> = Fold(
      input: v.states.states, initial: .Ok(c.next(uc: c.getUserContext().append(initial)))
    ) { (current, acc) in
      return switch acc {
      case .Ok(let acc):
        switch self.visit(current, acc) {
        case .Ok(let gend): .Ok(gend.next(uc: gend.getUserContext().append(",")))
        case .Error(let e): .Error(e)
        }
      case .Error(let e): .Error(e)
      }
    }

    return result.map {
      .Ok($0.next(uc: $0.getUserContext().append("]}")))
    }
  }

  public func visit(
    _ v: ParserState, _ c: VisitorContext<Generated>
  ) -> Result<VisitorContext<Generated>> {
    let direct_transition_codegen = {
      (
        _: ParserStateDirectTransition, c: VisitorContext<Generated>
      ) -> Result<VisitorContext<Generated>> in
      return .Ok(c.next(uc: c.getUserContext().append("[]")))
    }

    let no_transition_codegen = {
      (
        _: ParserStateNoTransition, c: VisitorContext<Generated>
      ) -> Result<VisitorContext<Generated>> in
      return .Ok(c.next(uc: c.getUserContext().append("[]")))
    }

    let select_transition_codegen = {
      (
        state: ParserStateSelectTransition, c: VisitorContext<Generated>
      ) -> Result<VisitorContext<Generated>> in
      return switch self.visit(state.te, c.next(uc: c.getUserContext().append("["))) {
      case .Ok(let res): .Ok(res.next(uc: res.getUserContext().append("]")))
      case .Error(let e): .Error(e)
      }
    }

    var initial = "{"
    initial += "name: \"\(v.getName())\","
    initial += "transitions: "

    let result: Result<VisitorContext<Generated>> =
      switch v {
      case let s as ParserStateSelectTransition:
        select_transition_codegen(s, c.next(uc: c.getUserContext().append(initial)))
      case let s as ParserStateDirectTransition:
        direct_transition_codegen(s, c.next(uc: c.getUserContext().append(initial)))
      case let s as ParserStateNoTransition:
        no_transition_codegen(s, c.next(uc: c.getUserContext().append(initial)))
      default:
        .Error(Error(withMessage: "Could not code gen \(self)"))
      }

    return result.map {
      .Ok($0.next(uc: $0.getUserContext().append("}")))
    }
  }

  public func visit(
    _ v: VariableDeclarationStatement, _ c: VisitorContext<Generated>
  ) -> Result<VisitorContext<Generated>> {

    return .Ok(c)
  }

  public func visit(
    _ v: ConditionalStatement, _ c: VisitorContext<Generated>
  ) -> Result<VisitorContext<Generated>> {
    return .Ok(c)
  }

  public func visit(
    _ v: BlockStatement, _ c: VisitorContext<Generated>
  ) -> Result<VisitorContext<Generated>> {
    return .Ok(c)
  }

  public func visit(
    _ v: ReturnStatement, _ c: VisitorContext<Generated>
  ) -> Result<VisitorContext<Generated>> {
    return .Ok(c)
  }

  public func visit(
    _ v: ApplyStatement, _ c: VisitorContext<Generated>
  ) -> Result<VisitorContext<Generated>> {
    return .Ok(c)
  }

  public func visit(
    _ v: KeysetExpression, _ c: VisitorContext<Generated>
  ) -> Result<VisitorContext<Generated>> {
    return .Ok(c)
  }

  public func visit(
    _ v: SelectCaseExpression, _ c: VisitorContext<Generated>
  ) -> Result<VisitorContext<Generated>> {
    return .Ok(c)
  }

  public func visit(
    _ v: SelectExpression, _ c: VisitorContext<Generated>
  ) -> Result<VisitorContext<Generated>> {
    return .Ok(c.next(uc: c.getUserContext().append("Select")))
  }

  public func visit(
    _ v: ArrayAccessExpression, _ c: VisitorContext<Generated>
  ) -> Result<VisitorContext<Generated>> {
    return .Ok(c)
  }

  public func visit(
    _ v: FieldAccessExpression, _ c: VisitorContext<Generated>
  ) -> Result<VisitorContext<Generated>> {
    return .Ok(c)
  }

  public func visit(
    _ v: FunctionCall, _ c: VisitorContext<Generated>
  ) -> Result<VisitorContext<Generated>> {
    return .Ok(c)
  }

  public func visit(
    _ v: BinaryOperatorExpression, _ c: VisitorContext<Generated>
  ) -> Result<VisitorContext<Generated>> {
    return .Ok(c)
  }

  public func visit(
    _ v: Declaration, _ c: VisitorContext<Generated>
  ) -> Result<VisitorContext<Generated>> {
    return .Ok(c)
  }

  public func visit(
    _ v: ExternDeclaration, _ c: VisitorContext<Generated>
  ) -> Result<VisitorContext<Generated>> {
    return .Ok(c)
  }

  public func visit(
    _ v: FunctionDeclaration, _ c: VisitorContext<Generated>
  ) -> Result<VisitorContext<Generated>> {
    return .Ok(c)
  }

  public func visit(
    _ v: Action, _ c: VisitorContext<Generated>
  ) -> Result<VisitorContext<Generated>> {
    return .Ok(c)
  }

  public func visit(
    _ v: TableKeyEntry, _ c: VisitorContext<Generated>
  ) -> Result<VisitorContext<Generated>> {
    return .Ok(c)
  }

  public func visit(
    _ v: TablePropertyList, _ c: VisitorContext<Generated>
  ) -> Result<VisitorContext<Generated>> {
    return .Ok(c)
  }

  public func visit(_ v: Table, _ c: VisitorContext<Generated>) -> Result<VisitorContext<Generated>>
  {
    return .Ok(c)
  }

  public func visit(
    _ v: Control, _ c: VisitorContext<Generated>
  ) -> Result<VisitorContext<Generated>> {
    return .Ok(c)
  }
  public func visit(
    _ parser_state: InstantiatedParserState, _ c: VisitorContext<Generated>
  ) -> Common.Result<VisitorContext<Generated>> {
    return .Ok(c)
  }
}