import Foundation
import os

/// Base protocol for all Rosetta binary packets.
protocol Packet {
    static var packetId: Int { get }
    func write(to stream: Stream)
    mutating func read(from stream: Stream) throws
}

// MARK: - Packet Registry

enum PacketRegistry {

    private static let logger = Logger(subsystem: "com.rosetta.messenger", category: "PacketRegistry")

    /// All known packet factories, keyed by packet ID.
    private static let factories: [Int: () -> any Packet] = [
        0x00: { PacketHandshake() },
        0x01: { PacketUserInfo() },
        0x02: { PacketResult() },
        0x03: { PacketSearch() },
        0x04: { PacketOnlineSubscribe() },
        0x05: { PacketOnlineState() },
        0x06: { PacketMessage() },
        0x07: { PacketRead() },
        0x08: { PacketDelivery() },
        0x09: { PacketDeviceNew() },
        0x0A: { PacketRequestUpdate() },
        0x0B: { PacketTyping() },
        0x0F: { PacketRequestTransport() },
        0x10: { PacketPushNotification() },
        0x11: { PacketCreateGroup() },
        0x12: { PacketGroupInfo() },
        0x13: { PacketGroupInviteInfo() },
        0x14: { PacketGroupJoin() },
        0x15: { PacketGroupLeave() },
        0x16: { PacketGroupBan() },
        0x17: { PacketDeviceList() },
        0x18: { PacketDeviceResolve() },
        0x19: { PacketSync() },
        0x1A: { PacketSignalPeer() },
        0x1B: { PacketWebRTC() },
        0x1C: { PacketIceServers() },
    ]

    /// Deserializes a packet from raw binary data.
    static func decode(from data: Data) -> (packetId: Int, packet: any Packet)? {
        guard data.count >= 2 else {
            logger.error("Rejecting packet: too short payload size=\(data.count)")
            return nil
        }

        let stream = Stream(data: data)
        let packetId = stream.readInt16()

        guard let factory = factories[packetId] else {
            let packetHex = String(format: "0x%02X", packetId)
            logger.warning("Rejecting packet: unknown packetId=\(packetHex) size=\(data.count)")
            return nil
        }

        var packet = factory()
        do {
            try packet.read(from: stream)
        } catch {
            let packetHex = String(format: "0x%02X", packetId)
            logger.error("Rejecting packet: parse failure packetId=\(packetHex) size=\(data.count) error=\(String(describing: error))")
            return nil
        }
        return (packetId, packet)
    }

    /// Serializes a packet to raw binary data (including the 2-byte packet ID header).
    static func encode(_ packet: any Packet) -> Data {
        let stream = Stream()
        stream.writeInt16(type(of: packet).packetId)
        packet.write(to: stream)
        return stream.toData()
    }
}

// MARK: - Attachment Types

enum AttachmentType: Int, Codable, Sendable {
    case image = 0
    case messages = 1
    case file = 2
    case avatar = 3
    case call = 4
    case voice = 5

    /// Android parity: `fromInt() ?: UNKNOWN`. Fallback to `.image` for unknown values
    /// so a single unknown type doesn't crash the entire [MessageAttachment] array decode.
    init(from decoder: Decoder) throws {
        let raw = try decoder.singleValueContainer().decode(Int.self)
        self = AttachmentType(rawValue: raw) ?? .image
    }
}

struct MessageAttachment: Codable, Equatable {
    var id: String = ""
    var preview: String = ""
    var blob: String = ""
    var type: AttachmentType = .image
    var transportTag: String = ""
    var transportServer: String = ""

    /// Effective download tag: prefers per-attachment `transportTag`, falls back to preview-embedded tag.
    /// Backward compat: old messages (before multi-transport) have empty `transportTag` but tag in preview.
    var effectiveDownloadTag: String {
        if !transportTag.isEmpty { return transportTag }
        return AttachmentPreviewCodec.downloadTag(from: preview)
    }

    // Explicit memberwise init (preserves all existing call sites).
    init(id: String = "", preview: String = "", blob: String = "",
         type: AttachmentType = .image,
         transportTag: String = "", transportServer: String = "") {
        self.id = id
        self.preview = preview
        self.blob = blob
        self.type = type
        self.transportTag = transportTag
        self.transportServer = transportServer
    }

