import 'dart:async';
import 'dart:convert';
import 'dart:io';
import 'dart:typed_data';

import 'package:libac_dart/encryption/aes.dart';
import 'package:libac_dart/encryption/xtea.dart';

import '../nbt/NbtIo.dart';
import '../nbt/Stream.dart';
import '../nbt/Tag.dart';
import '../nbt/impl/CompoundTag.dart';
import '../nbt/impl/StringTag.dart';

enum EncryptionType {
  AES(value: 2),
  XTEA(value: 1),
  NONE(value: 0);

  final int value;
  const EncryptionType({required this.value});

  static EncryptionType valueOf(int val) {
    switch (val) {
      case 1:
        {
          return EncryptionType.XTEA;
        }
      case 2:
        {
          return EncryptionType.AES;
        }
      default:
        {
          return EncryptionType.NONE;
        }
    }
  }
}

class PacketsConsts {
  /// Version of the packets system. Bumped when there is a major change to protocol
  static const VERSION = 2;
}

class PacketServer {
  static ServerSocket? socket;
  static bool shouldRestart = true;
  static EncryptionType encryptionType = EncryptionType.NONE;
  static Uint8List AESKey = Uint8List(0);
  static String PSK = "";
  static String TEA_SALT = "Harbinger 01/05/2025 @ 03:59:17 AM";

  /// Packet Data Format:
  ///
  /// 4 bytes (int)  - Version
  /// 8 Bytes (Long) - Total expected bytes of packet
  /// 8 bytes (Long) - Sequence ID
  /// 1 byte         - Encryption Type
  /// 8 bytes (Long) - Number of bytes to read
  /// <arbitrary> - NBT Data / Encrypted NBT Data
  ///
  /// Response Format:
  ///
  /// 4 Bytes (int)  - Version
  /// 8 bytes (Long) - Total expected bytes in packet
  /// 1 byte - Success flag, Zero or 255 currently.
  /// 1 byte - Encryption Type
  /// 8 byes (Long) - Packet Length
  /// <arbitrary> - NBT Data / Encrypted NBT Data
  ///
  static Future<void> start(int port) async {
    socket = await ServerSocket.bind(InternetAddress.anyIPv4, port);
    print("Server now listening on port ${port}");

    await for (var sock in socket!) {
      S2CResponse response = S2CResponse();
      print(
          "New connection from ${sock.remoteAddress.address}:${sock.remotePort}");

      ByteLayer layer = ByteLayer();

      try {
        sock.listen((data) async {
          layer.writeBytes(data);
          var oldPos = layer.currentPosition;
          layer.resetPosition();
          int version = layer.readInt();

          int pktTotalExpected = layer.readLong();
          if (pktTotalExpected <= layer.length) {
            // Allow Processing
          } else {
            layer.restorePosition(oldPos);
            return;
          }

          layer.resetPosition();
          layer.readInt();
          layer.readLong(); // This is unused outside of the above sanity check.
          try {
            int encryptType = layer.readByte();
            EncryptionType ENCType = EncryptionType.valueOf(encryptType);

            int sequenceID = layer.readLong();
            print("Sequence ID in request: $sequenceID");

            int numBytes = layer.readLong();

            List<int> remainingBytes = layer.readBytes(numBytes);
            if (ENCType == EncryptionType.AES) {
              int ivLen = layer.readInt();
              List<int> ivBytes = layer.readBytes(ivLen);
              AESData encData = AESData(iv: ivBytes, data: remainingBytes);

              AESCipher aes = await AESCipher.useKey(AESKey);
              remainingBytes = await aes.decrypt(encData);
            } else if (ENCType == EncryptionType.XTEA) {
              XTEA xtea = await XTEA(PSK);

              remainingBytes =
                  xtea.decipher(Uint8List.fromList(remainingBytes));
            }

            CompoundTag tag =
                await NbtIo.readFromStream(Uint8List.fromList(remainingBytes));
            StringBuilder builder = StringBuilder();
            Tag.writeStringifiedNamedTag(tag, builder, 0);

            print("Request from client: \n${builder}");

            C2SRequestPacket request = C2SRequestPacket();
            request.decodeTag(tag);

            PacketResponse reply = await request.handleServerPacket();

            // Server uses NBT to communicate
            builder = StringBuilder();
            Tag.writeStringifiedNamedTag(reply.replyDataTag, builder, 0);

            print("Response to client: \n${builder}");
            Uint8List nbtData = await NbtIo.writeToStream(reply.replyDataTag);

            layer.clear();
            layer.writeLong(sequenceID);
            layer.writeByte(0xFF); // Successful receipt
            layer.writeByte(ENCType.value);

            // Encryption Subroutine
            if (ENCType == EncryptionType.AES) {
              AESCipher aes = AESCipher.useKey(AESKey);
              var encData = await aes.encrypt(nbtData);
              nbtData = Uint8List.fromList(encData.data);

              layer.writeInt(encData.iv.length);
              layer.writeBytes(encData.iv);
            } else if (ENCType == EncryptionType.XTEA) {
              XTEA tea = await XTEA(PSK);
              nbtData = Uint8List.fromList(tea.encipher(nbtData));
            }

            layer.writeLong(nbtData.lengthInBytes);
            layer.writeBytes(nbtData);
            nbtData = layer.bytes;

            // NOTE: Added a length indicator because SocketServer is apparently... really really dumb in its impl, and has no way to know when all data has been received, so no special event. We just have to check for it based on this initial value.
            layer.clear();
            layer.writeInt(PacketsConsts.VERSION);
            layer.writeLong(nbtData.lengthInBytes + layer.currentPosition + 8);
            layer.writeBytes(nbtData);

            sock.add(layer.bytes);
          } catch (E, stack) {
            response.contents
                .put("error", StringTag.valueOf("Malformed request packet"));

            print(
                "Something went wrong. Malformed request? \n\n${E}\n\n${stack}\n\n\n\n");
          } finally {
            await sock.flush();
            sock.close();
          }
          layer.clear();
        }, onDone: () {
          layer.clear();
        }, onError: (E) {
          print("ERROR: ${E}");
          sock.close();
          layer.clear();
        });
      } catch (E) {
        sock.close();
      }
    }
  }
}

