// 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 SwiftCompilerPlugin
import SwiftSyntax
@_spi(ExperimentalLanguageFeature) import SwiftSyntaxMacros

public func remove_embedded_quotes(_ from: String) -> String {
  let result = from.replacing("\"", with: [])
  return result
}

struct MacroError: Error, CustomStringConvertible {
  var message: String
  var description: String {
    return message
  }
  public init(withMessage _message: String) {
    message = _message
  }
}

public struct UseOkResult: ExpressionMacro {
  public static func expansion(
    of node: some FreestandingMacroExpansionSyntax,
    in context: some MacroExpansionContext
  ) throws -> ExprSyntax {

    guard let argument = node.arguments.first?.expression else {
      throw Require.Error.SyntaxError
    }

    return """
      {
          switch \(argument) {
              case Result.Ok(let __good): return __good
              case Result.Error(let __error):
                  print("Unexpected result: \\(__error)")
                  throw Require.Error.UnexpectedResult
          }
      }()
      """
  }
}

public struct UseErrorResult: ExpressionMacro {
  public static func expansion(
    of node: some FreestandingMacroExpansionSyntax,
    in context: some MacroExpansionContext
  ) throws -> ExprSyntax {

    guard let argument = node.arguments.first?.expression else {
      throw Require.Error.SyntaxError
    }

    return """
      {
          switch \(argument) {
              case Result.Error(let __error): return __error
              case Result.Ok(let __good):
                  print("Unexpected result: \\(__good)")
                  throw Require.Error.UnexpectedResult
          }
      }()
      """
  }
}

public struct Require {
  public enum Error: Swift.Error {
    case UnexpectedResult
    case SyntaxError
  }
}

public struct RequireResult: ExpressionMacro {
  public static func expansion(
    of node: some FreestandingMacroExpansionSyntax,
    in context: some MacroExpansionContext
  ) throws -> ExprSyntax {

    guard let argument = node.arguments.first?.expression else {
      throw Require.Error.SyntaxError
    }

    return """
      {
          switch \(argument) {
              case Result.Ok(_): return true
              case Result.Error(let __error):
                  print("Unexpected result: \\(__error)")
                  return false
          }
      }()
      """
  }
}

public struct RequireErrorResult: ExpressionMacro {
  public static func expansion(
    of node: some FreestandingMacroExpansionSyntax,
    in context: some MacroExpansionContext
  ) throws -> ExprSyntax {

    let arguments = node.arguments.indices
    let expected_error = node.arguments[arguments.startIndex].expression
    let error_producer = node.arguments[arguments.index(after: arguments.startIndex)].expression

    return ExprSyntax(
      """
      {
          let __expected_error = \(expected_error)
          let __actual_error = \(error_producer)
          if case Result.Error(let __found_error) = __actual_error {
              if !__expected_error.eq(__found_error) {
                print("Expected Error: \\(__expected_error) but got Error: \\(__found_error)")
                return false
              }
              return true
          } else {
              print("Expected error, but got Ok")
              return false
          }
      }()
      """)
  }
}

