//===----------------------------------------------------------------------===//
// 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 ContainerAPIClient
import ContainerOS
import ContainerPersistence
import ContainerizationError
import ContainerizationExtras
import DNSServer
import Foundation
import Logging
import Synchronization
import SystemPackage

actor LocalhostDNSHandler: DNSHandler {
    private let ttl: UInt32
    private let watcher: DirectoryWatcher

    private var dns: [DNSName: IPv4Address]

    public init(configPath: FilePath = HostDNSResolver.defaultConfigPath, ttl: UInt32 = 5, log: Logger) {
        self.ttl = ttl

        self.watcher = DirectoryWatcher(directoryPath: configPath, log: log)
        self.dns = [DNSName: IPv4Address]()
    }

    public func monitorResolvers() async throws {
        var readyIterator = await self.watcher.readyEvents.makeAsyncIterator()

        try await self.watcher.startWatching { [weak self] filePaths in
            var dns: [DNSName: IPv4Address] = [:]
            let regex = try Regex(HostDNSResolver.localhostOptionsRegex)

            for file in filePaths.filter({
                $0.lastComponent?.string.starts(with: HostDNSResolver.containerizationPrefix) == true
            }) {
                let content = try String(contentsOfFile: file.string, encoding: .utf8)

                if let match = content.firstMatch(of: regex),
                    let ipv4 = (match[1].substring.flatMap { try? IPv4Address(String($0)) })
                {
                    guard let lastName = file.lastComponent?.string else {
                        continue
                    }
                    let name = String(lastName.dropFirst(HostDNSResolver.containerizationPrefix.count))
                    guard let dnsName = try? DNSName(name) else {
                        continue
                    }
                    dns[dnsName] = ipv4
                }
            }
            if let self {
                Task { await self.updateDNS(dns) }
            }
        }

        // Wait for the watcher to actually be watching before returning, so callers can rely on
        // resolver file changes being observed once this call completes, instead of racing the
        // watcher's own poll cadence.
        await readyIterator.next()
    }

    public func answer(query: Message) async throws -> Message? {
        guard let question = query.questions.first else {
            return nil
        }
        let n = question.name.hasSuffix(".") ? String(question.name.dropLast()) : question.name
        let key = try DNSName(labels: n.isEmpty ? [] : n.split(separator: ".", omittingEmptySubsequences: false).map(String.init))
        var record: ResourceRecord?
        switch question.type {
        case ResourceRecordType.host:
            if let ip = dns[key] {
                record = HostRecord<IPv4Address>(name: question.name, ttl: ttl, ip: ip)
            }
        case ResourceRecordType.host6:
            guard dns[key] != nil else {
                return nil
            }
            return Message(
                id: query.id,
                type: .response,
                returnCode: .noError,
                questions: query.questions,
                answers: []
            )
        default:
            return Message(
                id: query.id,
                type: .response,
                returnCode: .notImplemented,
                questions: query.questions,
                answers: []
            )
        }

        guard let record else {
            return nil
        }

        return Message(
            id: query.id,
            type: .response,
            returnCode: .noError,
            questions: query.questions,
            answers: [record]
        )
    }

    private func updateDNS(_ dns: [DNSName: IPv4Address]) {
        self.dns = dns
    }
}