class PacketClient {
  Socket? socket;
  bool connected = false;
  String lastIP = "";
  int port = 25306;
  int packetSequence = 0;
  EncryptionType encryptionType = EncryptionType.NONE;
  Uint8List aesKey = Uint8List(0);

  /// Used for xtea encryption
  String PSK = "";

  PacketClient();

  Future<void> startConnect(String IPAddress, int port) async {
    try {
      socket = await Socket.connect(IPAddress, port);
      connected = true;
      lastIP = IPAddress;
      this.port = port;
    } catch (E) {
      connected = false;
      socket = null;
    }
  }

  /// Tries to send a packet to the connected server
  ///
  /// On success, returns either, the decoded [S2CResponse], or on error a S2CResponse containing an error and a stacktrace as [StringTag]
  ///
  /// Packet Data Format:
  ///
  /// 8 Bytes (Long) - Total expected bytes of packet minus the 8 bytes here.
  /// 8 bytes (Long) - Sequence ID
  /// 8 bytes (Long) - Number of bytes to read
  /// <arbitrary> - NBT Data
  ///
  /// Response Format:
  ///
  /// 8 bytes (Long) - Total expected bytes in packet minus the 8 bytes here.
  /// 1 byte - Success flag, Zero or 255 currently.
  /// 8 byes (Long) - Packet Length
  /// <arbitrary> - NBT Data
  ///
  Future<S2CResponse> send(IPacket packet, bool shouldReconnect) async {
    if (!connected) {
      return S2CResponse();
    }
    C2SRequestPacket request = C2SRequestPacket();
    request.payload = packet;
    request.cap = packet.getChannelID();
    bool success = false;

    ByteLayer layer = ByteLayer();

    Uint8List nbtData =
        await NbtIo.writeToStream(request.encodeTag().asCompoundTag());

    ByteLayer reply = ByteLayer();
    CompoundTag NBTTag = CompoundTag();

    while (!success) {
      layer.clear();
      layer.writeInt(PacketsConsts.VERSION);
      layer.writeLong(packetSequence);
      layer.writeByte(encryptionType.value);

      if (encryptionType == EncryptionType.AES) {
        AESCipher aes = AESCipher.useKey(aesKey);
        AESData encData = await aes.encrypt(nbtData);
        layer.writeInt(encData.iv.length);
        layer.writeBytes(encData.iv);

        nbtData = Uint8List.fromList(encData.data);
      } else if (encryptionType == EncryptionType.XTEA) {
        XTEA xtea = XTEA(PSK);
        nbtData = Uint8List.fromList(await xtea.encipher(nbtData.toList()));
      }

      layer.writeLong(nbtData.lengthInBytes);
      layer.writeBytes(nbtData);
      var tmpBytes = layer.bytes;
      layer.clear();
      layer.writeLong(tmpBytes.lengthInBytes);
      layer.writeBytes(tmpBytes);

      Completer responseWait = Completer();

      socket!.add(layer.bytes);

      socket!.listen((data) async {
        reply.writeBytes(data);

        var oldPos = reply.currentPosition;
        reply.resetPosition();
        int lenOfReply = reply.readLong();
        if (lenOfReply + 8 <= reply.length) {
          // We can now process the data
        } else {
          reply.restorePosition(oldPos);
          return;
        }

        // Validate response validity
        reply.resetPosition();
        reply.readLong(); // This is unused outside of the sanity check above.
        int sequence = reply.readLong();
        int successReceipt = reply.readByte();
        EncryptionType encType = EncryptionType.valueOf(layer.readByte());
        List<int> encIV = [];

        if (encType == EncryptionType.AES) {
          int ivLen = layer.readInt();
          encIV = layer.readBytes(ivLen);
        }

        int numBytes = reply.readLong();
        List<int> pktBytes = reply.readBytes(numBytes);

        if (successReceipt == 0xFF && packetSequence == sequence)
          success = true;

        if (success) {
          if (encType == EncryptionType.AES) {
            AESCipher aes = AESCipher.useKey(aesKey);
            AESData encData = AESData(iv: encIV, data: pktBytes);
            pktBytes = (await aes.decrypt(encData)).toList();
          } else if (encType == EncryptionType.XTEA) {
            XTEA xtea = XTEA(PSK);
            pktBytes = await xtea.decipher(pktBytes);
          }

          NBTTag = await NbtIo.readFromStream(Uint8List.fromList(pktBytes));
        }

        responseWait.complete();
      }, onError: (err) {
        if (!responseWait.isCompleted) responseWait.complete();
      });

      await responseWait.future;

      packetSequence++;
      if (!success) await Future.delayed(Duration(seconds: 5));
    }

    CompoundTag ct = CompoundTag();
    StringBuilder builder = StringBuilder();
    Tag.writeStringifiedNamedTag(NBTTag, builder, 0);

    print("Response from server: \n${builder}");
    ct.put("result", NBTTag);

    await close();

    if (shouldReconnect) await startConnect(lastIP, port);

    S2CResponse replyPkt = S2CResponse();
    try {
      replyPkt.decodeTag(ct.get("result")!.asCompoundTag());
    } catch (E, stack) {
      replyPkt.contents = CompoundTag(); // This is essentially a null response
      replyPkt.contents.put("error", StringTag.valueOf(E.toString()));
      replyPkt.contents.put("stacktrace", StringTag.valueOf(stack.toString()));
    }

    return replyPkt;
  }

