import CryptoKit
import Foundation

// MARK: - Configuration

struct Config {
    let issuerId: String
    let keyId: String
    let privateKeyPath: String

    static func fromEnvironment() throws -> Config {
        guard let issuerId = ProcessInfo.processInfo.environment["AS_ISSUER_ID"] else {
            throw AppError.missingEnvVar("AS_ISSUER_ID")
        }
        guard let keyId = ProcessInfo.processInfo.environment["AS_KEY_ID"] else {
            throw AppError.missingEnvVar("AS_KEY_ID")
        }
        guard let keyPath = ProcessInfo.processInfo.environment["AS_PRIVATE_KEY_PATH"] else {
            throw AppError.missingEnvVar("AS_PRIVATE_KEY_PATH")
        }
        return Config(issuerId: issuerId, keyId: keyId, privateKeyPath: keyPath)
    }
}

// MARK: - Errors

enum AppError: LocalizedError {
    case missingEnvVar(String)
    case invalidPrivateKey
    case apiError(Int, String)
    case noSubscriptionFound(String)
    case noGroupFound(String)
    case missingUsage

    var errorDescription: String? {
        switch self {
        case .missingEnvVar(let name):
            return "Missing environment variable: \(name)"
        case .invalidPrivateKey:
            return "Failed to load private key from AS_PRIVATE_KEY_PATH"
        case .apiError(let status, let body):
            return "API error (HTTP \(status)): \(body)"
        case .noSubscriptionFound(let productId):
            return "No subscription found for product ID: \(productId)"
        case .noGroupFound(let name):
            return "No subscription group found matching: \(name)"
        case .missingUsage:
            return "Usage: swift run pricing <product-id>"
        }
    }
}

// MARK: - JWT Generation

func generateJWT(config: Config) throws -> String {
    let pemString = try String(contentsOfFile: config.privateKeyPath, encoding: .utf8)

    // Strip PEM headers and whitespace
    let base64Key = pemString
        .replacingOccurrences(of: "-----BEGIN PRIVATE KEY-----", with: "")
        .replacingOccurrences(of: "-----END PRIVATE KEY-----", with: "")
        .replacingOccurrences(of: "\n", with: "")
        .replacingOccurrences(of: "\r", with: "")
        .trimmingCharacters(in: .whitespaces)

    guard let keyData = Data(base64Encoded: base64Key) else {
        throw AppError.invalidPrivateKey
    }

    let privateKey = try P256.Signing.PrivateKey(derRepresentation: keyData)

    let now = Date()
    let header: [String: Any] = [
        "alg": "ES256",
        "kid": config.keyId,
        "typ": "JWT",
    ]
    let payload: [String: Any] = [
        "iss": config.issuerId,
        "iat": Int(now.timeIntervalSince1970),
        "exp": Int(now.addingTimeInterval(20 * 60).timeIntervalSince1970),
        "aud": "appstoreconnect-v1",
    ]

    let headerData = try JSONSerialization.data(withJSONObject: header)
    let payloadData = try JSONSerialization.data(withJSONObject: payload)

    let headerB64 = headerData.base64URLEncoded()
    let payloadB64 = payloadData.base64URLEncoded()

    let signingInput = "\(headerB64).\(payloadB64)"
    let signature = try privateKey.signature(for: Data(signingInput.utf8))
    let signatureB64 = signature.rawRepresentation.base64URLEncoded()

    return "\(headerB64).\(payloadB64).\(signatureB64)"
}

// MARK: - Base64URL Encoding

extension Data {
    func base64URLEncoded() -> String {
        base64EncodedString()
            .replacingOccurrences(of: "+", with: "-")
            .replacingOccurrences(of: "/", with: "_")
            .replacingOccurrences(of: "=", with: "")
    }
}

// MARK: - API Client

let baseURL = "https://api.appstoreconnect.apple.com/v1"

