#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <signal.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <sys/select.h>
#include <sys/wait.h>
#include <fcntl.h>
#include <errno.h>
#include <time.h>

// --- Configuration ---
#define PORT_VIDEO 5000
#define PORT_AUDIO 5001
#define DEVICE_VIDEO "/dev/video20"
#define SINK_AUDIO "AirLink_Sink"
#define SOURCE_MIC "AirLink_Mic"
#define BUFFER_SIZE 65535

// --- Globals ---
pid_t pid_video_process = 0;
pid_t pid_audio_process = 0;
int pipe_video_fd = -1;
int pipe_audio_fd = -1;
int socket_video = -1;
int socket_audio = -1;
int is_running = 1;

// Stats tracking
typedef struct {
    long long bytes_video;
    long long bytes_audio;
    int packets_video;
    int packets_audio;
    time_t last_log_time;
} TrafficStats;

TrafficStats stats = {0, 0, 0, 0, 0};

// --- Helpers ---

void kill_process(pid_t *pid) {
    if (*pid > 0) {
        kill(*pid, SIGTERM);
        waitpid(*pid, NULL, 0); // Wait to prevent zombies
        *pid = 0;
    }
}

// Generates a black screen on the virtual camera (used for Pause/Stop)
void render_black_screen() {
    kill_process(&pid_video_process);
    if (pipe_video_fd != -1) { close(pipe_video_fd); pipe_video_fd = -1; }

    printf("[STATE] Rendering black screen (Privacy Mode)\n");
    // We run this as a blocking system call because it's a short "state change" operation
    char cmd[512];
    snprintf(cmd, sizeof(cmd), 
        "ffmpeg -hide_banner -loglevel error -f lavfi -i color=black:s=1280x720 "
        "-pix_fmt yuyv422 -frames:v 5 -f v4l2 %s > /dev/null 2>&1", 
        DEVICE_VIDEO);
    system(cmd);
}

void cleanup(int sig) {
    is_running = 0;
    printf("\n[INFO] Shutting down services...\n");
    
    kill_process(&pid_video_process);
    kill_process(&pid_audio_process);

    // Clean up PulseAudio modules
    char cmd[256];
    snprintf(cmd, sizeof(cmd), "pactl unload-module %s > /dev/null 2>&1", SINK_AUDIO);
    system(cmd);
    system("pactl unload-module AirLink_Mic > /dev/null 2>&1");
    
    // Stop Avahi
    system("pkill -f 'avahi-publish-service -s AirLink'");
    
    // Final Black Screen
    render_black_screen();
    
    exit(0);
}

// Spawns a process and returns the Write FD for its STDIN
pid_t spawn_worker(char *const cmd[], int *fd_in) {
    int pipefd[2];
    if (pipe(pipefd) == -1) return -1;

    pid_t pid = fork();
    if (pid == 0) {
        // Child
        close(pipefd[1]); // Close write end
        dup2(pipefd[0], STDIN_FILENO); // Redirect pipe read to STDIN
        close(pipefd[0]);
        
        // Close other file descriptors to be safe
        for (int i = 3; i < 256; i++) close(i);
        
        execvp(cmd[0], cmd);
        exit(1); // Should not reach here
    } else if (pid > 0) {
        // Parent
        close(pipefd[0]); // Close read end
        *fd_in = pipefd[1]; // Keep write end
        return pid;
    }
    return -1;
}

void initialize_system() {
    printf("[INIT] Configuring Kernel Modules and Audio Subsystem...\n");
    
    // Video: Reload v4l2loopback
    system("sudo modprobe -r v4l2loopback 2>/dev/null");
    char v4l2_cmd[256];
    snprintf(v4l2_cmd, sizeof(v4l2_cmd), 
        "sudo modprobe v4l2loopback video_nr=20 card_label=AirLink_Cam exclusive_caps=1");
    system(v4l2_cmd);
    
    // Audio: Null Sink
    char pa_cmd[512];
    snprintf(pa_cmd, sizeof(pa_cmd), 
        "pactl load-module module-null-sink sink_name=%s "
        "sink_properties=device.description='AirLink_Audio' rate=48000 channels=1 >/dev/null 2>&1", 
        SINK_AUDIO);
    system(pa_cmd);
    
    // Audio: Remap Source (Virtual Mic)
    snprintf(pa_cmd, sizeof(pa_cmd), 
        "pactl load-module module-remap-source source_name=%s "
        "master=%s.monitor source_properties=device.description='AirLink_Mic' channels=1 >/dev/null 2>&1", 
        SOURCE_MIC, SINK_AUDIO);
    system(pa_cmd);
}

