//===----------------------------------------------------------------------===//
// Copyright © 2025 Apple Inc. and the Containerization project authors. All rights reserved.
//
// 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 AsyncHTTPClient
import ContainerizationError
import ContainerizationOS
import Foundation
import Logging
import NIO
import NIOHTTP1

#if os(macOS)
import Network
#endif

/// Data used to control retry behavior for `RegistryClient`.
public struct RetryOptions: Sendable {
    /// The maximum number of retries to attempt before failing.
    public var maxRetries: Int
    /// The retry interval in nanoseconds.
    public var retryInterval: UInt64
    /// A provided closure to handle if a given HTTP response should be
    /// retried.
    public var shouldRetry: (@Sendable (HTTPClientResponse) -> Bool)?

    public init(maxRetries: Int, retryInterval: UInt64, shouldRetry: (@Sendable (HTTPClientResponse) -> Bool)? = nil) {
        self.maxRetries = maxRetries
        self.retryInterval = retryInterval
        self.shouldRetry = shouldRetry
    }
}

/// A client for interacting with OCI compliant container registries.
public final class RegistryClient: ContentClient {
    private static let defaultRetryOptions = RetryOptions(
        maxRetries: 3,
        retryInterval: 1_000_000_000,
        shouldRetry: ({ response in
            response.status.code >= 500
        })
    )

    let client: HTTPClient
    let base: URLComponents
    let clientID: String
    let authentication: Authentication?
    let retryOptions: RetryOptions?
    let bufferSize: Int

    public convenience init(
        reference: String,
        insecure: Bool = false,
        auth: Authentication? = nil,
        logger: Logger? = nil
    ) throws {
        let ref = try Reference.parse(reference)
        guard let domain = ref.resolvedDomain else {
            throw ContainerizationError(.invalidArgument, message: "Invalid domain for image reference \(reference)")
        }
        let scheme = insecure ? "http" : "https"
        let _url = "\(scheme)://\(domain)"
        guard let url = URL(string: _url) else {
            throw ContainerizationError(.invalidArgument, message: "Cannot convert \(_url) to URL")
        }
        guard let host = url.host else {
            throw ContainerizationError(.invalidArgument, message: "Invalid host \(domain)")
        }
        let port = url.port
        self.init(
            host: host,
            scheme: scheme,
            port: port,
            authentication: auth,
            retryOptions: Self.defaultRetryOptions
        )
    }

    public init(
        host: String,
        scheme: String? = "https",
        port: Int? = nil,
        authentication: Authentication? = nil,
        clientID: String? = nil,
        retryOptions: RetryOptions? = nil,
        bufferSize: Int = Int(4.mib()),
        logger: Logger? = nil
    ) {
        var components = URLComponents()
        components.scheme = scheme
        components.host = host
        components.port = port

        self.base = components
        self.clientID = clientID ?? "containerization-registry-client"
        self.authentication = authentication
        self.retryOptions = retryOptions
        self.bufferSize = bufferSize
        var httpConfiguration = HTTPClient.Configuration()
        let proxyConfig: HTTPClient.Configuration.Proxy? = {
            let proxyEnv = ProcessInfo.processInfo.environment["HTTP_PROXY"]
            guard let proxyEnv else {
                return nil
            }
            guard let url = URL(string: proxyEnv), let host = url.host(), let port = url.port else {
                return nil
            }
            return .server(host: host, port: port)
        }()
        httpConfiguration.proxy = proxyConfig
        if let logger {
            self.client = HTTPClient(eventLoopGroupProvider: .singleton, configuration: httpConfiguration, backgroundActivityLogger: logger)
        } else {
            self.client = HTTPClient(eventLoopGroupProvider: .singleton, configuration: httpConfiguration)
        }
    }

    deinit {
        _ = client.shutdown()
    }

    func host() -> String {
        base.host ?? ""
    }