func apiRequest(path: String, token: String, retries: Int = 3) async throws -> (Data, HTTPURLResponse) {
    let urlString = path.hasPrefix("http") ? path : "\(baseURL)\(path)"
    guard let url = URL(string: urlString) else {
        throw URLError(.badURL)
    }

    var request = URLRequest(url: url)
    request.timeoutInterval = 30
    request.setValue("Bearer \(token)", forHTTPHeaderField: "Authorization")
    request.setValue("application/json", forHTTPHeaderField: "Content-Type")

    var lastError: Error?
    for attempt in 1...retries {
        do {
            let (data, response) = try await URLSession.shared.data(for: request)
            guard let httpResponse = response as? HTTPURLResponse else {
                throw URLError(.badServerResponse)
            }

            if httpResponse.statusCode == 429 || (httpResponse.statusCode >= 500 && httpResponse.statusCode < 600) {
                let delay = Double(attempt) * 2.0
                fputs("  Rate limited/server error (HTTP \(httpResponse.statusCode)), retrying in \(Int(delay))s...\n", stderr)
                try await Task.sleep(nanoseconds: UInt64(delay * 1_000_000_000))
                continue
            }

            if httpResponse.statusCode >= 400 {
                let body = String(data: data, encoding: .utf8) ?? "<no body>"
                throw AppError.apiError(httpResponse.statusCode, body)
            }

            return (data, httpResponse)
        } catch let error as AppError {
            throw error // don't retry client errors
        } catch {
            lastError = error
            if attempt < retries {
                let delay = Double(attempt) * 2.0
                fputs("  Request failed (\(error.localizedDescription)), retrying in \(Int(delay))s...\n", stderr)
                try await Task.sleep(nanoseconds: UInt64(delay * 1_000_000_000))
            }
        }
    }

    throw lastError ?? URLError(.timedOut)
}

// MARK: - API Response Models

struct APIResponse<T: Decodable>: Decodable {
    let data: [T]
    let included: [IncludedResource]?
    let links: PageLinks?
}

struct SingleResponse<T: Decodable>: Decodable {
    let data: T
}

struct PageLinks: Decodable {
    let next: String?
}

struct IncludedResource: Decodable {
    let type: String
    let id: String
    let attributes: IncludedAttributes?
}

struct IncludedAttributes: Decodable {
    // SubscriptionPricePoint attributes
    let customerPrice: String?
    let proceeds: String?
    let proceedsYear2: String?
    // Territory attributes
    let currency: String?
    // Shared
    let name: String?
}

struct SubscriptionGroup: Decodable {
    let type: String
    let id: String
    let attributes: SubscriptionGroupAttributes?
}

struct SubscriptionGroupAttributes: Decodable {
    let referenceName: String?
}

struct Subscription: Decodable {
    let type: String
    let id: String
    let attributes: SubscriptionAttributes
}

struct SubscriptionAttributes: Decodable {
    let productId: String?
    let name: String?
    let state: String?
}

struct SubscriptionPrice: Decodable {
    let type: String
    let id: String
    let attributes: SubscriptionPriceAttributes?
    let relationships: SubscriptionPriceRelationships?
}

struct SubscriptionPriceAttributes: Decodable {
    let startDate: String?
    let preserved: Bool?
}

struct SubscriptionPriceRelationships: Decodable {
    let subscriptionPricePoint: RelationshipData?
    let territory: RelationshipData?
}

struct RelationshipData: Decodable {
    let data: ResourceIdentifier?
}

struct ResourceIdentifier: Decodable {
    let type: String
    let id: String
}

// MARK: - Price Entry (for display)

struct PriceEntry {
    let territory: String
    let currency: String
    let customerPrice: String
    let proceeds: String
}

// MARK: - Fetch Logic

func findSubscription(appId: String, productId: String, token: String) async throws -> Subscription {
    let groupsPath = "/apps/\(appId)/subscriptionGroups?limit=200"
    let (groupsData, _) = try await apiRequest(path: groupsPath, token: token)
    let groupsResponse = try JSONDecoder().decode(APIResponse<SubscriptionGroup>.self, from: groupsData)

    if groupsResponse.data.isEmpty {
        throw AppError.noSubscriptionFound(productId)
    }

    for group in groupsResponse.data {
        let subsPath = "/subscriptionGroups/\(group.id)/subscriptions?limit=200"
        let (subsData, _) = try await apiRequest(path: subsPath, token: token)
        let subsResponse = try JSONDecoder().decode(APIResponse<Subscription>.self, from: subsData)

        if let match = subsResponse.data.first(where: { $0.attributes.productId == productId }) {
            return match
        }
    }

    throw AppError.noSubscriptionFound(productId)
}

func findSubscriptionsInGroup(appId: String, groupName: String, token: String) async throws -> (SubscriptionGroup, [Subscription]) {
    let groupsPath = "/apps/\(appId)/subscriptionGroups?limit=200"
    let (groupsData, _) = try await apiRequest(path: groupsPath, token: token)
    let groupsResponse = try JSONDecoder().decode(APIResponse<SubscriptionGroup>.self, from: groupsData)

    // Case-insensitive match on referenceName
    guard let group = groupsResponse.data.first(where: {
        $0.attributes?.referenceName?.lowercased() == groupName.lowercased()
    }) else {
        let available = groupsResponse.data.compactMap { $0.attributes?.referenceName }.joined(separator: ", ")
        fputs("Available groups: \(available)\n", stderr)
        throw AppError.noGroupFound(groupName)
    }

    let subsPath = "/subscriptionGroups/\(group.id)/subscriptions?limit=200"
    let (subsData, _) = try await apiRequest(path: subsPath, token: token)
    let subsResponse = try JSONDecoder().decode(APIResponse<Subscription>.self, from: subsData)

    return (group, subsResponse.data)
}