public struct RequireNodeType: CodeItemMacro {
  public static func expansion(
    of node: some FreestandingMacroExpansionSyntax, in context: some MacroExpansionContext
  ) throws -> [CodeBlockItemSyntax] {
    let arguments = node.arguments.indices
    var arg_index = arguments.startIndex
    let node_to_check = node.arguments[arg_index].expression
    arg_index = arguments.index(after: arg_index)
    let expected_type = node.arguments[arg_index].expression
    arg_index = arguments.index(after: arg_index)
    let expected_type_nice_name = node.arguments[arg_index].expression

    let error_message =
      "Did not find " + remove_embedded_quotes(expected_type_nice_name.description)

    return [
      CodeBlockItemSyntax(
        """
        if \(node_to_check).nodeType != \(expected_type) {
          return Result.Error(
            ErrorWithLocation(sourceLocation: \(node_to_check).toSourceLocation(), withError: "\(raw: error_message)"))
        }
        """)
    ]
  }
}
public struct RequireNodesType: CodeItemMacro {
  public static func expansion(
    of node: some FreestandingMacroExpansionSyntax, in context: some MacroExpansionContext
  ) throws -> [CodeBlockItemSyntax] {
    let arguments = node.arguments.indices
    var arg_index = arguments.startIndex

    let node_to_check = node.arguments[arg_index].expression

    arg_index = arguments.index(after: arg_index)
    guard let expected_types = node.arguments[arg_index].expression.as(ArrayExprSyntax.self) else {
      throw MacroError(withMessage: "Node(s) to check must be in an array")
    }

    arg_index = arguments.index(after: arg_index)
    guard
      let expected_type_nice_names = node.arguments[arg_index].expression.as(ArrayExprSyntax.self)
    else {
      throw MacroError(withMessage: "Node nice names must be in an array")
    }

    let error_message =
      "Did not find one of the expected types: "
      + expected_type_nice_names.elements.map { l in
        remove_embedded_quotes("\(l.expression)")
      }.joined(separator: ",")

    let ifs = expected_types.elements.map { l in
      "\(node_to_check).nodeType != \(l.expression)"
    }.joined(separator: " && ")

    return [
      CodeBlockItemSyntax(
        """
        if \(raw: ifs) {
          return Result.Error(
            ErrorWithLocation(sourceLocation: \(node_to_check).toSourceLocation(), withError: "\(raw: error_message)"))
        }
        """)
    ]
  }
}
public struct SkipUnlessNodeType: CodeItemMacro {
  public static func expansion(
    of node: some FreestandingMacroExpansionSyntax, in context: some MacroExpansionContext
  ) throws -> [CodeBlockItemSyntax] {
    let arguments = node.arguments.indices
    var arg_index = arguments.startIndex
    let node_to_check = node.arguments[arg_index].expression
    arg_index = arguments.index(after: arg_index)
    let expected_type = node.arguments[arg_index].expression

    return [
      CodeBlockItemSyntax(
        """
        if \(node_to_check).nodeType != \(expected_type) {
          return Result.Ok(.none)
        }
        """)
    ]
  }
}

public struct SkipUnlessNodesTypes: CodeItemMacro {
  public static func expansion(
    of node: some FreestandingMacroExpansionSyntax, in context: some MacroExpansionContext
  ) throws -> [CodeBlockItemSyntax] {

    let arguments = node.arguments.indices
    var arg_index = arguments.startIndex

    let node_to_check = node.arguments[arg_index].expression

    arg_index = arguments.index(after: arg_index)
    guard let expected_types = node.arguments[arg_index].expression.as(ArrayExprSyntax.self) else {
      throw MacroError(withMessage: "Node(s) to check must be in an array")
    }

    arg_index = arguments.index(after: arg_index)

    let ifs = expected_types.elements.map { l in
      "\(node_to_check).nodeType != \(l.expression)"
    }.joined(separator: " && ")

    return [
      CodeBlockItemSyntax(
        """
        if \(raw: ifs) {
          return Result.Ok(.none)
        }
        """)
    ]
  }
}

public struct MustOr: CodeItemMacro {
  public static func expansion(
    of node: some FreestandingMacroExpansionSyntax, in context: some MacroExpansionContext
  ) throws -> [CodeBlockItemSyntax] {
    let arguments = node.arguments.indices

    var arg_index = arguments.startIndex
    let result = node.arguments[arg_index].expression

    arg_index = arguments.index(after: arg_index)
    let thing = node.arguments[arg_index].expression

    arg_index = arguments.index(after: arg_index)
    let or = node.arguments[arg_index].expression

    return [
      CodeBlockItemSyntax(
        """
        if let __thing = \(thing) {
          \(result) = __thing 
        } else {
          return \(or)
        }
        """)
    ]
  }
}

