 import UIKit
 
 public protocol SplooshSearchDelegate: class {
-    func foundPeersChanged(foundPeers: [String])
-    
-    func invitationReceived(peer: String, handleInvitation: (Bool) -> Void)
-    
-    func connectedPeersChanged(peers: [String])
+    func foundHostsChanged(foundHosts: [MCPeerID])
+    
+    func invitationToConnectReceived(peer: MCPeerID, handleInvitation: (Bool) -> Void)
+    
+    func connectionsChanged(peers: [MCPeerID])
+    
+    func connectedToPeer(peer: MCPeerID)
+    
+    func handleDataPacket(data: NSData, peerID: MCPeerID)
 }
 
 public class SplooshSearch: NSObject {
+    var autoAcceptGuests = true
+    
+    static let defaultTimeout: NSTimeInterval = 30
     private var serviceAdvertiser: MCNearbyServiceAdvertiser?
     private var serviceBrowser: MCNearbyServiceBrowser?
-    private var isSearching = false
-    private var foundPeers = [String]()
+    private var foundHosts = [MCPeerID]()
     
     public weak var delegate: SplooshSearchDelegate?
     

        stopPeerSearch()
    }
    
    public func startPeerSearch() {
        isSearching = true
        if let advertiser = serviceAdvertiser {
            advertiser.startAdvertisingPeer()
        } else {
            serviceAdvertiser = MCNearbyServiceAdvertiser(peer: myPeerId,
                discoveryInfo: nil, serviceType: serviceType)
            serviceBrowser = MCNearbyServiceBrowser(peer: myPeerId, serviceType: serviceType)
            
            serviceAdvertiser?.delegate = self
            serviceBrowser?.delegate = self
            
            serviceAdvertiser?.startAdvertisingPeer()
            serviceBrowser?.startBrowsingForPeers()
        }
    }
    
    public func cancelPeerSearch() {
        stopPeerSearch()
        session.disconnect()
    }
    
    public func finishPeerSearch() -> MCSession? {
        stopPeerSearch()
        return session
    }
    
    public func stopPeerSearch() {
        if let advertiser = serviceAdvertiser {
            advertiser.stopAdvertisingPeer()
        }
        if let browser = serviceBrowser {
            browser.stopBrowsingForPeers()
        }
        isSearching = false
    }
    
    public func connectToPeer(peer: String) {
        let peerId = MCPeerID(displayName: peer)
        serviceBrowser?.invitePeer(peerId, toSession: session, withContext: nil, timeout: 30)
    }
}

extension SplooshSearch: MCNearbyServiceAdvertiserDelegate {
    // Invitation is received from peer
    public func advertiser(advertiser: MCNearbyServiceAdvertiser,
        didReceiveInvitationFromPeer peerID: MCPeerID,
        withContext context: NSData?, invitationHandler: (Bool, MCSession) -> Void) {
            
            delegate?.invitationReceived(peerID.displayName) {
                (approved) in
                invitationHandler(approved, self.session)
            }
    }
}

    // Peer is found in browser
    public func browser(browser: MCNearbyServiceBrowser, foundPeer peerID: MCPeerID,
        withDiscoveryInfo info: [String : String]?) {
            NSLog("%@", "foundPeer: \(peerID)")
            
            guard let discoveryInfo = info else {
                return
            }
            guard discoveryInfo["peerType"] == "host" else {
                return
            }
            
            NSLog("%@", "invitePeer: \(peerID)")
            foundHosts.append(peerID)
            delegate?.foundHostsChanged(foundHosts)
    }
    
    // Peer is lost in browser
    public func browser(browser: MCNearbyServiceBrowser, lostPeer peerID: MCPeerID) {
        guard foundHosts.contains(peerID) else {
            return
        }
        
        let index = foundHosts.indexOf(peerID)!
        foundHosts.removeAtIndex(index)
        delegate?.foundHostsChanged(foundHosts)
    }
}