struct RawPriceEntry {
    let territory: String
    let currency: String
    let customerPrice: String
    let proceeds: String
    let startDate: Date? // nil = original price (no scheduled change)
}

private let dateFormatter: DateFormatter = {
    let df = DateFormatter()
    df.dateFormat = "yyyy-MM-dd"
    df.locale = Locale(identifier: "en_US_POSIX")
    df.timeZone = TimeZone(identifier: "UTC")
    return df
}()

func fetchPrices(subscriptionId: String, token: String) async throws -> [PriceEntry] {
    var rawEntries: [RawPriceEntry] = []
    var nextURL: String? =
        "/subscriptions/\(subscriptionId)/prices?include=subscriptionPricePoint,territory&limit=200"

    let today = Date()

    while let currentPath = nextURL {
        let (data, _) = try await apiRequest(path: currentPath, token: token)
        let response = try JSONDecoder().decode(APIResponse<SubscriptionPrice>.self, from: data)

        var pricePoints: [String: IncludedAttributes] = [:]
        var territories: [String: IncludedAttributes] = [:]

        for resource in response.included ?? [] {
            switch resource.type {
            case "subscriptionPricePoints":
                pricePoints[resource.id] = resource.attributes
            case "territories":
                territories[resource.id] = resource.attributes
            default:
                break
            }
        }

        for price in response.data {
            let pricePointId = price.relationships?.subscriptionPricePoint?.data?.id
            let territoryId = price.relationships?.territory?.data?.id

            let pp = pricePointId.flatMap { pricePoints[$0] }
            let terr = territoryId.flatMap { territories[$0] }

            let startDate = price.attributes?.startDate.flatMap { dateFormatter.date(from: $0) }

            // Skip prices scheduled for the future
            if let startDate, startDate > today {
                continue
            }

            rawEntries.append(RawPriceEntry(
                territory: terr?.name ?? territoryId ?? "Unknown",
                currency: terr?.currency ?? "???",
                customerPrice: pp?.customerPrice ?? "-",
                proceeds: pp?.proceeds ?? "-",
                startDate: startDate
            ))
        }

        nextURL = response.links?.next
    }

    // Keep only the most recent price per territory
    var latestByTerritory: [String: RawPriceEntry] = [:]
    for entry in rawEntries {
        if let existing = latestByTerritory[entry.territory] {
            let existingDate = existing.startDate ?? .distantPast
            let newDate = entry.startDate ?? .distantPast
            if newDate > existingDate {
                latestByTerritory[entry.territory] = entry
            }
        } else {
            latestByTerritory[entry.territory] = entry
        }
    }

    return latestByTerritory.values
        .map { PriceEntry(territory: $0.territory, currency: $0.currency, customerPrice: $0.customerPrice, proceeds: $0.proceeds) }
        .sorted { $0.territory < $1.territory }
}

// MARK: - Display

func padRight(_ s: String, _ width: Int) -> String {
    s.count >= width ? String(s.prefix(width)) : s + String(repeating: " ", count: width - s.count)
}

func padLeft(_ s: String, _ width: Int) -> String {
    s.count >= width ? s : String(repeating: " ", count: width - s.count) + s
}

func printPrices(_ entries: [PriceEntry], productId: String) {
    print("\nSubscription Prices for: \(productId)")
    print(String(repeating: "─", count: 72))
    print("\(padRight("Territory", 30))  \(padRight("Curr", 6))  \(padLeft("Price", 10))  \(padLeft("Proceeds", 10))")
    print(String(repeating: "─", count: 72))

    for entry in entries {
        print("\(padRight(String(entry.territory.prefix(30)), 30))  \(padRight(entry.currency, 6))  \(padLeft(entry.customerPrice, 10))  \(padLeft(entry.proceeds, 10))")
    }

    print(String(repeating: "─", count: 72))
    print("Total territories: \(entries.count)")
}

// MARK: - CSV Export

func csvEscape(_ field: String) -> String {
    if field.contains(",") || field.contains("\"") || field.contains("\n") {
        return "\"\(field.replacingOccurrences(of: "\"", with: "\"\""))\""
    }
    return field
}

func buildCSVRows(productId: String, entries: [PriceEntry]) -> [String] {
    entries.map { entry in
        [productId, entry.territory, entry.currency, entry.customerPrice, entry.proceeds]
            .map { csvEscape($0) }
            .joined(separator: ",")
    }
}

