Flights/FlightsTests/LoadFactorServiceTests.swift

import XCTest
@testable import Flights

/// TDD **red phase** for ``LoadFactorService``.
///
/// These tests pin down behaviour the current implementation gets wrong
/// (timezone handling, peak-season detection in airport-local time,
/// equipment-swap edge cases, clamping) plus the correctness it already
/// has (confidence buckets, nil-on-missing-record). Every test is written
/// against the *future* `estimate(...)` signature that takes an
/// ``AirportDatabase`` so the service can resolve the origin airport's
/// timezone instead of leaning on a fixed UTC calendar.
///
/// Expected initial state when this file lands:
///   - Tests calling the new signature fail to compile (the new
///     `database:` parameter doesn't exist yet). That's the failing red.
///   - Phase 3 adds the parameter + timezone lookup + edge-case guards;
///     these tests then go green.
///
/// All assertions rely on the bundled ``bts_bundle.json``. Records used:
///   - ``WN_1701_OAK_BUR``  (leisure, OAK origin, Pacific TZ)
///   - ``UA_1_SFO_EWR``     (business, SFO origin, Pacific TZ)
///   - ``WN_5_DAL_HOU``     (leisure, high baseline → clamping)
///   - ``WN_61_DAL_HOU``    (leisure, mid-bucket confidence)
///   - ``AA_1000_ORD_DFW``  (high-bucket confidence)
@MainActor
final class LoadFactorServiceTests: XCTestCase {

    // Shared so we don't reload the BTS bundle / airports JSON per test.
    private static let airportDatabase = AirportDatabase()
    private static let service = LoadFactorService()

    private var airportDatabase: AirportDatabase { Self.airportDatabase }
    private var service: LoadFactorService { Self.service }

    // MARK: - Helpers

    /// Builds a Date from an ISO-8601 string with explicit offset, e.g.
    /// "2026-06-07T18:00:00-07:00".
    private func date(_ iso: String, file: StaticString = #file, line: UInt = #line) -> Date {
        let formatter = ISO8601DateFormatter()
        formatter.formatOptions = [.withInternetDateTime]
        guard let d = formatter.date(from: iso) else {
            XCTFail("Could not parse ISO date: \(iso)", file: file, line: line)
            return Date()
        }
        return d
    }

    // MARK: - 1. Timezone correctness (weekend detection)

    /// 6 PM Sunday at OAK (PDT) is *Sunday* in airport-local time, even
    /// though it's already Monday UTC. The current implementation uses a
    /// UTC calendar (LoadFactorService.swift:69-72), so it misses the
    /// weekend bump for west-coast late-evening departures.
    ///
    /// Both the weekend leisure (+5%) and the peak-season (+7%) bumps
    /// should fire here. Against the current bug, only peak fires —
    /// asserting `predicted >= base + 0.10` proves both bumps stacked.
    func test_weekendBump_appliesInAirportLocalTime_notUTC() async throws {
        let carrier = "WN"
        let flight = 1701
        let origin = "OAK"
        let dest = "BUR"
        // Sunday 6 PM PDT == Monday 1 AM UTC.
        let depart = date("2026-06-07T18:00:00-07:00")

        guard let base = await BTSDataStore.shared.record(
            carrier: carrier, flightNumber: flight, origin: origin, dest: dest
        ) else {
            throw XCTSkip("Bundled BTS bundle missing \(carrier)\(flight) \(origin)→\(dest); cannot run timezone test")
        }

        let estimate = await service.estimate(
            carrier: carrier,
            flightNumber: flight,
            origin: origin,
            dest: dest,
            date: depart,
            database: airportDatabase,
            liveSeats: nil
        )

        let result = try XCTUnwrap(estimate, "estimate(...) returned nil for a record that exists in the bundle")

        // Weekend leisure (+5%) + peak-season June (+7%) = at least +12%
        // on top of the base. Account for tiny FP drift with a 0.5% slack.
        let expected = base.avgLoadFactor + 0.05 + 0.07
        XCTAssertGreaterThanOrEqual(
            result.predicted,
            min(1.0, expected) - 0.005,
            "OAK 6 PM Sunday PDT should pick up the weekend leisure bump; current UTC-only code drops it."
        )
        XCTAssertTrue(
            result.basis.lowercased().contains("weekend"),
            "Basis string should mention the weekend adjustment, got: \(result.basis)"
        )
    }

