import Foundation

final class FutureModificationDateClamp: @unchecked Sendable {
    private let lock = NSLock()
    private var clampDate: Date?

    init(clampDate: Date? = nil) {
        self.clampDate = clampDate
    }

    func clamp(url _: URL, modifiedAt: Date, now: Date) -> Date {
        guard modifiedAt > now else { return modifiedAt }
        return self.lock.withLock {
            if let clampDate = self.clampDate {
                return min(clampDate, now)
            }
            self.clampDate = now
            return now
        }
    }
}

public struct LocalAgentSessionScanner: Sendable {
    typealias ProcessOutputProvider = @Sendable ([String: String]) async -> String
    typealias CWDProvider = @Sendable ([Int32], [String: String]) async -> [Int32: String]

    private struct Rollout: Sendable {
        let url: URL
        let modifiedAt: Date
        let metadata: CodexRolloutMetadata
    }

    private struct ScanContext: Sendable {
        let homeDirectory: URL
        let host: String
        let now: Date
        let codexAppServerPresent: Bool
        let includeFileOnlySessions: Bool
        let threadMetadata: [String: CodexThreadMetadata]
    }

    public let config: SessionScanConfig
    private let futureModificationDateClamp = FutureModificationDateClamp()
    private let processOutputProvider: ProcessOutputProvider?
    private let cwdProvider: CWDProvider?
    private let didVisitDirectoryEntry: (@Sendable () -> Void)?

    public init(config: SessionScanConfig = SessionScanConfig()) {
        self.config = config
        self.processOutputProvider = nil
        self.cwdProvider = nil
        self.didVisitDirectoryEntry = nil
    }

    init(
        config: SessionScanConfig = SessionScanConfig(),
        processOutputProvider: @escaping ProcessOutputProvider,
        cwdProvider: @escaping CWDProvider,
        didVisitDirectoryEntry: (@Sendable () -> Void)? = nil)
    {
        self.config = config
        self.processOutputProvider = processOutputProvider
        self.cwdProvider = cwdProvider
        self.didVisitDirectoryEntry = didVisitDirectoryEntry
    }

    @concurrent
    public func scan(
        now: Date = Date(),
        environment: [String: String] = ProcessInfo.processInfo.environment,
        includeFileOnlySessions: Bool = true) async -> [AgentSession]
    {
        let processOutput = if let processOutputProvider = self.processOutputProvider {
            await processOutputProvider(environment)
        } else {
            await self.processOutput(environment: environment)
        }
        let allProcesses = AgentPSOutputParser.parse(processOutput)
        let processes = Array(AgentSessionCorrelation.newestProcessesFirst(
            AgentPSOutputParser.agentProcesses(from: allProcesses))
            .prefix(max(0, self.config.maxProcessCount)))
        guard Self.shouldScanSessionMetadata(
            hasAgentProcesses: !processes.isEmpty,
            includeFileOnlySessions: includeFileOnlySessions)
        else { return [] }
        let codexAppServerPresent = AgentPSOutputParser.hasCodexAppServer(in: allProcesses)
        let cwdByPID = if let cwdProvider = self.cwdProvider {
            await cwdProvider(processes.map(\ .pid), environment)
        } else {
            await self.cwdByPID(processes.map(\ .pid), environment: environment)
        }
        let codexCWDs = processes.compactMap { process -> String? in
            guard AgentPSOutputParser.provider(for: process) == .codex else { return nil }
            return cwdByPID[process.pid]
        }
        let homeDirectory = URL(fileURLWithPath: environment["HOME"] ?? NSHomeDirectory(), isDirectory: true)
        let codexHomeDirectory = URL(
            fileURLWithPath: environment["CODEX_HOME"] ?? homeDirectory.appendingPathComponent(".codex").path,
            isDirectory: true)
        let host = ProcessInfo.processInfo.hostName
        var directoryBudget = DirectoryMetadataScanBudget(
            maxEntryCount: self.config.maxDirectoryEntryCount,
            maxDepth: self.config.maxDirectoryDepth,
            timeLimit: includeFileOnlySessions
                ? self.config.directoryScanBudget
                : min(self.config.directoryScanBudget, self.config.adaptiveDirectoryScanBudget),
            didVisitEntry: self.didVisitDirectoryEntry)
        let rollouts: [Rollout] = if includeFileOnlySessions || !codexCWDs.isEmpty {
            self.codexRollouts(
                now: now,
                codexHomeDirectory: codexHomeDirectory,
                matchingCWDs: includeFileOnlySessions ? nil : codexCWDs,
                directoryBudget: &directoryBudget)
        } else {
            []
        }
        let threadMetadata = Self.codexThreadMetadata(
            rollouts: rollouts,
            codexHomeDirectory: codexHomeDirectory,
            environment: environment)
        return self.sessions(
            processes: processes,
            cwdByPID: cwdByPID,
            rollouts: rollouts,
            context: ScanContext(
                homeDirectory: homeDirectory,
                host: host,
                now: now,
                codexAppServerPresent: codexAppServerPresent,
                includeFileOnlySessions: includeFileOnlySessions,
                threadMetadata: threadMetadata),
            directoryBudget: &directoryBudget)
    }

