#!/usr/bin/env swift
//
//  generate-enterprise-config.swift
//  VoxyAI Enterprise Config Generator
//
//  CLI tool to generate signed and encrypted enterprise configuration plists
//  for deployment via MDM or config management.
//
//  Signing is performed remotely via VoxyAI's cloud signing service. The tool
//  sends only a SHA-256 hash of the encrypted payload -- your API keys and
//  sensitive data never leave your machine.
//
//  Usage:
//    swift generate-enterprise-config.swift \
//      --org "Acme Corp" \
//      --license "XXXX-XXXX-XXXX-XXXX" \
//      --api-key anthropic:sk-ant-xxx \
//      --api-key openai:sk-xxx \
//      --allowed-providers anthropic,openai \
//      --enforce-secret-scanning \
//      --enable-cloud-logging \
//      --blocked-keyword "acme" \
//      --blocked-keyword "project-phoenix" \
//      --custom-pattern "Internal Domain:.*\\.acme\\.com:regex" \
//      --output enterprise-config.plist
//
//  To generate a new Ed25519 keypair:
//    swift generate-enterprise-config.swift --generate-keys
//

import Foundation
import CryptoKit

// MARK: - Constants

let defaultSigningURL = "https://us-central1-voxyai-prod.cloudfunctions.net/sign-config"

// MARK: - Models

struct CustomScanPatternInput {
    let name: String
    let pattern: String
    let isRegex: Bool
}

struct MandatoryMCPServerInput: Codable {
    let name: String
    let transportType: String
    let command: String?
    let args: [String]?
    let url: String?
    let headers: [String: String]?
    let isEnabled: Bool
    let timeoutSeconds: Int?
    let isEnterpriseLocked: Bool
}

struct SigningResponse: Decodable {
    let success: Bool
    let signature: String?
    let error: String?
}

// MARK: - Key Generation

func generateEd25519Keypair() {
    let privateKey = Curve25519.Signing.PrivateKey()
    let publicKey = privateKey.publicKey

    let privateKeyBase64 = privateKey.rawRepresentation.base64EncodedString()
    let publicKeyBase64 = publicKey.rawRepresentation.base64EncodedString()

    print("Ed25519 Keypair Generated")
    print("=========================")
    print("")
    print("PRIVATE KEY (store in cloud function env var ENTERPRISE_SIGNING_PRIVATE_KEY):")
    print(privateKeyBase64)
    print("")
    print("PUBLIC KEY (embed in app as EnterpriseConstants.signaturePublicKey):")
    print(publicKeyBase64)
    print("")
    print("Store the private key in the cloud function environment.")
    print("The public key goes into EnterpriseConfig.swift.")
}

// MARK: - Remote Signing

func requestRemoteSignature(hashHex: String, licenseKey: String, orgName: String, signingURL: String) async throws -> Data {
    guard let url = URL(string: signingURL) else {
        throw NSError(domain: "ConfigGen", code: 1, userInfo: [NSLocalizedDescriptionKey: "Invalid signing URL: \(signingURL)"])
    }

    var request = URLRequest(url: url)
    request.httpMethod = "POST"
    request.setValue("application/json", forHTTPHeaderField: "Content-Type")
    request.timeoutInterval = 30

    let body: [String: String] = [
        "hash": hashHex,
        "licenseKey": licenseKey,
        "organizationName": orgName
    ]

    request.httpBody = try JSONSerialization.data(withJSONObject: body)

    let (data, response) = try await URLSession.shared.data(for: request)

    guard let httpResponse = response as? HTTPURLResponse else {
        throw NSError(domain: "ConfigGen", code: 2, userInfo: [NSLocalizedDescriptionKey: "Invalid response from signing service"])
    }

    let signingResponse = try JSONDecoder().decode(SigningResponse.self, from: data)

    guard httpResponse.statusCode == 200, signingResponse.success, let signatureBase64 = signingResponse.signature else {
        let errorMessage = signingResponse.error ?? "HTTP \(httpResponse.statusCode)"
        throw NSError(domain: "ConfigGen", code: 3, userInfo: [NSLocalizedDescriptionKey: "Signing failed: \(errorMessage)"])
    }

    guard let signatureData = Data(base64Encoded: signatureBase64) else {
        throw NSError(domain: "ConfigGen", code: 4, userInfo: [NSLocalizedDescriptionKey: "Invalid signature format from signing service"])
    }

    return signatureData
}