    // MARK: - 2. Peak-season detection in airport-local time

    /// Midnight UTC on 1 July from an SFO origin is *17:00 on 30 June* in
    /// airport-local time, which means **no peak-season bump should
    /// apply**. Current code reads the month from a UTC calendar and
    /// over-counts this as July.
    func test_peakSeason_usesAirportLocalMonth_notUTC() async throws {
        let carrier = "UA"
        let flight = 1
        let origin = "SFO"
        let dest = "EWR"
        // Midnight UTC on 1 July → 5 PM PDT on 30 June at SFO.
        let depart = date("2026-07-01T00:00:00Z")

        guard let base = await BTSDataStore.shared.record(
            carrier: carrier, flightNumber: flight, origin: origin, dest: dest
        ) else {
            throw XCTSkip("Bundled BTS bundle missing \(carrier)\(flight) \(origin)→\(dest); cannot run peak-season test")
        }

        let estimate = await service.estimate(
            carrier: carrier,
            flightNumber: flight,
            origin: origin,
            dest: dest,
            date: depart,
            database: airportDatabase,
            liveSeats: nil
        )

        let result = try XCTUnwrap(estimate)

        // Tuesday 30 June in airport-local time: weekday, not peak season.
        // UA is "business" but the day is a weekday so no weekday bump
        // either. Prediction should equal the base within FP tolerance.
        XCTAssertEqual(
            result.predicted,
            base.avgLoadFactor,
            accuracy: 0.005,
            "30 June local should not trigger the +7% peak-season bump"
        )
        XCTAssertFalse(
            result.basis.lowercased().contains("peak season"),
            "Basis should not mention peak season, got: \(result.basis)"
        )
    }

    // MARK: - 3. Equipment-swap edge cases

    /// When the live aircraft has *more* seats than the historical avg,
    /// we should not bump up — the ratio path is only meant to scale
    /// predictions higher when a smaller jet is operating the segment.
    func test_equipmentSwap_largerAircraftDoesNotBumpUp() async throws {
        let carrier = "WN"
        let flight = 61
        let origin = "DAL"
        let dest = "HOU"
        let depart = date("2026-09-15T14:00:00-05:00") // Tuesday, non-peak, non-weekend

        guard let base = await BTSDataStore.shared.record(
            carrier: carrier, flightNumber: flight, origin: origin, dest: dest
        ) else {
            throw XCTSkip("Bundled BTS bundle missing \(carrier)\(flight) \(origin)→\(dest)")
        }

        // Live seats far above the historical avg (175). A bigger plane
        // should NOT push the prediction higher.
        let bigger = base.avgSeats + 200

        let estimate = await service.estimate(
            carrier: carrier,
            flightNumber: flight,
            origin: origin,
            dest: dest,
            date: depart,
            database: airportDatabase,
            liveSeats: bigger
        )

        let result = try XCTUnwrap(estimate)

        XCTAssertLessThanOrEqual(
            result.predicted,
            base.avgLoadFactor + 0.005,
            "Bigger live aircraft must not bump prediction up"
        )
    }

    /// liveSeats == 0 used to be a divide-by-zero hazard. We must guard
    /// so the call returns a normal estimate (no ratio applied).
    func test_equipmentSwap_liveSeatsZeroDoesNotDivideByZero() async throws {
        let carrier = "WN"
        let flight = 61
        let origin = "DAL"
        let dest = "HOU"
        let depart = date("2026-09-15T14:00:00-05:00")

        guard await BTSDataStore.shared.record(
            carrier: carrier, flightNumber: flight, origin: origin, dest: dest
        ) != nil else {
            throw XCTSkip("Bundled BTS bundle missing \(carrier)\(flight) \(origin)→\(dest)")
        }

        let estimate = await service.estimate(
            carrier: carrier,
            flightNumber: flight,
            origin: origin,
            dest: dest,
            date: depart,
            database: airportDatabase,
            liveSeats: 0
        )

        let result = try XCTUnwrap(estimate, "Service should still return an estimate when liveSeats == 0")
        XCTAssertTrue(result.predicted.isFinite, "Prediction must not be NaN/Inf when liveSeats == 0")
        XCTAssertFalse(result.basis.lowercased().contains("smaller aircraft"),
                       "liveSeats == 0 must not trigger the smaller-aircraft path")
    }