public struct CliTestDeclarationMacro: PeerMacro, Sendable {

  private static func doc_shrink(_ from: String) -> String {
    return from.replacing(Regex(#/^.*\/\/\/[\s]+/#), with: "")
  }

  public static func expansion(
    of node: AttributeSyntax,
    providingPeersOf declaration: some DeclSyntaxProtocol,
    in context: some MacroExpansionContext
  ) throws -> [DeclSyntax] {

    let test_name = declaration.cast(FunctionDeclSyntax.self).name

    let cli_test_expected_output = node.leadingTrivia.filter({ $0.isComment }).map({
      doc_shrink("\($0)")
    }).joined(separator: "\\n")

    let cli_test_driver_thunk_name = context.makeUniqueName(
      declaration.cast(FunctionDeclSyntax.self).name.text + "_thunk_")

    let (expected_decl, expected_label) = if cli_test_expected_output.isEmpty {
      ("let expected = \"\(test_name)\"", "withExpectedPath:")
    } else {
      ("let expected = \"\(cli_test_expected_output)\"", "withExpected:")
    }

    let cli_test_driver_thunk: DeclSyntax = """
      @Sendable private func \(cli_test_driver_thunk_name)() async throws {
        \(raw: expected_decl)

        _ = unsafe try await Testing.__requiringUnsafe(
          Testing.__requiringTry(
            Testing.__requiringAwait(swiftCliTestRunner(\(test_name), \(raw: expected_label) expected))))
      }
      """

    let cli_test_driver_generator_name = context.makeUniqueName(test_name.text + "_generator_")
    let cli_test_driver_generator: DeclSyntax = """
        @Sendable private func \(cli_test_driver_generator_name)() async -> Testing.Test {
          return .__function(
            named: "\(test_name)",
            in: nil,
            xcTestCompatibleSelector: Testing.__xcTestCompatibleSelector("\(test_name)"),
            traits: [],
            sourceLocation: Testing.SourceLocation.__here(),
            parameters: [],
            testFunction: \(cli_test_driver_thunk_name)
          )
        }
      """

    let cli_test_driver_accessor_name = context.makeUniqueName(test_name.text + "_accessor_")
    let cli_test_driver_accessor: DeclSyntax = """
        private nonisolated let \(cli_test_driver_accessor_name): Accessor = { outValue, type, _, _ in

          Testing.Test.__store(\(cli_test_driver_generator_name), into: outValue, asTypeAt: type)
        }
      """

    let cli_test_driver_content_record_name = context.makeUniqueName(
      declaration.cast(FunctionDeclSyntax.self).name.text + "_cr_")
    let cli_test_driver_cr: DeclSyntax = """
        private nonisolated let \(cli_test_driver_content_record_name): TestContentRecord = (
          0x7465_7374, /* indicate a test */
          0,
          unsafe \(cli_test_driver_accessor_name),
          0,
          0
        )
      """

    let cli_test_driver_content_container_name = context.makeUniqueName(test_name.text + "__🟡$_container_")
    let cli_test_driver_container: DeclSyntax = """
        struct \(cli_test_driver_content_container_name): Testing.__TestContentRecordContainer {
          nonisolated static var __testContentRecord: TestContentRecord {
            unsafe \(cli_test_driver_content_record_name)
          }
        }
      """
    return [
      cli_test_driver_thunk, cli_test_driver_generator, cli_test_driver_accessor,
      cli_test_driver_cr, cli_test_driver_container,
    ]
  }
}

@main
struct P4Macros: CompilerPlugin {
  var providingMacros: [Macro.Type] = [
    RequireResult.self, RequireErrorResult.self, UseOkResult.self, UseErrorResult.self,
    RequireNodeType.self, SkipUnlessNodeType.self, SkipUnlessNodesTypes.self, RequireNodesType.self,
    MustOr.self, CliTestDeclarationMacro.self,
  ]
}