    internal func request<T>(
        components: URLComponents,
        method: HTTPMethod = .GET,
        bodyClosure: () throws -> HTTPClientRequest.Body? = { nil },
        headers: [(String, String)]? = nil,
        closure: (HTTPClientResponse) async throws -> T
    ) async throws -> T {
        guard let path = components.url?.absoluteString else {
            throw ContainerizationError(.invalidArgument, message: "Invalid url \(components.path)")
        }

        var request = HTTPClientRequest(url: path)
        request.method = method

        var currentToken: TokenResponse?
        let token: String? = try await {
            if let basicAuth = authentication {
                return try await basicAuth.token()
            }
            return nil
        }()

        if let token {
            request.headers.add(name: "Authorization", value: "\(token)")
        }

        // Add any arbitrary headers
        headers?.forEach { (k, v) in request.headers.add(name: k, value: v) }
        var retryCount = 0
        var response: HTTPClientResponse?
        while true {
            request.body = try bodyClosure()
            do {
                let _response = try await client.execute(request, deadline: .distantFuture)
                response = _response
                if _response.status == .unauthorized || _response.status == .forbidden {
                    let authHeader = _response.headers[TokenRequest.authenticateHeaderName]
                    let tokenRequest: TokenRequest
                    do {
                        tokenRequest = try self.createTokenRequest(parsing: authHeader)
                    } catch {
                        // The server did not tell us how to authenticate our requests,
                        // Or we do not support scheme the server is requesting for.
                        // Throw the 401/403 to the caller, and let them decide how to proceed.
                        throw RegistryClient.Error.invalidStatus(url: path, _response.status)
                    }
                    if let ct = currentToken, ct.isValid(scope: tokenRequest.scope) {
                        break
                    }
                    let _currentToken = try await fetchToken(request: tokenRequest)
                    guard let token = _currentToken.getToken() else {
                        throw ContainerizationError(.internalError, message: "Failed to fetch Bearer token")
                    }
                    currentToken = _currentToken
                    request.headers.replaceOrAdd(name: "Authorization", value: token)
                    retryCount += 1
                    continue
                }
                guard let retryOptions = self.retryOptions else {
                    break
                }
                guard retryCount < retryOptions.maxRetries else {
                    break
                }
                guard let shouldRetry = retryOptions.shouldRetry, shouldRetry(_response) else {
                    break
                }
                retryCount += 1
                try await Task.sleep(nanoseconds: retryOptions.retryInterval)
                continue
            } catch let err as RegistryClient.Error {
                throw err
            } catch {
                #if os(macOS)
                if let err = error as? NWError {
                    if err.errorCode == kDNSServiceErr_NoSuchRecord {
                        throw ContainerizationError(.internalError, message: "No Such DNS Record \(host())")
                    }
                }
                #endif
                guard let retryOptions = self.retryOptions, retryCount < retryOptions.maxRetries else {
                    throw error
                }
                retryCount += 1
                try await Task.sleep(nanoseconds: retryOptions.retryInterval)
            }
        }
        guard let response else {
            throw ContainerizationError(.internalError, message: "Invalid response")
        }
        return try await closure(response)
    }

    internal func requestData(
        components: URLComponents,
        headers: [(String, String)]? = nil
    ) async throws -> Data {
        try await request(components: components, method: .GET, headers: headers) { response in
            guard response.status == .ok else {
                let url = components.url?.absoluteString ?? "unknown"
                throw Error.invalidStatus(url: url, response.status)
            }

            var body = try await response.body.collect(upTo: self.bufferSize)
            guard let bytes = body.readBytes(length: body.readableBytes) else {
                throw ContainerizationError(.internalError, message: "Cannot read bytes from HTTP response")
            }
            return Data(bytes)
        }
    }

    internal func requestJSON<T: Decodable>(
        components: URLComponents,
        headers: [(String, String)]? = nil
    ) async throws -> T {
        let data = try await self.requestData(components: components, headers: headers)
        return try JSONDecoder().decode(T.self, from: data)
    }

    /// A minimal endpoint, mounted at /v2/ will provide version support information based on its response statuses.
    /// See https://distribution.github.io/distribution/spec/api/#api-version-check
    public func ping() async throws {
        var components = base
        components.path = "/v2/"

        try await request(components: components) { response in
            guard response.status == .ok else {
                let url = components.url?.absoluteString ?? "unknown"
                throw Error.invalidStatus(url: url, response.status)
            }
        }
    }
}