    /// If, in some future BTS record, ``avgSeats`` is 0 we must not
    /// crash. The bundled bundle has no such record today, so this test
    /// just exercises the code path with a sane liveSeats and a real
    /// record and asserts no crash + finite output. The Phase 3 fix
    /// should guard `base.avgSeats > 0` before doing the ratio math.
    func test_equipmentSwap_zeroAvgSeatsDoesNotCrash() async throws {
        let carrier = "WN"
        let flight = 61
        let origin = "DAL"
        let dest = "HOU"
        let depart = date("2026-09-15T14:00:00-05:00")

        guard await BTSDataStore.shared.record(
            carrier: carrier, flightNumber: flight, origin: origin, dest: dest
        ) != nil else {
            throw XCTSkip("Bundled BTS bundle missing \(carrier)\(flight) \(origin)→\(dest)")
        }

        // Small but positive — exercises the ratio branch on a real
        // record so any future regression that drops the avgSeats > 0
        // guard would surface here when paired with a zero-seats record.
        let estimate = await service.estimate(
            carrier: carrier,
            flightNumber: flight,
            origin: origin,
            dest: dest,
            date: depart,
            database: airportDatabase,
            liveSeats: 1
        )

        let result = try XCTUnwrap(estimate)
        XCTAssertTrue(result.predicted.isFinite,
                      "Prediction must remain finite even when the seat ratio is extreme")
    }

    // MARK: - 4. Clamping

    /// Sunday 7 June 2026 at DAL is a leisure-carrier weekend in peak
    /// season — stacking +5% + +7% + a smaller-aircraft ratio bump must
    /// clamp at 1.0, never exceed it.
    func test_predictionClampsAtOne_evenAfterStackedBumps() async throws {
        let carrier = "WN"
        let flight = 5
        let origin = "DAL"
        let dest = "HOU"
        // Sunday 2026-06-07 at noon CDT — Sunday in both UTC and local.
        let depart = date("2026-06-07T12:00:00-05:00")

        guard await BTSDataStore.shared.record(
            carrier: carrier, flightNumber: flight, origin: origin, dest: dest
        ) != nil else {
            throw XCTSkip("Bundled BTS bundle missing \(carrier)\(flight) \(origin)→\(dest)")
        }

        // Aggressive smaller-aircraft ratio to make sure stacked bumps
        // would otherwise blow past 1.0.
        let estimate = await service.estimate(
            carrier: carrier,
            flightNumber: flight,
            origin: origin,
            dest: dest,
            date: depart,
            database: airportDatabase,
            liveSeats: 100
        )

        let result = try XCTUnwrap(estimate)
        XCTAssertLessThanOrEqual(result.predicted, 1.0,
                                 "Predicted load factor must clamp at 1.0")
        XCTAssertGreaterThanOrEqual(result.predicted, 0.0,
                                    "Predicted load factor must clamp at 0.0")
    }

    // MARK: - 5. Confidence buckets

    /// sampleSize >= 60 → 0.85 confidence.
    /// Picks the first record in the bundle with totalFlights >= 60.
    func test_confidence_highBucket_when60OrMoreFlights() async throws {
        let all = await BTSDataStore.shared.allRecordsKeyed()
        guard let (key, _) = all
            .filter({ $0.value.totalFlights >= 60 })
            .first
        else {
            throw XCTSkip("Bundled BTS bundle has no record with totalFlights >= 60")
        }

        let parts = key.split(separator: "_").map(String.init)
        guard parts.count == 4, let fn = Int(parts[1]) else {
            XCTFail("Unexpected BTS key shape: \(key)")
            return
        }

        let depart = date("2026-09-15T14:00:00-05:00") // Tuesday, non-peak
        let estimate = await service.estimate(
            carrier: parts[0],
            flightNumber: fn,
            origin: parts[2],
            dest: parts[3],
            date: depart,
            database: airportDatabase,
            liveSeats: nil
        )
        let result = try XCTUnwrap(estimate)
        XCTAssertEqual(result.confidence, 0.85, accuracy: 0.0001,
                       "totalFlights >= 60 must map to 0.85 confidence")
    }