    // Custom Decodable — `decodeIfPresent` for backward compat with old DB JSON
    // that doesn't have transportTag/transportServer fields.
    init(from decoder: Decoder) throws {
        let container = try decoder.container(keyedBy: CodingKeys.self)
        id = try container.decodeIfPresent(String.self, forKey: .id) ?? ""
        preview = try container.decodeIfPresent(String.self, forKey: .preview) ?? ""
        blob = try container.decodeIfPresent(String.self, forKey: .blob) ?? ""
        type = try container.decodeIfPresent(AttachmentType.self, forKey: .type) ?? .image
        transportTag = try container.decodeIfPresent(String.self, forKey: .transportTag) ?? ""
        transportServer = try container.decodeIfPresent(String.self, forKey: .transportServer) ?? ""
    }
}

// MARK: - Reply Message Data (desktop parity)

/// JSON structure for reply/forward message blobs (AttachmentType.messages).
/// Desktop: `MessageReply` in `useReplyMessages.ts`.
/// Android: `ReplyMessage` in `ChatViewModel.kt`.
struct ReplyMessageData: Codable {
    let message_id: String
    let publicKey: String
    let message: String
    let timestamp: Int64
    let attachments: [ReplyAttachmentData]
    /// Desktop commit `aaa4b42`: hex-encoded plainKeyAndNonce of the ORIGINAL message.
    /// Allows recipients to decrypt forwarded attachments without re-upload.
    let chacha_key_plain: String
    /// Android parity: explicit forward flag. Without it, single-message forwards
    /// in groups are misidentified as replies (third-party heuristic is disabled for groups).
    let forwarded: Bool
    /// Android parity: sender display name. Avoids server lookup on Android.
    let senderName: String

    init(message_id: String, publicKey: String, message: String,
         timestamp: Int64, attachments: [ReplyAttachmentData],
         chacha_key_plain: String = "",
         forwarded: Bool = false,
         senderName: String = "") {
        self.message_id = message_id
        self.publicKey = publicKey
        self.message = message
        self.timestamp = timestamp
        self.attachments = attachments
        self.chacha_key_plain = chacha_key_plain
        self.forwarded = forwarded
        self.senderName = senderName
    }

    // Custom Decodable — `decodeIfPresent` for backward compat with old reply blobs
    // that don't have `chacha_key_plain`, `forwarded`, or `senderName`.
    init(from decoder: Decoder) throws {
        let c = try decoder.container(keyedBy: CodingKeys.self)
        message_id = try c.decode(String.self, forKey: .message_id)
        publicKey = try c.decode(String.self, forKey: .publicKey)
        message = try c.decode(String.self, forKey: .message)
        timestamp = try c.decode(Int64.self, forKey: .timestamp)
        attachments = try c.decode([ReplyAttachmentData].self, forKey: .attachments)
        chacha_key_plain = try c.decodeIfPresent(String.self, forKey: .chacha_key_plain) ?? ""
        forwarded = try c.decodeIfPresent(Bool.self, forKey: .forwarded) ?? false
        senderName = try c.decodeIfPresent(String.self, forKey: .senderName) ?? ""
    }
}

/// Transport info inside a reply/forward attachment.
/// Desktop: nested `transport: { transport_tag, transport_server }` in `Attachment`.
struct ReplyAttachmentTransport: Codable {
    let transport_tag: String
    let transport_server: String

    init(transport_tag: String = "", transport_server: String = "") {
        self.transport_tag = transport_tag
        self.transport_server = transport_server
    }
}

/// Attachment inside a reply/forward blob.
/// Desktop: `Attachment` interface in `packet.message.ts`.
struct ReplyAttachmentData: Codable {
    let id: String
    let type: Int
    let preview: String
    let blob: String
    /// Desktop commit `aaa4b42`: per-attachment transport for forwarded attachments.
    let transport: ReplyAttachmentTransport

    init(id: String, type: Int, preview: String, blob: String,
         transport: ReplyAttachmentTransport = ReplyAttachmentTransport()) {
        self.id = id
        self.type = type
        self.preview = preview
        self.blob = blob
        self.transport = transport
    }