    public static func shouldScanSessionMetadata(
        hasAgentProcesses: Bool,
        includeFileOnlySessions: Bool) -> Bool
    {
        hasAgentProcesses || includeFileOnlySessions
    }

    private static func codexThreadMetadata(
        rollouts: [Rollout],
        codexHomeDirectory: URL,
        environment: [String: String])
        -> [String: CodexThreadMetadata]
    {
        let sessionIDs = Set(rollouts.map(\.metadata.sessionID))
        let indexedNames = CodexThreadMetadataReader.indexedThreadNames(
            codexHomeDirectory: codexHomeDirectory,
            sessionIDs: sessionIDs)
        var groups: [String: (reader: CodexThreadMetadataReader, sessionIDs: Set<String>)] = [:]
        for rollout in rollouts {
            let resolvedWorkingDirectory = rollout.metadata.cwd.map {
                URL(fileURLWithPath: $0, isDirectory: true).standardizedFileURL
            }
            let reader = CodexThreadMetadataReader(
                codexHomeDirectory: codexHomeDirectory,
                environment: environment,
                resolvedWorkingDirectory: resolvedWorkingDirectory)
            let key = reader.databaseURL.path
            if var group = groups[key] {
                group.sessionIDs.insert(rollout.metadata.sessionID)
                groups[key] = group
            } else {
                groups[key] = (reader, [rollout.metadata.sessionID])
            }
        }

        var metadata: [String: CodexThreadMetadata] = [:]
        for group in groups.values {
            metadata.merge(group.reader.metadata(for: group.sessionIDs, indexedNames: indexedNames)) { _, latest in
                latest
            }
        }
        return metadata
    }

