//===----------------------------------------------------------------------===//
// Copyright © 2026 Apple Inc. and the container project authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//   https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//===----------------------------------------------------------------------===//

import ArgumentParser
import ContainerAPIClient
import ContainerLog
import ContainerPersistence
import ContainerResource
import ContainerizationError
import Darwin
import Foundation
import Logging
import TerminalProgress

public struct ClusterCreate: AsyncParsableCommand {
    public init() {}

    public static let configuration = CommandConfiguration(
        commandName: "create",
        abstract: "Create and start a local Kubernetes cluster and optional worker nodes"
    )

    @Option(name: .long, help: "Cluster name")
    var name: String = ClusterHelper.defaultName

    @Flag(name: [.customLong("rm"), .long], help: "Remove the cluster container after it stops")
    var remove: Bool = false

    @OptionGroup(title: "Resource options")
    var resourceFlags: Flags.Resource

    @OptionGroup(title: "Registry options")
    var registryFlags: Flags.Registry

    @OptionGroup(title: "Image fetch options")
    var imageFetchFlags: Flags.ImageFetch

    @Option(help: "Node image reference")
    var nodeImage: String = ClusterHelper.nodeImage

    @Option(name: .long, help: "Optional path to a CNI manifest to apply, or \"NONE\" to skip installing a CNI.")
    var cni: String?

    @Option(name: .long, help: "Number of worker nodes to create")
    var workers: UInt = 0

    /// Task count is `4 + workers` as an `Int`, so keep `workers` far enough below `Int.max` to avoid overflow.
    static let maxWorkers: UInt = 1024

    public func validate() throws {
        guard workers <= Self.maxWorkers else {
            throw ValidationError("--workers must be at most \(Self.maxWorkers)")
        }
    }

