tv-anarchy/Sources/TVAnarchyCore/Metadata/LocalLLMEpisodeRefiner.swift

import Foundation

/// Episode-title refiner backed by MLX on Apple Silicon (`episode_refiner.py`).
/// Invoked only on Title Library cache miss; successful answers are persisted by
/// `LibraryDisplayNames` — this type does not maintain its own filename cache.
public struct LocalLLMEpisodeRefiner: EpisodeTitleRefiner {
    public init() {}

    public func refineEpisode(showName: String, season: Int, episode: Int, path: String,
                              regexTitle: String?) -> EpisodeRefinement? {
        guard Self.session.healthy else { return nil }
        let filename = (path as NSString).lastPathComponent
        let payload: [String: Any] = [
            "contentKey": TitleLibrary.contentKey(showName: showName, season: season, episode: episode),
            "showName": showName,
            "season": season,
            "episode": episode,
            "filename": filename,
            "regexTitle": regexTitle ?? "",
        ]
        guard let data = try? JSONSerialization.data(withJSONObject: payload),
              let json = String(data: data, encoding: .utf8) else { return nil }

        let dir = RepoPaths.recommender.path
        let cmd = "cd \(Self.shq(dir)) && uv run python -m media_rec.episode_refiner \(Self.shq(json))"
        let r = ProcessRunner.runShell(cmd, timeout: 90, cwd: dir)
        guard r.ok,
              let out = r.stdout.trimmingCharacters(in: .whitespacesAndNewlines).data(using: .utf8),
              let decoded = try? JSONDecoder().decode(Refined.self, from: out) else {
            if !r.ok { Log.warn("episode refiner failed (exit \(r.status)): \(r.stderr.suffix(160))") }
            Self.session.recordFailure()
            return nil
        }
        Self.session.recordSuccess()
        let title = decoded.episodeTitle.trimmingCharacters(in: .whitespaces)
        guard !title.isEmpty else { return nil }
        return EpisodeRefinement(title: title, confidence: decoded.confidence)
    }

    private struct Refined: Decodable {
        let episodeTitle: String
        let confidence: Double
    }

    private static func shq(_ s: String) -> String {
        "'" + s.replacingOccurrences(of: "'", with: "'\\''") + "'"
    }

    private static let session = EpisodeRefinerSession()
}

/// Result of an episode-title refinement attempt.
public struct EpisodeRefinement: Sendable, Equatable {
    public var title: String
    public var confidence: Double

    public init(title: String, confidence: Double) {
        self.title = title
        self.confidence = confidence
    }
}

/// Seam for model-backed episode title extraction (MLX). Implementations return nil
/// when unavailable or unconfident — callers keep deterministic fallbacks.
public protocol EpisodeTitleRefiner: Sendable {
    func refineEpisode(showName: String, season: Int, episode: Int, path: String,
                       regexTitle: String?) -> EpisodeRefinement?
}

/// Session kill-switch — consecutive subprocess failures disable MLX for the process.
final class EpisodeRefinerSession: @unchecked Sendable {
    private let lock = NSLock()
    private var consecutiveFailures = 0
    private static let maxFailures = 2

    var healthy: Bool {
        lock.lock(); defer { lock.unlock() }
        return consecutiveFailures < Self.maxFailures
    }

    func recordFailure() {
        lock.lock(); defer { lock.unlock() }
        consecutiveFailures += 1
    }

    func recordSuccess() {
        lock.lock(); defer { lock.unlock() }
        consecutiveFailures = 0
    }

    /// Test hook — re-enable after simulated failures.
    func resetForTests() {
        lock.lock(); defer { lock.unlock() }
        consecutiveFailures = 0
    }
}