//
//  OpenHaystack – Tracking personal Bluetooth devices via Apple's Find My network
//
//  Copyright © 2021 Secure Mobile Networking Lab (SEEMOO)
//  Copyright © 2021 The Open Wireless Link Project
//
//  SPDX-License-Identifier: AGPL-3.0-only
//

import Combine
import Foundation
import SwiftUI

class FindMyController: ObservableObject {
  static let shared = FindMyController()

  @Published var error: Error?
  @Published var devices = [FindMyDevice]()

  func loadPrivateKeys(
    from data: Data, with searchPartyToken: Data, completion: @escaping (Error?) -> Void
  ) {
    do {
      let devices = try PropertyListDecoder().decode([FindMyDevice].self, from: data)

      self.devices.append(contentsOf: devices)
      self.fetchReports(with: searchPartyToken, completion: completion)
    } catch {
      self.error = FindMyErrors.decodingPlistFailed(message: String(describing: error))
    }
  }

  func importReports(reports: [FindMyReport], and keys: Data, completion: @escaping () -> Void)
    throws
  {
    var devices = try PropertyListDecoder().decode([FindMyDevice].self, from: keys)

    // Decrypt the reports with the imported keys
    DispatchQueue.global(qos: .background).async {
      // Add the reports to the according device by finding the right key for the report
      for report in reports {

        guard
          let deviceIndex = devices.firstIndex(where: { (device) -> Bool in
            device.keys.contains { (key) -> Bool in
              key.hashedKey.base64EncodedString() == report.id
            }
          })
        else {
          print("No device found for id")
          continue
        }
        if var reports = devices[deviceIndex].reports {
          reports.append(report)
          devices[deviceIndex].reports = reports
        } else {
          devices[deviceIndex].reports = [report]
        }
      }
      self.devices = devices

      // Decrypt the reports
      self.decryptReports {
        self.exportDevices()
        DispatchQueue.main.async {
          completion()
        }
      }

    }
  }

  func importDevices(devices: Data) throws {
    var devices = try PropertyListDecoder().decode([FindMyDevice].self, from: devices)

    // Delete the decrypted reports
    for idx in devices.startIndex..<devices.endIndex {
      devices[idx].decryptedReports = nil
    }

    self.devices = devices

    // Decrypt reports again with additional information
    self.decryptReports {

    }
  }

  func fetchReports(with searchPartyToken: Data, completion: @escaping (Error?) -> Void) {

    DispatchQueue.global(qos: .background).async {
      let fetchReportGroup = DispatchGroup()

      let fetcher = ReportsFetcher()

      var devices = self.devices
      for deviceIndex in 0..<devices.count {
        fetchReportGroup.enter()
        devices[deviceIndex].reports = []

        // Only use the newest keys for testing
        let keys = devices[deviceIndex].keys

        let keyHashes = keys.map({ $0.hashedKey.base64EncodedString() })

        // 21 days
        let duration: Double = (24 * 60 * 60) * 21
        let startDate = Date() - duration

        fetcher.query(
          forHashes: keyHashes,
          start: startDate,
          duration: duration,
          searchPartyToken: searchPartyToken
        ) { jd in
          guard let jsonData = jd else {
            fetchReportGroup.leave()
            return
          }

          do {
            // Decode the report
            let report = try JSONDecoder().decode(FindMyReportResults.self, from: jsonData)
            devices[deviceIndex].reports = report.results

          } catch {
            print("Failed with error \(error)")
            devices[deviceIndex].reports = []
          }
          fetchReportGroup.leave()
        }

      }

      // Completion Handler
      fetchReportGroup.notify(queue: .main) {
        print("Finished loading the reports. Now decrypt them")

        // Export the reports to the desktop
        var reports = [FindMyReport]()
        for device in devices {
          for report in device.reports! {
            reports.append(report)
          }
        }

        #if EXPORT
          if let encoded = try? JSONEncoder().encode(reports) {
            let outputDirectory = FileManager.default.urls(
              for: .desktopDirectory, in: .userDomainMask
            ).first!
            try? encoded.write(to: outputDirectory.appendingPathComponent("reports.json"))
          }
        #endif

        DispatchQueue.main.async {
          self.devices = devices

          self.decryptReports {
            completion(nil)
          }

        }
      }
    }

  }

  func decryptReports(completion: () -> Void) {
    print("Decrypting reports")

    // Iterate over all devices
    for deviceIdx in 0..<devices.count {
      devices[deviceIdx].decryptedReports = []
      let device = devices[deviceIdx]

      // Map the keys in a dictionary for faster access
      guard let reports = device.reports else { continue }
      let keyMap = device.keys.reduce(
        into: [String: FindMyKey](), { $0[$1.hashedKey.base64EncodedString()] = $1 })

      let accessQueue = DispatchQueue(
        label: "threadSafeAccess",
        qos: .userInitiated,
        attributes: .concurrent,
        autoreleaseFrequency: .workItem, target: nil)
      var decryptedReports = [FindMyLocationReport](
        repeating:
          FindMyLocationReport(lat: 0, lng: 0, acc: 0, dP: Date(), t: Date(), c: 0),
        count: reports.count)
      DispatchQueue.concurrentPerform(iterations: reports.count) { (reportIdx) in
        let report = reports[reportIdx]
        guard let key = keyMap[report.id] else { return }
        do {
          // Decrypt the report
          let locationReport = try DecryptReports.decrypt(report: report, with: key)
          accessQueue.async(flags: .barrier) {
            decryptedReports[reportIdx] = locationReport
          }
        } catch {
          return
        }
      }

      accessQueue.sync {
        devices[deviceIdx].decryptedReports = decryptedReports
      }
    }

    completion()

  }

  func exportDevices() {

    if let encoded = try? PropertyListEncoder().encode(self.devices) {
      let outputDirectory = FileManager.default.urls(for: .desktopDirectory, in: .userDomainMask)
        .first!
      try? encoded.write(to: outputDirectory.appendingPathComponent("devices-\(Date()).plist"))
    }
  }

}

struct FindMyControllerKey: EnvironmentKey {
  static var defaultValue: FindMyController = .shared
}

extension EnvironmentValues {
  var findMyController: FindMyController {
    get { self[FindMyControllerKey.self] }
    set { self[FindMyControllerKey.self] = newValue }
  }
}

enum FindMyErrors: Error {
  case decodingPlistFailed(message: String)
}