  Future<void> close() async {
    await socket!.close();
    connected = false;
  }
}

abstract class IPacket with NbtEncodable, JsonEncodable {
  String getChannelID();

  // This function handles the packet
  Future<PacketResponse> handleServerPacket();
  Future<void> handleClientPacket();
  NetworkDirection direction();
}

class StopServerPacket extends IPacket {
  @override
  void decodeJson(String params) {}

  @override
  void decodeTag(Tag tag) {}

  @override
  NetworkDirection direction() {
    return NetworkDirection.ClientToServer;
  }

  @override
  String encodeJson() {
    return json.encode({});
  }

  @override
  Tag encodeTag() {
    return CompoundTag();
  }

  @override
  void fromJson(Map<String, dynamic> js) {}

  @override
  String getChannelID() {
    return "StopServer";
  }

  @override
  Future<PacketResponse> handleServerPacket() async {
    // We're now on the server. Handle the packet with a response to the client
    PacketServer.shouldRestart = false;

    S2CResponse response = S2CResponse();

    return PacketResponse(replyDataTag: response.encodeTag().asCompoundTag());
  }

  @override
  Map<String, dynamic> toJson() {
    return {};
  }

  @override
  Future<void> handleClientPacket() {
    throw UnimplementedError();
  }
}

class PacketResponse {
  static final nil = PacketResponse(replyDataTag: CompoundTag());

  PacketResponse({required this.replyDataTag});

  CompoundTag replyDataTag = CompoundTag();
}

class PacketRegistry {
  Map<String, IPacket Function()> _registry = {};
  static PacketRegistry _inst = PacketRegistry._();

  PacketRegistry._() {
    registerDefaults();
  }

  factory PacketRegistry() {
    return _inst;
  }

  int get count => _registry.length;

  void register(IPacket packet, IPacket Function() packetResolver) {
    _registry[packet.getChannelID()] = packetResolver;
  }

  IPacket getPacket(String channel) {
    if (_registry.containsKey(channel)) {
      IPacket Function() callback = _registry[channel]!;
      return callback();
    } else
      throw Exception("No such channel has been registered");
  }