    private func sessions(
        processes: [AgentProcessRecord],
        cwdByPID: [Int32: String],
        rollouts: [Rollout],
        context: ScanContext,
        directoryBudget: inout DirectoryMetadataScanBudget) -> [AgentSession]
    {
        var sessions: [AgentSession] = []
        var matchedRolloutPaths = Set<String>()
        let claudeProcesses = processes.filter { AgentPSOutputParser.provider(for: $0) == .claude }
        let claudeCWDs = Set(claudeProcesses.compactMap { cwdByPID[$0.pid] })
        var claudeTranscriptsByCWD: [String: [ClaudeSessionProjectMapper.Transcript]] = [:]
        for cwd in claudeCWDs {
            claudeTranscriptsByCWD[cwd] = ClaudeSessionProjectMapper.transcripts(
                cwd: cwd,
                homeDirectory: context.homeDirectory,
                limit: self.config.maxClaudeTranscriptCountPerProject,
                now: context.now,
                budget: &directoryBudget,
                clampModificationDate: self.futureModificationDateClamp.clamp)
        }
        let claudeTranscripts = AgentSessionCorrelation.assignClaudeTranscripts(
            processes: claudeProcesses,
            cwdByPID: cwdByPID,
            transcriptsByCWD: claudeTranscriptsByCWD)
        let codexProcesses = processes.filter { AgentPSOutputParser.provider(for: $0) == .codex }
        let codexDescriptiveNamePIDs = AgentSessionCorrelation.unambiguousProcessIDs(
            processes: codexProcesses,
            cwdByPID: cwdByPID)

        for process in processes {
            guard let provider = AgentPSOutputParser.provider(for: process) else { continue }
            let cwd = cwdByPID[process.pid]
            switch provider {
            case .claude:
                let transcript = claudeTranscripts[process.pid]
                sessions.append(AgentSession(
                    id: transcript?.url.deletingPathExtension().lastPathComponent ?? "pid:\(process.pid)",
                    provider: .claude,
                    source: AgentPSOutputParser.source(for: process),
                    state: self.config.state(
                        lastActivityAt: transcript?.modifiedAt,
                        now: context.now,
                        hasLiveProcess: true),
                    pid: process.pid,
                    cwd: cwd,
                    projectName: Self.projectName(cwd),
                    startedAt: process.startedAt,
                    lastActivityAt: transcript?.modifiedAt,
                    transcriptPath: transcript?.url.path,
                    host: context.host))
            case .codex:
                let rollout = rollouts.first { candidate in
                    !matchedRolloutPaths.contains(candidate.url.path) &&
                        AgentSessionCorrelation.codexWorkingDirectoriesMatch(candidate.metadata.cwd, cwd)
                }
                if let rollout {
                    matchedRolloutPaths.insert(rollout.url.path)
                }
                let rolloutSource = rollout?.metadata.sessionSource
                sessions.append(AgentSession(
                    id: rollout?.metadata.sessionID ?? "pid:\(process.pid)",
                    provider: .codex,
                    source: rolloutSource == nil || rolloutSource == .unknown ? .cli : rolloutSource ?? .cli,
                    state: self.config.state(
                        lastActivityAt: rollout?.modifiedAt,
                        now: context.now,
                        hasLiveProcess: true),
                    pid: process.pid,
                    cwd: cwd ?? rollout?.metadata.cwd,
                    projectName: Self.projectName(cwd ?? rollout?.metadata.cwd),
                    sessionName: codexDescriptiveNamePIDs.contains(process.pid)
                        ? rollout?.metadata.descriptiveName(
                            threadMetadata: rollout.flatMap { context.threadMetadata[$0.metadata.sessionID] })
                        : nil,
                    startedAt: process.startedAt,
                    lastActivityAt: rollout?.modifiedAt,
                    transcriptPath: rollout?.url.path,
                    host: context.host))
            }
        }

        for rollout in rollouts
            where context.includeFileOnlySessions && !matchedRolloutPaths.contains(rollout.url.path)
        {
            guard var session = CodexRolloutFirstLineParser.makeSession(
                metadata: rollout.metadata,
                transcriptURL: rollout.url,
                modifiedAt: rollout.modifiedAt,
                host: context.host,
                config: self.config,
                now: context.now)
            else { continue }
            session.sessionName = rollout.metadata.descriptiveName(
                threadMetadata: context.threadMetadata[rollout.metadata.sessionID])
            session.source = AgentSessionCorrelation.fileOnlyCodexSource(
                metadataSource: session.source,
                appServerPresent: context.codexAppServerPresent)
            sessions.append(session)
        }

        var seen = Set<String>()
        return sessions
            .sorted { lhs, rhs in
                if lhs.state != rhs.state {
                    return lhs.state == .active
                }
                return (lhs.lastActivityAt ?? lhs.startedAt ?? .distantPast) >
                    (rhs.lastActivityAt ?? rhs.startedAt ?? .distantPast)
            }
            .filter { seen.insert("\($0.host):\($0.id)").inserted }
    }

    private func processOutput(environment: [String: String]) async -> String {
        let binary = ["/bin/ps", "/usr/bin/ps"].first { FileManager.default.isExecutableFile(atPath: $0) }
        guard let binary,
              let result = try? await SubprocessRunner.run(
                  binary: binary,
                  arguments: ["-axo", "pid=,ppid=,lstart=,command="],
                  environment: environment,
                  timeout: 5,
                  label: "agent session process scan")
        else { return "" }
        return result.stdout
    }