void start_video_stream() {
    // Check if already running
    if (pid_video_process > 0 && waitpid(pid_video_process, NULL, WNOHANG) == 0) return;

    // FFmpeg command to read raw H264 from stdin and write to V4L2
    char *cmd[] = {
        "ffmpeg", 
        "-fflags", "nobuffer", 
        "-flags", "low_delay",
        "-analyzeduration", "0", 
        "-probesize", "32",
        "-hide_banner", 
        "-loglevel", "quiet",
        "-f", "h264", 
        "-i", "pipe:0",
        "-vf", "scale=1280:720:force_original_aspect_ratio=decrease,pad=1280:720:(ow-iw)/2:(oh-ih)/2",
        "-vsync", "0", 
        "-copyts", 
        "-pix_fmt", "yuyv422",
        "-f", "v4l2", DEVICE_VIDEO, 
        NULL
    };
    
    pid_video_process = spawn_worker(cmd, &pipe_video_fd);
    signal(SIGPIPE, SIG_IGN); // Prevent crash if ffmpeg dies
}

void start_audio_stream() {
    if (pid_audio_process > 0 && waitpid(pid_audio_process, NULL, WNOHANG) == 0) return;

    // Use a shell to pipe ffmpeg decoder directly to pacat
    // Input: Raw PCM or Encoded Audio -> FFmpeg -> Raw PCM -> Pacat
    // Based on Python: Raw s16le incoming -> Pacat
    // If incoming is already raw PCM (s16le, 48k, mono), we can skip ffmpeg or just use it as pass-through.
    // We will use the shell wrapper for simplicity of piping.
    
    int pipefd[2];
    if (pipe(pipefd) == -1) return;

    pid_audio_process = fork();
    if (pid_audio_process == 0) {
        close(pipefd[1]);
        dup2(pipefd[0], STDIN_FILENO);
        
        // Command: Decode (ensure format) -> Play
        // We assume the Android app sends raw PCM (s16le) as per your Python script audio logic
        execl("/bin/sh", "sh", "-c", 
            "ffmpeg -hide_banner -loglevel quiet -fflags nobuffer -f s16le -ar 48000 -ac 1 -i pipe:0 "
            "-f s16le -ar 48000 -ac 1 pipe:1 | "
            "pacat --playback --device=AirLink_Sink --rate=48000 --channels=1 --format=s16le --latency-msec=20", 
            NULL);
        exit(1);
    }
    close(pipefd[0]);
    pipe_audio_fd = pipefd[1];
}

void print_stats() {
    time_t now = time(NULL);
    if (stats.last_log_time == 0) stats.last_log_time = now;

    double delta = difftime(now, stats.last_log_time);
    if (delta >= 1.0) {
        double video_mbps = (stats.bytes_video * 8.0) / (1000.0 * 1000.0) / delta;
        double audio_kbps = (stats.bytes_audio * 8.0) / 1000.0 / delta;
        
        printf("[STATS] Video: %6.2f Mbps (%d pps) | Audio: %6.2f Kbps (%d pps)\n",
            video_mbps, stats.packets_video,
            audio_kbps, stats.packets_audio);
            
        // Reset
        stats.bytes_video = 0;
        stats.bytes_audio = 0;
        stats.packets_video = 0;
        stats.packets_audio = 0;
        stats.last_log_time = now;
    }
}

