import Foundation
import Testing
@testable import IrdP2P
@testable import IrdP2PSignaling
import IrdP2PCrypto

/// An in-process relay with the `/ws` hub semantics: one join per socket, rooms
/// namespaced by id, `from` stamped by the relay, unknown recipients dropped,
/// rejoin-with-same-id replaces the old socket, `ping` answered with `pong`.
final class FakeRelay: WebSocketTransport, @unchecked Sendable {
    final class Socket: WebSocketConnection, @unchecked Sendable {
        let relay: FakeRelay
        let id = UUID()
        var room: RoomId?
        var peer: PeerId?
        private let inbound: AsyncStream<WebSocketMessage>
        private let feed: AsyncStream<WebSocketMessage>.Continuation
        private var iterator: AsyncStream<WebSocketMessage>.AsyncIterator?
        private(set) var closed = false

        init(relay: FakeRelay) {
            self.relay = relay
            (inbound, feed) = AsyncStream<WebSocketMessage>.makeStream(bufferingPolicy: .unbounded)
        }

        func send(text: String) async throws {
            guard !closed else { throw P2PError.closed }
            relay.handle(text, from: self)
        }

        func receive() async throws -> WebSocketMessage {
            if iterator == nil { iterator = inbound.makeAsyncIterator() }
            if let m = await iterator!.next() { return m }
            throw P2PError.closed
        }

        func close(code: Int, reason: String?) async {
            relay.disconnect(self)
        }

        func deliver(_ text: String) { feed.yield(.text(text)) }
        func terminate() { closed = true; feed.finish() }
    }

    private let lock = NSLock()
    private var rooms: [String: [PeerId: Socket]] = [:]
    private(set) var frames: [(from: PeerId?, text: String)] = []
    var maxRoomPeers = 16
    /// When set, every forwarded `signal` frame is also delivered a second time (replay attack).
    var duplicateSignals = false
    /// Drops the next `n` signal frames.
    var dropNextSignals = 0

    func connect(_ request: URLRequest) async throws -> any WebSocketConnection {
        Socket(relay: self)
    }