    /// sampleSize 20-59 → 0.65 confidence.
    /// Picks the first record in the bundle with 20 <= totalFlights < 60.
    func test_confidence_midBucket_whenBetween20And59Flights() async throws {
        let all = await BTSDataStore.shared.allRecordsKeyed()
        guard let (key, _) = all
            .filter({ $0.value.totalFlights >= 20 && $0.value.totalFlights < 60 })
            .first
        else {
            throw XCTSkip("Bundled BTS bundle has no record with totalFlights in 20…59")
        }

        let parts = key.split(separator: "_").map(String.init)
        guard parts.count == 4, let fn = Int(parts[1]) else {
            XCTFail("Unexpected BTS key shape: \(key)")
            return
        }

        let depart = date("2026-09-15T14:00:00-05:00") // Tuesday, non-peak
        let estimate = await service.estimate(
            carrier: parts[0],
            flightNumber: fn,
            origin: parts[2],
            dest: parts[3],
            date: depart,
            database: airportDatabase,
            liveSeats: nil
        )
        let result = try XCTUnwrap(estimate)
        XCTAssertEqual(result.confidence, 0.65, accuracy: 0.0001,
                       "totalFlights in 20…59 must map to 0.65 confidence")
    }

    /// sampleSize < 20 → 0.40 confidence.
    ///
    /// The bundled `bts_bundle.json` currently has no record with
    /// totalFlights < 20. We probe every record and run the assertion
    /// against the lowest-sample record if (and only if) it falls into
    /// the < 20 bucket; otherwise we XCTSkip with a note. Phase 2 may
    /// add real low-sample data and unfreeze this test.
    func test_confidence_lowBucket_whenFewerThan20Flights() async throws {
        let all = await BTSDataStore.shared.allRecordsKeyed()
        guard let (key, record) = all
            .filter({ $0.value.totalFlights < 20 })
            .min(by: { $0.value.totalFlights < $1.value.totalFlights })
        else {
            throw XCTSkip("Bundled BTS bundle has no record with totalFlights < 20; can't pin the 0.40 bucket against real data yet")
        }

        // Re-split the bundle key (CARRIER_FLIGHTNUM_ORIGIN_DEST) — the
        // format is fixed by BTSDataStore.makeKey.
        let parts = key.split(separator: "_").map(String.init)
        guard parts.count == 4, let fn = Int(parts[1]) else {
            XCTFail("Unexpected BTS key shape: \(key)")
            return
        }

        let depart = date("2026-09-15T14:00:00-05:00") // Tuesday, non-peak
        let estimate = await service.estimate(
            carrier: parts[0],
            flightNumber: fn,
            origin: parts[2],
            dest: parts[3],
            date: depart,
            database: airportDatabase,
            liveSeats: nil
        )
        let result = try XCTUnwrap(estimate)
        XCTAssertEqual(result.confidence, 0.40, accuracy: 0.0001,
                       "totalFlights < 20 (record \(key), n=\(record.totalFlights)) must map to 0.40 confidence")
    }

    // MARK: - 6. No record → nil

    /// When the BTS bundle has no matching key the service must return
    /// nil — callers hide the load-factor UI rather than guess.
    func test_estimate_returnsNil_whenNoMatchingBTSRecord() async {
        let depart = date("2026-09-15T14:00:00-05:00")
        let estimate = await service.estimate(
            carrier: "ZZ",
            flightNumber: 99999,
            origin: "AAA",
            dest: "BBB",
            date: depart,
            database: airportDatabase,
            liveSeats: nil
        )
        XCTAssertNil(estimate, "Nonsense carrier/route must yield nil, not a guessed estimate")
    }
}