import Foundation

// MARK: - KeySignature

/// Musical key signature defining the tonic and mode, and the resulting sharps/flats.
struct KeySignature: Equatable {
    let tonic: Pitch.NoteName
    let accidental: Accidental   // for the tonic itself (e.g., F# major, Bb minor)
    let mode: Mode

    enum Mode: String, CaseIterable {
        case major
        case minor
    }

    /// The sharps or flats in this key signature, in standard order.
    /// Positive count = sharps, negative = flats.
    /// Returns the note names that are sharped or flatted.
    var sharpsAndFlats: [Pitch.NoteName] {
        let info = KeySignature.keyInfo(tonic: tonic, accidental: accidental, mode: mode)
        return info.affectedNotes
    }

    /// Whether this key uses sharps (true) or flats (false).
    /// Keys with no sharps/flats return true by convention.
    var usesSharps: Bool {
        let info = KeySignature.keyInfo(tonic: tonic, accidental: accidental, mode: mode)
        return info.usesSharps
    }

    /// Number of sharps or flats in the key signature
    var count: Int {
        return sharpsAndFlats.count
    }

    /// Check if a given note name is sharped or flatted in this key
    func isAffected(_ noteName: Pitch.NoteName) -> Bool {
        return sharpsAndFlats.contains(noteName)
    }

    /// Determine what accidental to display for a given pitch in this key.
    /// Returns nil if no accidental needs to be shown (note matches the key signature).
    func displayAccidental(for pitch: Pitch) -> Accidental? {
        let isInKey = isAffected(pitch.noteName)

        if isInKey {
            // This note is sharped/flatted by the key signature
            let keyAccidental: Accidental = usesSharps ? .sharp : .flat
            if pitch.accidental == keyAccidental {
                // Matches key signature — no accidental needed
                return nil
            } else {
                // Different from key signature — show the actual accidental
                return pitch.accidental
            }
        } else {
            // This note is natural in the key signature
            if pitch.accidental == .natural {
                return nil  // Natural in a natural-key context — no display needed
            } else {
                return pitch.accidental  // Show the accidental
            }
        }
    }

    /// Parse from string like "C major", "B minor", "F# major", "Bb minor"
    static func parse(_ str: String) -> KeySignature? {
        let trimmed = str.trimmingCharacters(in: .whitespaces).lowercased()
        let parts = trimmed.split(separator: " ")
        guard parts.count == 2 else { return nil }

        let keyStr = String(parts[0])
        let modeStr = String(parts[1])

        guard let mode = Mode(rawValue: modeStr) else { return nil }

        // Parse the key name: first char is note, rest is accidental
        guard !keyStr.isEmpty else { return nil }
        let noteChar = String(keyStr.prefix(1)).uppercased()
        guard let noteName = Pitch.NoteName(rawValue: noteChar) else { return nil }

        let accStr = String(keyStr.dropFirst())
        guard let acc = Accidental.parse(accStr) else { return nil }

        return KeySignature(tonic: noteName, accidental: acc, mode: mode)
    }

    // MARK: - Key signature lookup

    private struct KeyInfo {
        let usesSharps: Bool
        let affectedNotes: [Pitch.NoteName]
    }

    /// Order of sharps: F C G D A E B
    private static let sharpOrder: [Pitch.NoteName] = [.F, .C, .G, .D, .A, .E, .B]
    /// Order of flats: B E A D G C F
    private static let flatOrder: [Pitch.NoteName] = [.B, .E, .A, .D, .G, .C, .F]

    /// Lookup table for major keys by (tonic + accidental) -> number of sharps (positive) or flats (negative)
    private static func majorKeySharps(tonic: Pitch.NoteName, accidental: Accidental) -> Int? {
        // Standard circle of fifths
        switch (tonic, accidental) {
        case (.C, .natural): return 0
        case (.G, .natural): return 1
        case (.D, .natural): return 2
        case (.A, .natural): return 3
        case (.E, .natural): return 4
        case (.B, .natural): return 5
        case (.F, .sharp):   return 6
        case (.C, .sharp):   return 7
        case (.F, .natural): return -1
        case (.B, .flat):    return -2
        case (.E, .flat):    return -3
        case (.A, .flat):    return -4
        case (.D, .flat):    return -5
        case (.G, .flat):    return -6
        case (.C, .flat):    return -7
        default: return nil
        }
    }

    /// Get the relative major tonic for a minor key
    private static func relativeMajor(tonic: Pitch.NoteName, accidental: Accidental) -> (Pitch.NoteName, Accidental)? {
        // Minor key is 3 semitones below its relative major
        // Map minor tonic to relative major
        switch (tonic, accidental) {
        case (.A, .natural): return (.C, .natural)
        case (.E, .natural): return (.G, .natural)
        case (.B, .natural): return (.D, .natural)
        case (.F, .sharp):   return (.A, .natural)
        case (.C, .sharp):   return (.E, .natural)
        case (.G, .sharp):   return (.B, .natural)
        case (.D, .sharp):   return (.F, .sharp)
        case (.A, .sharp):   return (.C, .sharp)
        case (.D, .natural): return (.F, .natural)
        case (.G, .natural): return (.B, .flat)
        case (.C, .natural): return (.E, .flat)
        case (.F, .natural): return (.A, .flat)
        case (.B, .flat):    return (.D, .flat)
        case (.E, .flat):    return (.G, .flat)
        case (.A, .flat):    return (.C, .flat)
        default: return nil
        }
    }

    private static func keyInfo(tonic: Pitch.NoteName, accidental: Accidental, mode: Mode) -> KeyInfo {
        let sharpCount: Int

        switch mode {
        case .major:
            sharpCount = majorKeySharps(tonic: tonic, accidental: accidental) ?? 0
        case .minor:
            if let (relMajTonic, relMajAcc) = relativeMajor(tonic: tonic, accidental: accidental) {
                sharpCount = majorKeySharps(tonic: relMajTonic, accidental: relMajAcc) ?? 0
            } else {
                sharpCount = 0
            }
        }

        if sharpCount >= 0 {
            let notes = Array(sharpOrder.prefix(sharpCount))
            return KeyInfo(usesSharps: true, affectedNotes: notes)
        } else {
            let notes = Array(flatOrder.prefix(-sharpCount))
            return KeyInfo(usesSharps: false, affectedNotes: notes)
        }
    }
}

extension KeySignature: CustomStringConvertible {
    var description: String {
        let accStr: String
        switch accidental {
        case .natural:     accStr = ""
        case .sharp:       accStr = "#"
        case .flat:        accStr = "b"
        case .doubleSharp: accStr = "##"
        case .doubleFlat:  accStr = "bb"
        }
        return "\(tonic.rawValue)\(accStr) \(mode.rawValue)"
    }
}