// MARK: - Config Generation

func generateConfig(
    orgName: String,
    licenseKey: String,
    apiKeys: [String: String],
    allowedProviders: [String]?,
    enforcedProvider: String?,
    allowUserOverride: Bool,
    disabledFeatures: [String]?,
    enforceSecretScanning: Bool,
    enableCloudLogging: Bool,
    blockedKeywords: [String]?,
    customPatterns: [CustomScanPatternInput]?,
    mcpEnabled: Bool?,
    allowedMCPServers: [String]?,
    blockedMCPServers: [String]?,
    mandatoryMCPServers: [MandatoryMCPServerInput]?,
    signingURL: String,
    outputPath: String
) async throws {
    // Build payload (sensitive data that gets encrypted)
    var payloadDict: [String: Any] = [
        "apiKeys": apiKeys
    ]

    if let keywords = blockedKeywords, !keywords.isEmpty {
        payloadDict["blockedKeywords"] = keywords
    }

    if let patterns = customPatterns, !patterns.isEmpty {
        payloadDict["customScanPatterns"] = patterns.map { pattern -> [String: Any] in
            return [
                "name": pattern.name,
                "pattern": pattern.pattern,
                "isRegex": pattern.isRegex
            ]
        }
    }

    if let mcpServers = mandatoryMCPServers, !mcpServers.isEmpty {
        let encoder = JSONEncoder()
        if let mcpData = try? encoder.encode(mcpServers),
           let mcpArray = try? JSONSerialization.jsonObject(with: mcpData) {
            payloadDict["mandatoryMCPServers"] = mcpArray
        }
    }

    let payloadData = try JSONSerialization.data(withJSONObject: payloadDict, options: [])

    // Derive encryption key from org name and license key
    let keyMaterial = "\(orgName)|\(licenseKey)|VoxyAI-Enterprise-Salt-2024"
    let keyData = SHA256.hash(data: Data(keyMaterial.utf8))
    let symmetricKey = SymmetricKey(data: keyData)

    // Encrypt payload with AES-256-GCM
    print("Encrypting payload... ", terminator: "")
    let sealedBox = try AES.GCM.seal(payloadData, using: symmetricKey)

    // Combine nonce + ciphertext + tag
    var encryptedData = Data()
    encryptedData.append(contentsOf: sealedBox.nonce)
    encryptedData.append(sealedBox.ciphertext)
    encryptedData.append(sealedBox.tag)
    print("done")

    // Compute SHA-256 hash of encrypted payload
    let hash = SHA256.hash(data: encryptedData)
    let hashHex = hash.compactMap { String(format: "%02x", $0) }.joined()

    // Request remote signature
    print("Requesting signature from VoxyAI signing service... ", terminator: "")
    let signatureData = try await requestRemoteSignature(
        hashHex: hashHex,
        licenseKey: licenseKey,
        orgName: orgName,
        signingURL: signingURL
    )
    print("done")
    print("Config signed successfully.")

    // Build plist structure
    var config: [String: Any] = [
        "OrganizationName": orgName,
        "LicenseKey": licenseKey,
        "EncryptedPayload": encryptedData,
        "PayloadSignature": signatureData,
        "AllowUserOverride": allowUserOverride,
        "ConfigVersion": 1
    ]

    if let providers = allowedProviders {
        config["AllowedProviders"] = providers
    }

    if let provider = enforcedProvider {
        config["EnforcedAIProvider"] = provider
    }

    if let features = disabledFeatures {
        config["DisabledFeatures"] = features
    }

    if enforceSecretScanning {
        config["EnforceSecretScanning"] = true
    }

    if enableCloudLogging {
        config["EnableCloudLogging"] = true
    }

    if let mcpEnabled = mcpEnabled {
        config["MCPEnabled"] = mcpEnabled
    }

    if let allowedMCP = allowedMCPServers, !allowedMCP.isEmpty {
        config["AllowedMCPServers"] = allowedMCP
    }

    if let blockedMCP = blockedMCPServers, !blockedMCP.isEmpty {
        config["BlockedMCPServers"] = blockedMCP
    }

    // Write plist
    let plistData = try PropertyListSerialization.data(fromPropertyList: config, format: .xml, options: 0)
    try plistData.write(to: URL(fileURLWithPath: outputPath))

    print("Enterprise config written to: \(outputPath)")
    print("  Organization: \(orgName)")
    print("  API keys: \(apiKeys.keys.sorted().joined(separator: ", "))")
    if let providers = allowedProviders {
        print("  Allowed providers: \(providers.joined(separator: ", "))")
    }
    if enforceSecretScanning {
        print("  Secret scanning: ENFORCED")
    }
    if enableCloudLogging {
        print("  Cloud logging: ENABLED")
    }
    if let mcpEnabled = mcpEnabled {
        print("  MCP: \(mcpEnabled ? "ENABLED" : "DISABLED")")
    }
    if let allowedMCP = allowedMCPServers {
        print("  Allowed MCP servers: \(allowedMCP.joined(separator: ", "))")
    }
    if let blockedMCP = blockedMCPServers {
        print("  Blocked MCP servers: \(blockedMCP.joined(separator: ", "))")
    }
    if let mcpServers = mandatoryMCPServers {
        print("  Mandatory MCP servers: \(mcpServers.count)")
    }
}