int main() {
    signal(SIGINT, cleanup);
    
    initialize_system();

    // Start Avahi Advertising (Background)
    system("avahi-publish-service -s AirLink _procam._udp 5000 &");

    // Create UDP Sockets
    socket_video = socket(AF_INET, SOCK_DGRAM, 0);
    socket_audio = socket(AF_INET, SOCK_DGRAM, 0);

    // Set large kernel buffers to prevent drops
    int rcvbuf_video = 4 * 1024 * 1024; // 4MB
    int rcvbuf_audio = 512 * 1024;      // 512KB
    setsockopt(socket_video, SOL_SOCKET, SO_RCVBUF, &rcvbuf_video, sizeof(rcvbuf_video));
    setsockopt(socket_audio, SOL_SOCKET, SO_RCVBUF, &rcvbuf_audio, sizeof(rcvbuf_audio));

    struct sockaddr_in addr;
    memset(&addr, 0, sizeof(addr));
    addr.sin_family = AF_INET;
    addr.sin_addr.s_addr = htonl(INADDR_ANY);

    addr.sin_port = htons(PORT_VIDEO);
    if (bind(socket_video, (struct sockaddr*)&addr, sizeof(addr)) < 0) {
        perror("Bind Video Failed"); return 1;
    }

    addr.sin_port = htons(PORT_AUDIO);
    if (bind(socket_audio, (struct sockaddr*)&addr, sizeof(addr)) < 0) {
        perror("Bind Audio Failed"); return 1;
    }

    printf("[READY] Server listening on Ports %d (Video) & %d (Audio)\n", PORT_VIDEO, PORT_AUDIO);

    start_video_stream();
    start_audio_stream();

    unsigned char buffer[BUFFER_SIZE];
    fd_set readfds;
    int max_fd = (socket_video > socket_audio ? socket_video : socket_audio) + 1;
    
    struct sockaddr_in src_addr;
    socklen_t src_len;

    while (is_running) {
        print_stats();

        // Check process health and restart if needed
        if (waitpid(pid_video_process, NULL, WNOHANG) > 0) start_video_stream();
        if (waitpid(pid_audio_process, NULL, WNOHANG) > 0) start_audio_stream();

        FD_ZERO(&readfds);
        FD_SET(socket_video, &readfds);
        FD_SET(socket_audio, &readfds);

        struct timeval tv = {1, 0}; // 1 second timeout (allows stats to print even if idle)
        int activity = select(max_fd, &readfds, NULL, NULL, &tv);

        if (activity < 0 && errno != EINTR) break;

        // --- Handle Video & Control ---
        if (FD_ISSET(socket_video, &readfds)) {
            src_len = sizeof(src_addr);
            int n = recvfrom(socket_video, buffer, sizeof(buffer), 0, (struct sockaddr*)&src_addr, &src_len);
            
            if (n > 0) {
                // Control Packet Handling
                if (n < 16) {
                    if (memcmp(buffer, "PING", 4) == 0) {
                        sendto(socket_video, "PONG", 4, 0, (struct sockaddr*)&src_addr, src_len);
                    } 
                    else if (memcmp(buffer, "HEARTBEAT", 9) == 0) {
                        sendto(socket_video, "ACK", 3, 0, (struct sockaddr*)&src_addr, src_len);
                    }
                    else if (memcmp(buffer, "PAUSE", 5) == 0 || memcmp(buffer, "BYE", 3) == 0) {
                        render_black_screen();
                    }
                } 
                else {
                    // Video Data
                    if (pid_video_process == 0) start_video_stream();
                    write(pipe_video_fd, buffer, n);
                    
                    // Update Stats
                    stats.bytes_video += n;
                    stats.packets_video++;
                }
            }
        }

        // --- Handle Audio ---
        if (FD_ISSET(socket_audio, &readfds)) {
            int n = recvfrom(socket_audio, buffer, sizeof(buffer), 0, NULL, NULL);
            if (n > 0) {
                // Audio Data (ignore control packets on this port usually)
                if (n > 16) {
                    write(pipe_audio_fd, buffer, n);
                    
                    // Update Stats
                    stats.bytes_audio += n;
                    stats.packets_audio++;
                }
            }
        }
    }

    return 0;
}