    private func cwdByPID(_ pids: [Int32], environment: [String: String]) async -> [Int32: String] {
        guard !pids.isEmpty else { return [:] }
        if let lsof = self.findExecutable("lsof", environment: environment) {
            let joinedPIDs = pids.map(String.init).joined(separator: ",")
            if let result = try? await SubprocessRunner.run(
                binary: lsof,
                arguments: ["-a", "-d", "cwd", "-Fn", "-p", joinedPIDs],
                environment: environment,
                timeout: 5,
                acceptsNonZeroExit: true,
                label: "agent session cwd scan")
            {
                return LSOFCWDOutputParser.parse(result.stdout)
            }
        }

        #if os(Linux)
        return Dictionary(uniqueKeysWithValues: pids.compactMap { pid in
            let path = "/proc/\(pid)/cwd"
            guard let destination = try? FileManager.default.destinationOfSymbolicLink(atPath: path) else { return nil }
            return (pid, destination)
        })
        #else
        return [:]
        #endif
    }

    private func codexRollouts(
        now: Date,
        codexHomeDirectory: URL,
        matchingCWDs: [String]?,
        directoryBudget: inout DirectoryMetadataScanBudget) -> [Rollout]
    {
        let root = codexHomeDirectory.appendingPathComponent("sessions", isDirectory: true)
        let calendar = Calendar(identifier: .gregorian)
        let days = [now, calendar.date(byAdding: .day, value: -1, to: now)].compactMap(\.self)
        let formatter = DateFormatter()
        formatter.locale = Locale(identifier: "en_US_POSIX")
        formatter.dateFormat = "yyyy/MM/dd"
        let fileManager = FileManager.default

        let candidates = days.flatMap { day -> [(url: URL, modifiedAt: Date)] in
            let directory = root.appendingPathComponent(formatter.string(from: day), isDirectory: true)
            return directoryBudget.files(in: directory, fileManager: fileManager).compactMap { file in
                guard file.lastPathComponent.hasPrefix("rollout-"), file.pathExtension == "jsonl",
                      let modifiedAt = try? file.resourceValues(
                          forKeys: [.contentModificationDateKey]).contentModificationDate
                else { return nil }
                return (file, self.futureModificationDateClamp.clamp(
                    url: file,
                    modifiedAt: modifiedAt,
                    now: now))
            }
        }.sorted { $0.modifiedAt > $1.modifiedAt }

        var remainingCWDs = matchingCWDs ?? []
        var rollouts: [Rollout] = []
        for candidate in candidates.prefix(max(0, self.config.maxCodexRolloutCount)) {
            guard directoryBudget.hasTimeRemaining() else { break }
            guard let metadata = CodexRolloutFirstLineParser.read(from: candidate.url) else { continue }
            rollouts.append(Rollout(url: candidate.url, modifiedAt: candidate.modifiedAt, metadata: metadata))
            if let index = remainingCWDs.firstIndex(where: {
                AgentSessionCorrelation.codexWorkingDirectoriesMatch(metadata.cwd, $0)
            }) {
                remainingCWDs.remove(at: index)
                if matchingCWDs != nil, remainingCWDs.isEmpty {
                    break
                }
            }
        }
        return rollouts
    }

    private func findExecutable(_ name: String, environment: [String: String]) -> String? {
        let path = environment["PATH"] ?? "/usr/local/bin:/usr/bin:/bin:/opt/homebrew/bin"
        return path.split(separator: ":")
            .map { String($0) + "/" + name }
            .first { FileManager.default.isExecutableFile(atPath: $0) }
    }

    private static func standardized(_ path: String?) -> String? {
        path.map { URL(fileURLWithPath: $0).standardizedFileURL.path }
    }

    private static func projectName(_ cwd: String?) -> String? {
        guard let cwd, !cwd.isEmpty else { return nil }
        return URL(fileURLWithPath: cwd).lastPathComponent
    }
}