// MARK: - Argument Parsing

func printUsage() {
    print("""
    VoxyAI Enterprise Config Generator
    ===================================

    Generate a new Ed25519 keypair:
      swift generate-enterprise-config.swift --generate-keys

    Generate an enterprise config plist:
      swift generate-enterprise-config.swift \\
        --org "Organization Name" \\
        --license "LICENSE-KEY" \\
        [--api-key provider:key] \\
        [--allowed-providers provider1,provider2] \\
        [--enforced-provider provider] \\
        [--no-user-override] \\
        [--disabled-features feature1,feature2] \\
        [--enforce-secret-scanning] \\
        [--enable-cloud-logging] \\
        [--blocked-keyword keyword] \\
        [--custom-pattern "Name:pattern:regex|literal"] \\
        [--mcp-enabled true|false] \\
        [--allowed-mcp-server name] \\
        [--blocked-mcp-server name] \\
        [--mandatory-mcp-server config.json] \\
        [--signing-url URL] \\
        [--output path.plist]

    Signing is performed remotely via VoxyAI's cloud signing service.
    Only a SHA-256 hash of the encrypted payload is sent to the server.
    Your API keys and sensitive data never leave your machine.

    Provider names: anthropic, openai, gemini, groq, perplexity, mistral
    """)
}

// Parse command line arguments
let args = CommandLine.arguments

if args.count < 2 {
    printUsage()
    exit(1)
}

if args.contains("--help") || args.contains("-h") {
    printUsage()
    exit(0)
}

if args.contains("--generate-keys") {
    generateEd25519Keypair()
    exit(0)
}

// Parse config generation arguments
var orgName: String?
var licenseKey: String?
var apiKeys: [String: String] = [:]
var allowedProviders: [String]?
var enforcedProvider: String?
var allowUserOverride = true
var disabledFeatures: [String]?
var enforceSecretScanning = false
var enableCloudLogging = false
var blockedKeywords: [String] = []
var customPatterns: [CustomScanPatternInput] = []
var mcpEnabled: Bool?
var allowedMCPServers: [String] = []
var blockedMCPServers: [String] = []
var mandatoryMCPServers: [MandatoryMCPServerInput] = []
var signingURL = defaultSigningURL
var outputPath = "enterprise-config.plist"