    public func run() async throws {
        LoggingSystem.bootstrap { _ in StderrLogHandler() }
        let log = Logger(label: ClusterHelper.pluginName)

        guard ManagedContainer.nameValid(name) else {
            throw ContainerizationError(.invalidArgument, message: "cluster name \(name) is not a valid container ID")
        }

        let skipCNI = cni?.uppercased() == ClusterHelper.noCNIName
        if let cni, !skipCNI {
            guard FileManager.default.fileExists(atPath: cni) else {
                throw ContainerizationError(.invalidArgument, message: "CNI manifest not found at \(cni)")
            }
        }

        // Fail before provisioning the node VM.
        _ = try ClusterHelper.kubernetesVersion(nodeImage: nodeImage)

        let isTTY = isatty(FileHandle.standardError.fileDescriptor) == 1
        // TotalTasks include:
        // start, kubeadm init, one per worker, wait, write kubeconfig.
        let progressConfig = try ProgressConfig(
            showSpinner: isTTY,
            showTasks: true,
            showItems: true,
            ignoreSmallSize: true,
            totalTasks: 4 + Int(workers),
            clearOnFinish: isTTY,
            outputMode: isTTY ? .ansi : .plain
        )

        let progress = ProgressBar(config: progressConfig)
        defer { progress.finish() }
        progress.start()

        let containerSystemConfig: ContainerSystemConfig = try await ConfigurationLoader.load()
        let fqdn = ClusterHelper.fqdn(for: name, domain: containerSystemConfig.dns.domain)

        let provisioner = try LinuxNodeProvisioner(
            clusterName: name,
            roles: workers == 0 ? [StandardRoles.controlPlane, StandardRoles.worker] : [StandardRoles.controlPlane],
            nodeImage: nodeImage,
            cpus: resourceFlags.cpus,
            memory: resourceFlags.memory,
            registryScheme: registryFlags.scheme,
            maxConcurrentDownloads: imageFetchFlags.maxConcurrentDownloads,
            remove: remove,
            fqdn: fqdn
        )

        progress.set(description: "Starting cluster")
        try await provisioner.provision(name: name, log: log)

        let client = ContainerClient()
        var workerNames: [String] = []
        do {
            let vmIP = try await provisioner.address(name: name, log: log)
            var sans = ["127.0.0.1"]
            if let fqdn { sans.append(contentsOf: [vmIP, fqdn]) }

            progress.set(description: "Running kubeadm init")
            try await ClusterHelper.prepareNode(nodeID: name, client: client, log: log)
            try await ClusterHelper.bootstrapControlPlane(
                nodeID: name, nodeImage: nodeImage, apiServerSANs: sans, advertiseAddress: vmIP,
                schedulable: provisioner.roles.contains(StandardRoles.worker),
                cniManifestPath: cni,
                client: client, log: log)

            if workers > 0 {
                let (token, caCertHash) = try await ClusterHelper.createJoinToken(nodeID: name, client: client)
                let controlPlaneEndpoint = "\(vmIP):\(ClusterHelper.clusterContainerPort)"
                for i in 1...workers {
                    let workerName = "\(name)-worker-\(i)"
                    progress.set(description: "Joining worker \(workerName)")
                    let workerProvisioner = try LinuxNodeProvisioner(
                        clusterName: name,
                        roles: [StandardRoles.worker],
                        nodeImage: nodeImage,
                        cpus: resourceFlags.cpus,
                        memory: resourceFlags.memory,
                        registryScheme: registryFlags.scheme,
                        maxConcurrentDownloads: imageFetchFlags.maxConcurrentDownloads,
                        remove: remove
                    )
                    workerNames.append(workerName)
                    try await workerProvisioner.provision(name: workerName, log: log)
                    try await workerProvisioner.join(
                        name: workerName, controlPlaneEndpoint: controlPlaneEndpoint,
                        token: token, caCertHash: caCertHash, log: log)
                    if skipCNI {
                        try await workerProvisioner.waitForRegistered(name: workerName, log: log)
                    } else {
                        try await workerProvisioner.waitForReady(name: workerName, log: log)
                    }
                }
            }

            if skipCNI {
                progress.set(description: "Waiting for API server")
                try await ClusterHelper.waitForAPIServer(containerId: name, client: client, log: log)
            } else {
                progress.set(description: "Waiting for cluster to be ready")
                try await ClusterHelper.waitForReady(containerId: name, client: client, log: log)
            }

            progress.set(description: "Writing kubeconfig")
            let rawConfig = try await ClusterHelper.fetchConfig(containerId: name, client: client, log: log)
            let kubeConfig = try await ClusterHelper.transformConfig(rawConfig, containerId: name, fqdn: fqdn, client: client)
            do {
                try ClusterHelper.mergeConfig(kubeConfig, containerId: name, setCurrentContext: true, log: log)
            } catch {
                log.warning("failed to write kubeconfig", metadata: ["name": "\(name)", "error": "\(error)"])
                log.info("cluster is running; use 'container cluster write-config --name \(name)' to write the kubeconfig")
            }
        } catch {
            for workerName in workerNames {
                await Self.teardownIfOwned(workerName, client: client, provisioner: provisioner, log: log)
            }
            await Self.teardownIfOwned(name, client: client, provisioner: provisioner, log: log)
            try? ClusterHelper.removeConfig(containerId: name, log: log)
            throw error
        }

        progress.finish()
        print(name)
    }

    /// Tears down `containerName` only if it's a container this plugin created
    private static func teardownIfOwned(
        _ containerName: String, client: ContainerClient, provisioner: LinuxNodeProvisioner, log: Logger
    ) async {
        guard let container = try? await client.get(id: containerName) else { return }
        guard container.configuration.labels[ResourceLabelKeys.plugin] == ClusterHelper.pluginName else {
            log.warning(
                "container exists but is not owned by this plugin, refusing teardown",
                metadata: ["name": "\(containerName)"])
            return
        }
        try? await provisioner.teardown(name: containerName, log: log)
    }
}