  void registerDefaults() {
    register(S2CResponse(), () {
      return S2CResponse();
    });
    register(C2SRequestPacket(), () {
      return C2SRequestPacket();
    });
    register(StopServerPacket(), () {
      return StopServerPacket();
    });
    register(C2SPing(), () {
      return C2SPing();
    });
  }
}

enum NetworkDirection { ClientToServer, ServerToClient }

enum PacketOperation { Encode, Decode }

class S2CResponse implements IPacket {
  CompoundTag contents = CompoundTag();

  @override
  NetworkDirection direction() {
    return NetworkDirection.ServerToClient;
  }

  @override
  String getChannelID() {
    return "Response";
  }

  @override
  Future<PacketResponse> handleServerPacket() async {
    // We can't predict handling for this type, it is a data packet response with no pre-defined structure.
    return PacketResponse.nil;
  }

  @override
  void decodeJson(String encoded) {
    fromJson(json.decode(encoded));
  }

  @override
  void decodeTag(Tag encoded) {
    CompoundTag ct = encoded as CompoundTag;
    contents = ct.get("contents")!.asCompoundTag();
  }

  @override
  String encodeJson() {
    return json.encode(toJson());
  }

  @override
  Tag encodeTag() {
    CompoundTag tag = CompoundTag();
    tag.put("contents", contents);

    return tag;
  }

  @override
  void fromJson(Map<String, dynamic> params) {}

  @override
  Map<String, dynamic> toJson() {
    return {}; // Operation is not supported at this time.
  }

  @override
  Future<void> handleClientPacket() async {
    // We haven't got anything to process. This is structured data
  }
}

class C2SRequestPacket implements IPacket {
  String cap = ""; // Packet channel
  late IPacket payload;

  @override
  void decodeJson(String encoded) {
    fromJson(json.decode(encoded));
  }

  @override
  void decodeTag(Tag encoded) {
    CompoundTag tag = encoded.asCompoundTag();
    String cap = tag.get("cap")!.asString();
    payload = PacketRegistry().getPacket(cap);
    payload.decodeTag(tag.get("payload")!.asCompoundTag());
  }

  @override
  NetworkDirection direction() {
    return NetworkDirection.ClientToServer;
  }

  @override
  String encodeJson() {
    return json.encode(toJson());
  }

  @override
  Tag encodeTag() {
    CompoundTag tag = CompoundTag();
    tag.put("cap", StringTag.valueOf(payload.getChannelID()));
    tag.put("payload", payload.encodeTag());

    return tag;
  }

  @override
  void fromJson(Map<String, dynamic> params) {
    String cap = params['cap'] as String;
    payload = PacketRegistry().getPacket(cap);
    payload.fromJson(params['payload']);
  }

  @override
  String getChannelID() {
    return "C2SRequest";
  }

  @override
  Future<PacketResponse> handleServerPacket() async {
    // This has no internal handling
    return payload.handleServerPacket();
  }

  @override
  Map<String, dynamic> toJson() {
    return {"cap": payload.getChannelID(), "payload": payload.toJson()};
  }

  @override
  Future<void> handleClientPacket() {
    throw UnimplementedError();
  }
}

class C2SPing implements IPacket {
  String clientVersion = "";

  @override
  void decodeJson(String params) {
    fromJson(json.decode(params));
  }

  @override
  void decodeTag(Tag tag) {
    clientVersion = tag.asCompoundTag().get("version")!.asString();
  }

  @override
  NetworkDirection direction() {
    return NetworkDirection.ClientToServer;
  }

  @override
  String encodeJson() {
    return json.encode(toJson());
  }

  @override
  Tag encodeTag() {
    CompoundTag tag = CompoundTag();
    tag.put("version", StringTag.valueOf(clientVersion));

    return tag;
  }

  @override
  void fromJson(Map<String, dynamic> js) {
    clientVersion = js['version'] as String;
  }

  @override
  String getChannelID() {
    return "Ping";
  }

  @override
  Future<PacketResponse> handleServerPacket() async {
    CompoundTag tag = CompoundTag();
    tag.put("pong", StringTag.valueOf(Platform.version));

    S2CResponse response = S2CResponse();
    response.contents = tag;

    PacketResponse reply =
        PacketResponse(replyDataTag: response.encodeTag().asCompoundTag());
    return reply;
  }

  @override
  Map<String, dynamic> toJson() {
    return {"version": clientVersion};
  }

  @override
  Future<void> handleClientPacket() {
    throw UnimplementedError();
  }
}

mixin JsonEncodable {
  String encodeJson();
  void decodeJson(String params);

  Map<String, dynamic> toJson();
  void fromJson(Map<String, dynamic> js);
}

mixin NbtEncodable {
  Tag encodeTag();
  void decodeTag(Tag tag);
}