var i = 1
while i < args.count {
    switch args[i] {
    case "--org":
        i += 1; orgName = args[i]
    case "--license":
        i += 1; licenseKey = args[i]
    case "--api-key":
        i += 1
        let parts = args[i].split(separator: ":", maxSplits: 1)
        if parts.count == 2 {
            apiKeys[String(parts[0]).lowercased()] = String(parts[1])
        }
    case "--allowed-providers":
        i += 1
        allowedProviders = args[i].split(separator: ",").map { String($0).trimmingCharacters(in: .whitespaces) }
    case "--enforced-provider":
        i += 1; enforcedProvider = args[i]
    case "--no-user-override":
        allowUserOverride = false
    case "--disabled-features":
        i += 1
        disabledFeatures = args[i].split(separator: ",").map { String($0).trimmingCharacters(in: .whitespaces) }
    case "--enforce-secret-scanning":
        enforceSecretScanning = true
    case "--enable-cloud-logging":
        enableCloudLogging = true
    case "--blocked-keyword":
        i += 1; blockedKeywords.append(args[i])
    case "--custom-pattern":
        i += 1
        let parts = args[i].split(separator: ":", maxSplits: 2)
        if parts.count == 3 {
            let isRegex = String(parts[2]).lowercased() == "regex"
            customPatterns.append(CustomScanPatternInput(
                name: String(parts[0]),
                pattern: String(parts[1]),
                isRegex: isRegex
            ))
        }
    case "--mcp-enabled":
        i += 1; mcpEnabled = args[i].lowercased() == "true"
    case "--allowed-mcp-server":
        i += 1; allowedMCPServers.append(args[i])
    case "--blocked-mcp-server":
        i += 1; blockedMCPServers.append(args[i])
    case "--mandatory-mcp-server":
        i += 1
        do {
            let fileURL = URL(fileURLWithPath: args[i])
            let data = try Data(contentsOf: fileURL)
            let server = try JSONDecoder().decode(MandatoryMCPServerInput.self, from: data)
            mandatoryMCPServers.append(server)
        } catch {
            print("ERROR: Failed to read MCP server config '\(args[i])': \(error.localizedDescription)")
            exit(1)
        }
    case "--signing-url":
        i += 1; signingURL = args[i]
    case "--output":
        i += 1; outputPath = args[i]
    default:
        break
    }
    i += 1
}

guard let org = orgName else {
    print("ERROR: --org is required")
    exit(1)
}

guard let license = licenseKey else {
    print("ERROR: --license is required")
    exit(1)
}

// Run async config generation
let semaphore = DispatchSemaphore(value: 0)
var exitCode: Int32 = 0

Task {
    do {
        try await generateConfig(
            orgName: org,
            licenseKey: license,
            apiKeys: apiKeys,
            allowedProviders: allowedProviders,
            enforcedProvider: enforcedProvider,
            allowUserOverride: allowUserOverride,
            disabledFeatures: disabledFeatures,
            enforceSecretScanning: enforceSecretScanning,
            enableCloudLogging: enableCloudLogging,
            blockedKeywords: blockedKeywords.isEmpty ? nil : blockedKeywords,
            customPatterns: customPatterns.isEmpty ? nil : customPatterns,
            mcpEnabled: mcpEnabled,
            allowedMCPServers: allowedMCPServers.isEmpty ? nil : allowedMCPServers,
            blockedMCPServers: blockedMCPServers.isEmpty ? nil : blockedMCPServers,
            mandatoryMCPServers: mandatoryMCPServers.isEmpty ? nil : mandatoryMCPServers,
            signingURL: signingURL,
            outputPath: outputPath
        )
    } catch {
        print("ERROR: \(error.localizedDescription)")
        exitCode = 1
    }
    semaphore.signal()
}

semaphore.wait()
exit(exitCode)