    fileprivate func handle(_ text: String, from socket: Socket) {
        guard let v = try? JSONValue.parse(text), let t = v["t"]?.stringValue else { return }
        lock.withLock { frames.append((socket.peer, text)) }
        switch t {
        case "join":
            guard socket.room == nil, let room = try? RoomId(v["room"]?.stringValue ?? ""), let peer = try? PeerId(v["peer"]?.stringValue ?? "") else {
                socket.deliver(#"{"t":"error","code":"bad-message"}"#); return
            }
            let (replaced, others): (Socket?, [Socket]) = lock.withLock {
                var members = rooms[room.string] ?? [:]
                let old = members[peer]
                if old == nil, members.count >= maxRoomPeers { return (nil, []) }
                members[peer] = socket
                rooms[room.string] = members
                return (old, members.filter { $0.key != peer }.map(\.value))
            }
            if !(lock.withLock { rooms[room.string]?[peer] === socket }) {
                socket.deliver(#"{"t":"error","code":"room-full"}"#); return
            }
            socket.room = room
            socket.peer = peer
            replaced?.terminate()
            socket.deliver(#"{"t":"peers","peers":[\#(others.compactMap { $0.peer?.hex }.map { "\"\($0)\"" }.joined(separator: ","))]}"#)
            for o in others { o.deliver(#"{"t":"peer-joined","peer":"\#(peer.hex)"}"#) }   // the hub re-announces a replacing connection too
        case "signal":
            guard let room = socket.room, let me = socket.peer else { socket.deliver(#"{"t":"error","code":"not-joined"}"#); return }
            guard let to = try? PeerId(v["to"]?.stringValue ?? ""), let payload = v["payload"] else { socket.deliver(#"{"t":"error","code":"bad-message"}"#); return }
            let (target, dup): (Socket?, Bool) = lock.withLock {
                if dropNextSignals > 0 { dropNextSignals -= 1; return (nil, false) }
                return (rooms[room.string]?[to], duplicateSignals)
            }
            let frame = JSONValue.object(JSONObject([("t", .string("signal")), ("from", .string(me.hex)), ("payload", payload)])).serialized
            target?.deliver(frame)
            if dup { target?.deliver(frame) }
        case "leave": disconnect(socket)
        case "ping": socket.deliver(#"{"t":"pong"}"#)
        default: socket.deliver(#"{"t":"error","code":"bad-message"}"#)
        }
    }

    fileprivate func disconnect(_ socket: Socket) {
        guard let room = socket.room, let peer = socket.peer else { socket.terminate(); return }
        let others: [Socket] = lock.withLock {
            guard rooms[room.string]?[peer] === socket else { return [] }
            rooms[room.string]?[peer] = nil
            return rooms[room.string]?.values.map { $0 } ?? []
        }
        socket.terminate()
        for o in others { o.deliver(#"{"t":"peer-left","peer":"\#(peer.hex)"}"#) }
    }

    /// Kills a peer's socket without a leave (network loss).
    func dropSocket(of peer: PeerId) {
        let s: Socket? = lock.withLock { rooms.values.compactMap { $0[peer] }.first }
        if let s { disconnect(s) }
    }

    func members(of room: RoomId) -> [PeerId] { lock.withLock { Array(rooms[room.string]?.keys ?? [:].keys).sorted() } }
    var signalFrameCount: Int { lock.withLock { frames.filter { $0.text.hasPrefix(#"{"t":"signal""#) }.count } }
}

/// Collects a room's events and lets tests wait for a matching one.
final class Collector: @unchecked Sendable {
    private let lock = NSLock()
    private var events: [RoomEvent] = []
    private var task: Task<Void, Never>?

    init(_ room: Room) {
        let stream = room.events
        task = Task { [weak self] in
            for await e in stream { self?.record(e) }
        }
    }

    private func record(_ e: RoomEvent) { lock.withLock { events.append(e) } }
    var all: [RoomEvent] { lock.withLock { events } }
    func contains(_ e: RoomEvent) -> Bool { all.contains(e) }
    func count(_ e: RoomEvent) -> Int { all.filter { $0 == e }.count }

    /// Waits (real time) until an event satisfying `predicate` has been recorded.
    @discardableResult
    func waitFor(timeout: Duration = .seconds(5), _ predicate: @escaping (RoomEvent) -> Bool) async -> RoomEvent? {
        let deadline = ContinuousClock.now + timeout
        while ContinuousClock.now < deadline {
            if let e = all.first(where: predicate) { return e }
            try? await Task.sleep(for: .milliseconds(5))
        }
        return all.first(where: predicate)
    }

    func waitFor(_ event: RoomEvent, timeout: Duration = .seconds(5)) async -> Bool {
        await waitFor(timeout: timeout) { $0 == event } != nil
    }

    /// Waits until `n` events match.
    func waitForCount(_ n: Int, timeout: Duration = .seconds(5), _ predicate: @escaping (RoomEvent) -> Bool) async -> Bool {
        let deadline = ContinuousClock.now + timeout
        while ContinuousClock.now < deadline {
            if all.filter(predicate).count >= n { return true }
            try? await Task.sleep(for: .milliseconds(5))
        }
        return all.filter(predicate).count >= n
    }

    /// Records the full event log of both rooms plus the relay frames as a test issue.
    static func dump(_ label: String, _ h: Harness, _ collectors: [(String, Collector)]) {
        var text = "\(label)\n"
        for (name, c) in collectors { text += "\(name): \(c.all)\n" }
        text += "relay frames: \(h.relay.frames.map { "\($0.from?.hex.prefix(4) ?? "-"):\($0.text.prefix(60))" })\n"
        let room = String(RoomCrypto.roomId(secret: h.secret, appSalt: h.appSalt).string.prefix(6))
        text += "trace:\n" + h.trace.lines(forRoom: room).joined(separator: "\n")
        Issue.record(Comment(rawValue: text))
    }

    /// Messages received so far, in order.
    var messages: [(PeerId, Message)] { all.compactMap { if case .message(let f, let m) = $0 { return (f, m) } else { return nil } } }
}

/// A test clock shared with the signaling tests' semantics (manual advance, cancellation-aware).
final class TestClock: P2PClock, @unchecked Sendable {
    private let lock = NSLock()
    private var time: Duration = .zero
    private var waiters: [(id: UUID, deadline: Duration, c: CheckedContinuation<Void, any Error>)] = []

    var now: Duration { lock.withLock { time } }

    func sleep(for duration: Duration) async throws {
        let id = UUID()
        try await withTaskCancellationHandler {
            try await withCheckedThrowingContinuation { (c: CheckedContinuation<Void, any Error>) in
                // a task cancelled before it got here already ran onCancel (finding nothing): it must not
                // register, or it would linger in `sleeperCount` until the clock passes its deadline
                let immediate: Result<Void, any Error>? = lock.withLock {
                    if Task.isCancelled { return .failure(CancellationError()) }
                    if duration <= .zero { return .success(()) }
                    waiters.append((id, time + duration, c))
                    return nil
                }
                if let immediate { c.resume(with: immediate) }
            }
        } onCancel: {
            let cancelled: CheckedContinuation<Void, any Error>? = lock.withLock {
                guard let i = waiters.firstIndex(where: { $0.id == id }) else { return nil }
                return waiters.remove(at: i).c
            }
            cancelled?.resume(throwing: CancellationError())
        }
    }

    func advance(by d: Duration) {
        let due: [CheckedContinuation<Void, any Error>] = lock.withLock {
            time += d
            let fired = waiters.filter { $0.deadline <= time }.map(\.c)
            waiters.removeAll { $0.deadline <= time }
            return fired
        }
        for c in due { c.resume() }
    }

    var sleeperCount: Int { lock.withLock { waiters.count } }
}

/// Answers the descriptor fetch with a fixed descriptor and everything else with 404.
final class DescriptorHTTP: HTTPTransport, @unchecked Sendable {
    let body: String
    init(host: String = "app.example.test", mailbox: Bool = true, openMint: Bool = true) {
        let mint = openMint ? #","mint_endpoint":"https://\#(host)/mailbox-credentials","mint_ip_per_min":10,"mint_app_per_min":600"# : ""
        let mbx = mailbox ? #","mailbox":{"endpoint":"https://\#(host)/mailbox/v1","ttl_hours":24,"max_mb":256,"max_message_kb":64,"open_mint":\#(openMint)\#(mint)}"# : ""
        body = #"{"protocol":"ap2p-signal","version":"1","endpoints":{"ws":"wss://\#(host)/ws","turn_credentials":"https://\#(host)/turn-credentials"}\#(mbx),"limits":{"max_rooms":500,"max_room_peers":16,"max_conns_per_ip":20,"max_bytes_per_sec":262144,"messages_per_sec":30,"message_burst":60,"max_frame_bytes":65536}}"#
    }
    func send(_ request: URLRequest) async throws -> HTTPResponse {
        if request.url?.path.hasSuffix("/.well-known/ap2p") == true { return HTTPResponse(status: 200, body: Data(body.utf8)) }
        return HTTPResponse(status: 404, body: Data(#"{"error":"not_found"}"#.utf8))
    }
}

/// Collects RoomTrace lines for the dumps (shared across parallel tests; lines are prefixed by peer id).
final class TraceSink: @unchecked Sendable {
    static let shared = TraceSink()
    private let lock = NSLock()
    private var lines: [String] = []
    init() { RoomTrace.sink = { [weak self] line in self?.append(line) } }
    private func append(_ l: String) { lock.withLock { lines.append(l); if lines.count > 20000 { lines.removeFirst(5000) } } }
    func lines(forRoom prefix: String) -> [String] { lock.withLock { lines.filter { $0.hasPrefix(prefix) } } }
}

/// Two apps sharing one fake relay, one in-memory WebRTC network and one test clock.
struct Harness {
    let trace = TraceSink.shared
    let relay = FakeRelay()
    let webrtc = InMemoryPeerConnectionFactory()
    let clock = TestClock()
    let secret = RoomSecret.random()
    let appSalt = "ird-swift-tests/room"

    func app(preferDataChannel: Bool = true, http: (any HTTPTransport)? = nil, peerConnections: (any PeerConnectionFactory)? = nil) -> P2PApp {
        var cfg = P2PAppConfiguration(host: "app.example.test", appSalt: appSalt, origin: "https://o.example")
        cfg.http = http ?? DescriptorHTTP()
        cfg.webSocket = relay
        cfg.clock = clock
        cfg.peerConnections = peerConnections ?? webrtc
        cfg.observeAppLifecycle = false
        return P2PApp(configuration: cfg)
    }

    func options(preferDataChannel: Bool = true) -> RoomOptions {
        var o = RoomOptions()
        o.preferDataChannel = preferDataChannel
        o.pacing = false
        return o
    }

    func join(peer: PeerId, identity: Identity? = nil, preferDataChannel: Bool = true, secret: RoomSecret? = nil) async throws -> (Room, Collector) {
        let room = try await app().join(secret: secret ?? self.secret, identity: identity, peerId: peer, options: options(preferDataChannel: preferDataChannel))
        return (room, Collector(room))
    }

    /// Joins with explicit options through a factory that records what the room asks of it.
    func join(peer: PeerId, options: RoomOptions) async throws -> (Room, Collector, RecordingFactory) {
        let factory = RecordingFactory(webrtc)
        let room = try await app(peerConnections: factory).join(secret: secret, peerId: peer, options: options)
        return (room, Collector(room), factory)
    }
}

/// Counts the peer connections one room creates on the shared in-memory network, and how it asked for them.
final class RecordingFactory: PeerConnectionFactory, @unchecked Sendable {
    private let inner: InMemoryPeerConnectionFactory
    private let lock = NSLock()
    private var made = 0
    private var flags: [Bool] = []

    init(_ inner: InMemoryPeerConnectionFactory) { self.inner = inner }

    var count: Int { lock.withLock { made } }
    /// The `continualGathering` of every request, in order.
    var continualGathering: [Bool] { lock.withLock { flags } }

    func makePeerConnection(ice: IceConfiguration) async throws -> any PeerConnectionHandle {
        lock.withLock { made += 1 }
        return try await inner.makePeerConnection(ice: ice)
    }

    func makePeerConnection(ice: IceConfiguration, continualGathering: Bool) async throws -> any PeerConnectionHandle {
        lock.withLock { flags.append(continualGathering) }
        return try await makePeerConnection(ice: ice)
    }
}

let lowPeer = try! PeerId("00000000000000000000000000000001")
let highPeer = try! PeerId("0f1e2d3c4b5a69788796a5b4c3d2e1f0")
let thirdPeer = try! PeerId("ffffffffffffffffffffffffffffff00")