using System;
using System.Diagnostics;
using System.Net.Sockets;
using System.Numerics;
using System.Text;
using System.Text.Json;
using System.Text.Json.Serialization;
using Encryption;

enum ConversationState
{
    Initial = 0,
    RsaPublicKeyGiven,
    AesKeySended,
    RsaPublicKeySended,
    AesKeyGiven
}

class TcpClientProgram
{
    static void Main(string[] args)
    {
        TcpClient client = new TcpClient();
        client.Connect("127.0.0.1", 5000);
        Console.WriteLine("Connected to server.");

        NetworkStream stream = client.GetStream();
        bool isRunning = true;
        ConversationState state = ConversationState.Initial;
        while (isRunning)
        {
            byte[] buffer = new byte[4096];
            int bytesRead;

            RsaPublicKey publicKey = null;
            RsaPrivateKey privateKey = null;
            string? aesKey = null;
            
            while ((bytesRead = stream.Read(buffer, 0, buffer.Length)) != 0)
            {
                string receivedData = Encoding.ASCII.GetString(buffer, 0, bytesRead);
                Console.WriteLine("Received: " + receivedData);

                // Split the message into command and data
                string[] messageParts = receivedData.Split(new char[] { ':' }, 2);
                if (messageParts.Length < 2) continue; // Ignore malformed messages

                string command = messageParts[0];
                string message = messageParts[1];

                switch (command)
                {
                    case "PUBKEY":
                        Console.WriteLine("Received RSA Public Key: " + message);
                        state = ConversationState.RsaPublicKeyGiven;
                        var json = Convert.FromBase64String(message);
                        var dict = JsonSerializer.Deserialize<Dictionary<string, string>>(json);
                        publicKey = new RsaPublicKey(BigInteger.Parse((string)dict["E"]),
                            BigInteger.Parse((string)dict["N"]));
                        break;
                    case "AESKEY":
                        Console.WriteLine("Received Encrypted AES Key: " + message);
                        state = ConversationState.AesKeyGiven;
                        aesKey = Encoding.UTF8.GetString(Convert.FromBase64String(message));
                        break;
                    case "MSG":
                        Console.WriteLine("Received AES Encrypted Message: " + message);
                        // todo: decrypt message and print it
                        var decodedMessage = Convert.FromBase64String(message);
                        var decryptedMessageBytes = RSA.Decrypt(privateKey, decodedMessage);
                        Console.WriteLine($"Decrypted message: {Encoding.UTF8.GetString(decryptedMessageBytes)}");
                        break;
                    case "EXIT":
                        Console.WriteLine("Client is disconnecting...");
                        client.Close();
                        return;
                    default:
                        Console.WriteLine("Unknown command received: " + command);
                        break;
                }
            }

            Console.WriteLine("1. Send RSA public key");
            Console.WriteLine("4. Send AES key encrypted with RSA");
            Console.WriteLine("5. Send/receive AES encrypted message");
            Console.WriteLine("6. Exit");

            string choice = Console.ReadLine();
            if (state == ConversationState.RsaPublicKeyGiven && choice == "1")
            {
                Console.WriteLine("You already recieved Public Key from participant. It's time to send him your AES key back");
                continue;
            }

            if (state == ConversationState.AesKeyGiven && choice == "4")
            {
                Console.WriteLine("You already have AES key from participant. It's time to send him encrypted message.");
            }

            switch (choice)
            {
                case "1":
                    // Placeholder: Send a simulated RSA public key
                    var json = publicKey.ToJsonString();
                    var message = Convert.ToBase64String(Encoding.UTF8.GetBytes(json));
                    SendMessage(stream, "PUBKEY", message);
                    state = ConversationState.RsaPublicKeySended;
                    break;
                case "4":
                    // Placeholder: Send an AES key encrypted with RSA
                    aesKey = ""; //todo: generate random
                    var aesKeyEncryptedBytes = RSA.Encrypt(publicKey, Encoding.UTF8.GetBytes(aesKey));
                    SendMessage(stream, "AESKEY", Convert.ToBase64String(aesKeyEncryptedBytes));
                    state = ConversationState.AesKeySended;
                    break;
                case "5":
                    string msg = Console.ReadLine();
                    // Placeholder: Send an AES encrypted message
                    SendMessage(stream, "MSG", Convert.ToBase64String(RSA.Encrypt(publicKey, Encoding.UTF8.GetBytes(msg))));
                    break;
                case "6":
                    SendMessage(stream, "EXIT", "Disconnecting");
                    isRunning = false;
                    Console.WriteLine("Exiting...");
                    break;
                default:
                    Console.WriteLine("Invalid choice. Please try again.");
                    break;
            }
        }

        client.Close();
    }

    static void SendMessage(NetworkStream stream, string command, string message)
    {
        string formattedMessage = $"{command}:{message}";
        byte[] data = Encoding.ASCII.GetBytes(formattedMessage);
        stream.Write(data, 0, data.Length);
        Console.WriteLine($"Sent [{command}]: {message}");
    }
}