    // Custom Decodable — `decodeIfPresent` for backward compat with old reply blobs
    // that don't have `transport`.
    init(from decoder: Decoder) throws {
        let c = try decoder.container(keyedBy: CodingKeys.self)
        id = try c.decode(String.self, forKey: .id)
        type = try c.decode(Int.self, forKey: .type)
        preview = try c.decode(String.self, forKey: .preview)
        blob = try c.decode(String.self, forKey: .blob)
        transport = try c.decodeIfPresent(ReplyAttachmentTransport.self, forKey: .transport)
            ?? ReplyAttachmentTransport()
    }
}

extension ReplyAttachmentData {
    /// CDN download tag — try per-attachment transport first, fallback to preview-embedded tag.
    var effectiveDownloadTag: String {
        if !transport.transport_tag.isEmpty { return transport.transport_tag }
        return AttachmentPreviewCodec.downloadTag(from: preview)
    }

    var attachmentType: AttachmentType? {
        AttachmentType(rawValue: type)
    }
}

// MARK: - Search User

struct SearchUser {
    var username: String = ""
    var title: String = ""
    var publicKey: String = ""
    var verified: Int = 0
    var online: Int = 0
}

// MARK: - Update Server Packet (0x0A)

/// Request/response packet for update server discovery.
/// Mirrors Android/Desktop `PacketRequestUpdate`.
struct PacketRequestUpdate: Packet {
    static let packetId = 0x0A

    var updateServer: String = ""

    func write(to stream: Stream) {
        stream.writeString(updateServer)
    }

    mutating func read(from stream: Stream) {
        updateServer = stream.readString()
    }
}

// MARK: - Group Packets (0x11...0x16)

enum GroupStatus: Int, Codable {
    case joined = 0
    case invalid = 1
    case notJoined = 2
    case banned = 3

    init(value: Int) {
        self = GroupStatus(rawValue: value) ?? .notJoined
    }
}

struct PacketCreateGroup: Packet {
    static let packetId = 0x11

    var groupId: String = ""

    func write(to stream: Stream) {
        stream.writeString(groupId)
    }

    mutating func read(from stream: Stream) {
        groupId = stream.readString()
    }
}

struct PacketGroupInfo: Packet {
    static let packetId = 0x12

    var groupId: String = ""
    var members: [String] = []

    func write(to stream: Stream) {
        stream.writeString(groupId)
        stream.writeInt16(members.count)
        for member in members {
            stream.writeString(member)
        }
    }

    mutating func read(from stream: Stream) {
        groupId = stream.readString()
        let count = max(stream.readInt16(), 0)
        var parsed: [String] = []
        parsed.reserveCapacity(count)
        for _ in 0..<count {
            parsed.append(stream.readString())
        }
        members = parsed
    }
}

struct PacketGroupInviteInfo: Packet {
    static let packetId = 0x13

    var groupId: String = ""
    var membersCount: Int = 0
    var status: GroupStatus = .notJoined

    func write(to stream: Stream) {
        stream.writeString(groupId)
        stream.writeInt16(membersCount)
        stream.writeInt8(status.rawValue)
    }

    mutating func read(from stream: Stream) {
        groupId = stream.readString()
        membersCount = max(stream.readInt16(), 0)
        status = GroupStatus(value: stream.readInt8())
    }
}

struct PacketGroupJoin: Packet {
    static let packetId = 0x14

    var groupId: String = ""
    var status: GroupStatus = .notJoined
    var groupString: String = ""

    func write(to stream: Stream) {
        stream.writeString(groupId)
        stream.writeInt8(status.rawValue)
        stream.writeString(groupString)
    }

    mutating func read(from stream: Stream) {
        groupId = stream.readString()
        status = GroupStatus(value: stream.readInt8())
        groupString = stream.readString()
    }
}

struct PacketGroupLeave: Packet {
    static let packetId = 0x15

    var groupId: String = ""

    func write(to stream: Stream) {
        stream.writeString(groupId)
    }

    mutating func read(from stream: Stream) {
        groupId = stream.readString()
    }
}

struct PacketGroupBan: Packet {
    static let packetId = 0x16

    var groupId: String = ""
    var publicKey: String = ""

    func write(to stream: Stream) {
        stream.writeString(groupId)
        stream.writeString(publicKey)
    }

    mutating func read(from stream: Stream) {
        groupId = stream.readString()
        publicKey = stream.readString()
    }
}