func sanitizeFilename(_ name: String) -> String {
    let allowed = CharacterSet.alphanumerics.union(CharacterSet(charactersIn: "-_"))
    return String(name.unicodeScalars.map { allowed.contains($0) ? Character($0) : Character("_") })
}

func writeCSV(rows: [String], name: String, to directory: String) throws {
    let fm = FileManager.default
    if !fm.fileExists(atPath: directory) {
        try fm.createDirectory(atPath: directory, withIntermediateDirectories: true)
    }
    let filename = sanitizeFilename(name) + ".csv"
    let path = (directory as NSString).appendingPathComponent(filename)
    let header = "Product ID,Territory,Currency,Customer Price,Proceeds"
    let content = ([header] + rows).joined(separator: "\n") + "\n"
    try content.write(toFile: path, atomically: true, encoding: .utf8)
    print("  CSV exported: \(path)")
}

// MARK: - Argument Parsing

func findArgValue(_ flag: String) -> String? {
    let args = CommandLine.arguments
    guard let idx = args.firstIndex(of: flag), idx + 1 < args.count else { return nil }
    return args[idx + 1]
}

func hasFlag(_ flag: String) -> Bool {
    CommandLine.arguments.contains(flag)
}

// Strip known flags from arguments to get positional args
func positionalArgs() -> [String] {
    let args = Array(CommandLine.arguments.dropFirst()) // drop executable
    let flagsWithValue = ["--group", "--csv"]
    var positional: [String] = []
    var i = 0
    while i < args.count {
        if flagsWithValue.contains(args[i]) {
            i += 2 // skip flag and its value
        } else if args[i].hasPrefix("--") {
            i += 1 // skip standalone flag
        } else {
            positional.append(args[i])
            i += 1
        }
    }
    return positional
}

// MARK: - Main

let pos = positionalArgs()

guard !pos.isEmpty else {
    fputs("Usage: swift run pricing <app-id> [product-id] [options]\n", stderr)
    fputs("       swift run pricing <app-id> --group <group-name> [--csv <file>]\n", stderr)
    fputs("\n", stderr)
    fputs("  app-id:      Numeric Apple ID (e.g. 6737844884)\n", stderr)
    fputs("  product-id:  Single subscription product ID\n", stderr)
    fputs("  --group:     Fetch all subscriptions in a named group\n", stderr)
    fputs("  --csv:       Export results to CSV files in this directory (one per subscription)\n", stderr)
    exit(1)
}

let appId = pos[0]
let csvDir = findArgValue("--csv")
let groupName = findArgValue("--group")
let singleProductId = pos.count >= 2 ? pos[1] : nil

do {
    let config = try Config.fromEnvironment()
    let token = try generateJWT(config: config)
    print("Authenticating with App Store Connect API...")

    if let groupName {
        let (group, subscriptions) = try await findSubscriptionsInGroup(
            appId: appId, groupName: groupName, token: token)

        print(
            "Group: \(group.attributes?.referenceName ?? group.id) — \(subscriptions.count) subscription(s)\n"
        )

        for (index, sub) in subscriptions.enumerated() {
            let name = sub.attributes.name ?? sub.attributes.productId ?? sub.id
            let pid = sub.attributes.productId ?? sub.id
            let state = sub.attributes.state ?? "unknown"
            print("[\(name)] (state: \(state))")

            let prices = try await fetchPrices(subscriptionId: sub.id, token: token)
            printPrices(prices, productId: pid)

            if let csvDir {
                let rows = buildCSVRows(productId: pid, entries: prices)
                try writeCSV(rows: rows, name: name, to: csvDir)
            }
            print("")

            // Brief pause between subscriptions to avoid rate limiting
            if index < subscriptions.count - 1 {
                try await Task.sleep(nanoseconds: 1_000_000_000)
            }
        }
    } else if let singleProductId {
        let subscription = try await findSubscription(
            appId: appId, productId: singleProductId, token: token)
        let name = subscription.attributes.name ?? singleProductId
        print(
            "Found subscription: \(name) (state: \(subscription.attributes.state ?? "unknown"))"
        )

        let prices = try await fetchPrices(subscriptionId: subscription.id, token: token)
        printPrices(prices, productId: singleProductId)

        if let csvDir {
            let rows = buildCSVRows(productId: singleProductId, entries: prices)
            try writeCSV(rows: rows, name: name, to: csvDir)
        }
    } else {
        fputs("Error: provide a product-id or --group <name>\n", stderr)
        exit(1)
    }
} catch {
    fputs("Error: \(error.localizedDescription)\n", stderr)
    exit(1)
}