#include "rtsp_client.h"
#include "video_decoder.h"

#include <android/log.h>
#include <media/NdkMediaCodec.h>
#include <sys/socket.h>
#include <sys/mman.h>
#include <sys/syscall.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <netdb.h>
#include <unistd.h>
#include <fcntl.h>
#include <pthread.h>
#include <cstring>
#include <cerrno>
#include <cstdio>
#include <cstdlib>
#include <poll.h>

#define LOG_TAG "RtspClient"
#define ALOGD(...) __android_log_print(ANDROID_LOG_DEBUG, LOG_TAG, __VA_ARGS__)
#define ALOGE(...) __android_log_print(ANDROID_LOG_ERROR, LOG_TAG, __VA_ARGS__)
#define ALOGI(...) __android_log_print(ANDROID_LOG_INFO, LOG_TAG, __VA_ARGS__)
#define ALOGW(...) __android_log_print(ANDROID_LOG_WARN, LOG_TAG, __VA_ARGS__)

static rtsp_client::RtspState g_rtsp = {};
static AMediaCodec* g_codec = nullptr;
static uint8_t g_rtp_buf[rtsp_client::RTP_BUFFER_SIZE];
static uint8_t g_nal_buf[256 * 1024];
static int g_nal_pos = 0;

static int memfd_create_compat(const char* name, unsigned int flags) {
#ifdef __NR_memfd_create
    return (int)syscall(__NR_memfd_create, name, flags);
#else
    (void)flags;
    char path[256];
    snprintf(path, sizeof(path), "/data/local/tmp/camera_magic_rtsp_%s_%d", name, getpid());
    int fd = open(path, O_CREAT | O_RDWR | O_TRUNC, 0666);
    if (fd >= 0) unlink(path);
    return fd;
#endif
}

static bool parse_rtsp_url(const char* url, std::string& host, int& port, std::string& path) {
    const char* p = url;
    if (strncmp(p, "rtsp://", 7) != 0) return false;
    p += 7;
    
    const char* colon = strchr(p, ':');
    const char* slash = strchr(p, '/');
    
    if (colon && (!slash || colon < slash)) {
        host.assign(p, colon - p);
        port = atoi(colon + 1);
        p = colon + 1;
        while (*p >= '0' && *p <= '9') p++;
    } else {
        if (slash) {
            host.assign(p, slash - p);
        } else {
            host = p;
        }
        port = rtsp_client::RTSP_DEFAULT_PORT;
    }
    
    if (*p == '/') {
        path = p;
    } else {
        path = "/";
    }
    
    return true;
}

static int rtsp_send_request(int fd, const char* method, const char* url,
                              const char* extra_headers, int* cseq) {
    char req[2048];
    int len = snprintf(req, sizeof(req),
        "%s %s RTSP/1.0\r\n"
        "CSeq: %d\r\n"
        "%s"
        "\r\n",
        method, url, *cseq, extra_headers ? extra_headers : "");
    
    (*cseq)++;
    
    int sent = send(fd, req, len, MSG_NOSIGNAL);
    if (sent != len) {
        ALOGE("send failed: %d/%d (%s)", sent, len, strerror(errno));
        return -1;
    }
    
    ALOGD("Sent: %s CSeq=%d", method, *cseq - 1);
    return 0;
}

static int rtsp_read_response(int fd, char* buf, int buf_size, int* status_code,
                               char* session_buf, int session_buf_size) {
    int total = 0;
    int timeout_ms = 5000;
    
    struct pollfd pfd;
    pfd.fd = fd;
    pfd.events = POLLIN;
    
    while (total < buf_size - 1) {
        int ret = poll(&pfd, 1, timeout_ms);
        if (ret <= 0) {
            ALOGE("poll timeout/error: %d", ret);
            return -1;
        }
        
        ret = recv(fd, buf + total, buf_size - 1 - total, 0);
        if (ret <= 0) {
            ALOGE("recv failed: %s", strerror(errno));
            return -1;
        }
        total += ret;
        buf[total] = '\0';
        
        if (strstr(buf, "\r\n\r\n")) break;
    }
    
    if (sscanf(buf, "RTSP/1.0 %d", status_code) != 1) {
        ALOGE("Cannot parse status from response");
        return -1;
    }
    
    if (session_buf && session_buf_size > 0) {
        const char* sess = strstr(buf, "Session: ");
        if (sess) {
            sess += 9;
            int i = 0;
            while (*sess && *sess != '\r' && *sess != ';' && i < session_buf_size - 1) {
                session_buf[i++] = *sess++;
            }
            session_buf[i] = '\0';
        }
    }
    
    ALOGD("Response: %d", *status_code);
    return total;
}

static int create_udp_socket(int* port) {
    int fd = socket(AF_INET, SOCK_DGRAM, 0);
    if (fd < 0) return -1;
    
    struct sockaddr_in addr;
    memset(&addr, 0, sizeof(addr));
    addr.sin_family = AF_INET;
    addr.sin_addr.s_addr = htonl(INADDR_ANY);
    
    for (int p = 50000; p < 65000; p += 2) {
        addr.sin_port = htons(p);
        if (bind(fd, (struct sockaddr*)&addr, sizeof(addr)) == 0) {
            *port = p;
            return fd;
        }
    }
    
    close(fd);
    return -1;
}

static bool init_rtsp_shmem() {
    int w = g_rtsp.width > 0 ? g_rtsp.width : 1920;
    int h = g_rtsp.height > 0 ? g_rtsp.height : 1080;
    
    size_t size = video_decoder::calc_shmem_size(w, h);
    int fd = memfd_create_compat("cam2magic_rtsp", 0);
    if (fd < 0) {
        ALOGE("memfd_create failed: %s", strerror(errno));
        return false;
    }
    
    if (ftruncate(fd, size) < 0) {
        ALOGE("ftruncate failed: %s", strerror(errno));
        close(fd);
        return false;
    }
    
    uint8_t* base = (uint8_t*)mmap(nullptr, size, PROT_READ | PROT_WRITE, MAP_SHARED, fd, 0);
    if (base == MAP_FAILED) {
        ALOGE("mmap failed: %s", strerror(errno));
        close(fd);
        return false;
    }
    
    memset(base, 0, size);
    
    auto* header = reinterpret_cast<video_decoder::RingBufferHeader*>(base);
    header->magic = video_decoder::RING_MAGIC;
    header->version = video_decoder::RING_VERSION;
    header->frame_count = video_decoder::RING_FRAME_COUNT;
    header->frame_width = w;
    header->frame_height = h;
    header->frame_size = (uint32_t)(w * h * 3 / 2);
    
    g_rtsp.shmem_fd = fd;
    g_rtsp.shmem_base = base;
    g_rtsp.shmem_size = size;
    
    ALOGI("RTSP shared memory: fd=%d, size=%zu, %dx%d", fd, size, w, h);
    return true;
}

static void release_rtsp_shmem() {
    if (g_rtsp.shmem_base && g_rtsp.shmem_base != MAP_FAILED) {
        munmap(g_rtsp.shmem_base, g_rtsp.shmem_size);
        g_rtsp.shmem_base = nullptr;
    }
    if (g_rtsp.shmem_fd >= 0) {
        close(g_rtsp.shmem_fd);
        g_rtsp.shmem_fd = -1;
    }
    g_rtsp.shmem_size = 0;
}

static void feed_nal_to_decoder(const uint8_t* nal, int nal_len) {
    if (!g_codec || nal_len <= 0) return;
    
    ssize_t idx = AMediaCodec_dequeueInputBuffer(g_codec, 1000);
    if (idx < 0) return;
    
    size_t buf_size = 0;
    uint8_t* buf = AMediaCodec_getInputBuffer(g_codec, idx, &buf_size);
    if (!buf || (size_t)nal_len > buf_size) return;
    
    memcpy(buf, nal, nal_len);
    
    uint32_t flags = 0;
    if (nal[0] == 0 && nal[1] == 0 && nal[2] == 0 && nal[3] == 1) {
        int type = nal[4] & 0x1F;
        if (type == 5 || type == 7 || type == 8) {
            flags = AMEDIACODEC_BUFFER_FLAG_CODEC_CONFIG;
        }
    }
    
    AMediaCodec_queueInputBuffer(g_codec, idx, 0, nal_len, 0, flags);
}

static void drain_decoder_output() {
    if (!g_codec) return;
    
    AMediaCodecBufferInfo info;
    ssize_t idx = AMediaCodec_dequeueOutputBuffer(g_codec, &info, 0);
    
    if (idx >= 0) {
        size_t out_size = 0;
        uint8_t* out_buf = AMediaCodec_getOutputBuffer(g_codec, idx, &out_size);
        
        if (out_buf && info.size > 0 && g_rtsp.shmem_base) {
            int w = g_rtsp.width;
            int h = g_rtsp.height;
            
            if (w > 0 && h > 0 && out_size >= (size_t)(w * h * 3 / 2)) {
                const uint8_t* y = out_buf;
                const uint8_t* u = out_buf + w * h;
                const uint8_t* v = u + (w / 2) * (h / 2);
                
                video_decoder::write_frame_to_ring(y, u, v, w, w / 2, w, h);
                g_rtsp.frame_count++;
                
                if (g_rtsp.frame_count % 300 == 0) {
                    ALOGI("RTSP decoded %u frames", g_rtsp.frame_count);
                }
            }
        }
        
        AMediaCodec_releaseOutputBuffer(g_codec, idx, false);
    } else if (idx == AMEDIACODEC_INFO_OUTPUT_FORMAT_CHANGED) {
        AMediaFormat* fmt = AMediaCodec_getOutputFormat(g_codec);
        int32_t w = 0, h = 0;
        AMediaFormat_getInt32(fmt, AMEDIAFORMAT_KEY_WIDTH, &w);
        AMediaFormat_getInt32(fmt, AMEDIAFORMAT_KEY_HEIGHT, &h);
        ALOGI("RTSP output format changed: %dx%d", w, h);
        if (w > 0 && h > 0) {
            g_rtsp.width = w;
            g_rtsp.height = h;
        }
        AMediaFormat_delete(fmt);
    }
}

static void* rtsp_receiver_thread(void* arg) {
    (void)arg;
    ALOGI("RTSP receiver thread started");
    
    while (g_rtsp.running) {
        struct pollfd pfd;
        pfd.fd = g_rtsp.rtp_fd;
        pfd.events = POLLIN;
        
        int ret = poll(&pfd, 1, 1000);
        if (ret <= 0) continue;
        
        ret = recv(g_rtsp.rtp_fd, g_rtp_buf, sizeof(g_rtp_buf), 0);
        if (ret <= 0) continue;
        
        if (ret < 12) continue;
        
        int payload_type = g_rtp_buf[1] & 0x7F;
        
        if (payload_type != 96 && payload_type != 97 && payload_type != 98 && payload_type != 99 &&
            payload_type != 100 && payload_type != 101 && payload_type != 102 && payload_type != 103 &&
            payload_type != 104 && payload_type != 105 && payload_type != 106 && payload_type != 107 &&
            payload_type != 108 && payload_type != 109 && payload_type != 110 && payload_type != 111 &&
            payload_type != 112 && payload_type != 113 && payload_type != 114 && payload_type != 115 &&
            payload_type != 116 && payload_type != 117 && payload_type != 118 && payload_type != 119 &&
            payload_type != 120 && payload_type != 121 && payload_type != 122 && payload_type != 123 &&
            payload_type != 124 && payload_type != 125 && payload_type != 126 && payload_type != 127) {
            continue;
        }
        
        const uint8_t* payload = g_rtp_buf + 12;
        int payload_len = ret - 12;
        
        if (payload_len < 2) continue;
        
        int nal_type = payload[0] & 0x1F;
        
        if (nal_type == 28 || nal_type == 29) {
            int fu_header = payload[1];
            int start = fu_header & 0x80;
            int end = fu_header & 0x40;
            int nal_hdr = (payload[0] & 0xE0) | (fu_header & 0x1F);
            
            if (start) {
                g_nal_buf[0] = 0;
                g_nal_buf[1] = 0;
                g_nal_buf[2] = 0;
                g_nal_buf[3] = 1;
                g_nal_buf[4] = nal_hdr;
                g_nal_pos = 5;
                
                int copy_len = payload_len - 2;
                if (g_nal_pos + copy_len < (int)sizeof(g_nal_buf)) {
                    memcpy(g_nal_buf + g_nal_pos, payload + 2, copy_len);
                    g_nal_pos += copy_len;
                }
            } else {
                int copy_len = payload_len - 2;
                if (g_nal_pos + copy_len < (int)sizeof(g_nal_buf)) {
                    memcpy(g_nal_buf + g_nal_pos, payload + 2, copy_len);
                    g_nal_pos += copy_len;
                }
            }
            
            if (end) {
                feed_nal_to_decoder(g_nal_buf, g_nal_pos);
                drain_decoder_output();
                g_nal_pos = 0;
                g_rtsp.bytes_received += ret;
            }
        } else if (nal_type >= 1 && nal_type <= 23) {
            g_nal_buf[0] = 0;
            g_nal_buf[1] = 0;
            g_nal_buf[2] = 0;
            g_nal_buf[3] = 1;
            
            int copy_len = payload_len;
            if (4 + copy_len < (int)sizeof(g_nal_buf)) {
                memcpy(g_nal_buf + 4, payload, copy_len);
                feed_nal_to_decoder(g_nal_buf, 4 + copy_len);
                drain_decoder_output();
                g_rtsp.bytes_received += ret;
            }
        }
    }
    
    ALOGI("RTSP receiver thread finished");
    return nullptr;
}

int rtsp_client::init_rtsp(const char* url) {
    if (!url) return -1;
    
    memset(&g_rtsp, 0, sizeof(g_rtsp));
    g_rtsp.url = url;
    g_rtsp.rtsp_fd = -1;
    g_rtsp.rtp_fd = -1;
    g_rtsp.rtcp_fd = -1;
    g_rtsp.shmem_fd = -1;
    g_rtsp.width = 1920;
    g_rtsp.height = 1080;
    
    if (!parse_rtsp_url(url, g_rtsp.host, g_rtsp.port, g_rtsp.path)) {
        ALOGE("Invalid RTSP URL: %s", url);
        return -1;
    }
    
    ALOGI("RTSP init: %s:%d%s", g_rtsp.host.c_str(), g_rtsp.port, g_rtsp.path.c_str());
    return 0;
}

int rtsp_client::connect_rtsp() {
    int cseq = 1;
    char resp_buf[4096];
    int status = 0;
    char session_buf[256] = {0};
    
    struct sockaddr_in server;
    memset(&server, 0, sizeof(server));
    server.sin_family = AF_INET;
    server.sin_port = htons(g_rtsp.port);
    
    struct hostent* he = gethostbyname(g_rtsp.host.c_str());
    if (!he) {
        ALOGE("Cannot resolve %s", g_rtsp.host.c_str());
        return -1;
    }
    memcpy(&server.sin_addr, he->h_addr, he->h_length);
    
    g_rtsp.rtsp_fd = socket(AF_INET, SOCK_STREAM, 0);
    if (g_rtsp.rtsp_fd < 0) {
        ALOGE("socket failed: %s", strerror(errno));
        return -1;
    }
    
    if (connect(g_rtsp.rtsp_fd, (struct sockaddr*)&server, sizeof(server)) < 0) {
        ALOGE("connect failed: %s", strerror(errno));
        close(g_rtsp.rtsp_fd);
        g_rtsp.rtsp_fd = -1;
        return -1;
    }
    
    ALOGI("Connected to %s:%d", g_rtsp.host.c_str(), g_rtsp.port);
    
    char desc_url[1024];
    snprintf(desc_url, sizeof(desc_url), "rtsp://%s:%d%s",
              g_rtsp.host.c_str(), g_rtsp.port, g_rtsp.path.c_str());
    
    if (rtsp_send_request(g_rtsp.rtsp_fd, "DESCRIBE", desc_url,
                           "Accept: application/sdp\r\n", &cseq) < 0) {
        return -1;
    }
    
    if (rtsp_read_response(g_rtsp.rtsp_fd, resp_buf, sizeof(resp_buf), &status,
                            session_buf, sizeof(session_buf)) < 0 || status != 200) {
        ALOGE("DESCRIBE failed: %d", status);
        return -1;
    }
    
    const char* control = strstr(resp_buf, "a=control:");
    if (control) {
        control += 10;
        while (*control == ' ') control++;
        char* end = strchr((char*)control, '\r');
        if (end) {
            g_rtsp.control_url.assign(control, end - control);
            ALOGI("Control URL: %s", g_rtsp.control_url.c_str());
        }
    }
    
    const char* sdp_w = strstr(resp_buf, "a=framesize:");
    if (sdp_w) {
        int w = 0, h = 0;
        if (sscanf(sdp_w, "a=framesize:%*d %d-%d", &w, &h) == 2 && w > 0 && h > 0) {
            g_rtsp.width = w;
            g_rtsp.height = h;
            ALOGI("SDP resolution: %dx%d", w, h);
        }
    }
    
    g_rtsp.rtp_fd = create_udp_socket(&g_rtsp.local_rtp_port);
    if (g_rtsp.rtp_fd < 0) {
        ALOGE("Cannot create RTP socket");
        return -1;
    }
    
    g_rtsp.local_rtcp_port = g_rtsp.local_rtp_port + 1;
    
    char transport[512];
    snprintf(transport, sizeof(transport),
             "Transport: RTP/AVP/UDP;unicast;client_port=%d-%d\r\n",
             g_rtsp.local_rtp_port, g_rtsp.local_rtcp_port);
    
    char setup_url[1024];
    if (!g_rtsp.control_url.empty() && g_rtsp.control_url.find("rtsp://") == 0) {
        snprintf(setup_url, sizeof(setup_url), "%s", g_rtsp.control_url.c_str());
    } else if (!g_rtsp.control_url.empty()) {
        snprintf(setup_url, sizeof(setup_url), "rtsp://%s:%d/%s",
                  g_rtsp.host.c_str(), g_rtsp.port, g_rtsp.control_url.c_str());
    } else {
        snprintf(setup_url, sizeof(setup_url), "%s", desc_url);
    }
    
    if (rtsp_send_request(g_rtsp.rtsp_fd, "SETUP", setup_url, transport, &cseq) < 0) {
        return -1;
    }
    
    memset(session_buf, 0, sizeof(session_buf));
    if (rtsp_read_response(g_rtsp.rtsp_fd, resp_buf, sizeof(resp_buf), &status,
                            session_buf, sizeof(session_buf)) < 0 || status != 200) {
        ALOGE("SETUP failed: %d", status);
        return -1;
    }
    
    if (session_buf[0]) {
        g_rtsp.session_id = session_buf;
        ALOGI("Session: %s", g_rtsp.session_id.c_str());
    }
    
    if (!init_rtsp_shmem()) {
        return -1;
    }
    
    const char* mime = "video/avc";
    g_codec = AMediaCodec_createDecoderByType(mime);
    if (!g_codec) {
        ALOGE("Cannot create MediaCodec decoder");
        return -1;
    }
    
    AMediaFormat* fmt = AMediaFormat_new();
    AMediaFormat_setString(fmt, AMEDIAFORMAT_KEY_MIME, mime);
    AMediaFormat_setInt32(fmt, AMEDIAFORMAT_KEY_WIDTH, g_rtsp.width);
    AMediaFormat_setInt32(fmt, AMEDIAFORMAT_KEY_HEIGHT, g_rtsp.height);
    
    media_status_t st = AMediaCodec_configure(g_codec, fmt, nullptr, nullptr, 0);
    AMediaFormat_delete(fmt);
    if (st != AMEDIA_OK) {
        ALOGE("AMediaCodec_configure failed: %d", st);
        AMediaCodec_delete(g_codec);
        g_codec = nullptr;
        return -1;
    }
    
    st = AMediaCodec_start(g_codec);
    if (st != AMEDIA_OK) {
        ALOGE("AMediaCodec_start failed: %d", st);
        AMediaCodec_delete(g_codec);
        g_codec = nullptr;
        return -1;
    }
    
    g_rtsp.connected = true;
    ALOGI("RTSP connected, codec started");
    return 0;
}

int rtsp_client::play_rtsp() {
    if (!g_rtsp.connected) return -1;
    
    int cseq = 4;
    char resp_buf[4096];
    int status = 0;
    char session_buf[256] = {0};
    
    char play_url[1024];
    if (!g_rtsp.control_url.empty() && g_rtsp.control_url.find("rtsp://") == 0) {
        snprintf(play_url, sizeof(play_url), "%s", g_rtsp.control_url.c_str());
    } else {
        snprintf(play_url, sizeof(play_url), "rtsp://%s:%d%s",
                  g_rtsp.host.c_str(), g_rtsp.port, g_rtsp.path.c_str());
    }
    
    char session_hdr[512];
    snprintf(session_hdr, sizeof(session_hdr), "Session: %s\r\n", g_rtsp.session_id.c_str());
    
    if (rtsp_send_request(g_rtsp.rtsp_fd, "PLAY", play_url, session_hdr, &cseq) < 0) {
        return -1;
    }
    
    if (rtsp_read_response(g_rtsp.rtsp_fd, resp_buf, sizeof(resp_buf), &status,
                            session_buf, sizeof(session_buf)) < 0 || status != 200) {
        ALOGE("PLAY failed: %d", status);
        return -1;
    }
    
    g_rtsp.running = true;
    g_rtsp.playing = true;
    
    int ret = pthread_create(&g_rtsp.receiver_thread, nullptr, rtsp_receiver_thread, nullptr);
    if (ret != 0) {
        ALOGE("pthread_create failed: %s", strerror(ret));
        g_rtsp.running = false;
        return -1;
    }
    
    ALOGI("RTSP playing, receiver thread started");
    return 0;
}

void rtsp_client::stop_rtsp() {
    if (!g_rtsp.running) return;
    
    ALOGI("Stopping RTSP");
    g_rtsp.running = false;
    g_rtsp.playing = false;
    
    if (g_rtsp.receiver_thread) {
        pthread_join(g_rtsp.receiver_thread, nullptr);
        g_rtsp.receiver_thread = 0;
    }
    
    if (g_codec) {
        AMediaCodec_stop(g_codec);
        AMediaCodec_delete(g_codec);
        g_codec = nullptr;
    }
    
    if (g_rtsp.rtsp_fd >= 0) {
        int cseq = 10;
        char teardown_url[1024];
        snprintf(teardown_url, sizeof(teardown_url), "rtsp://%s:%d%s",
                  g_rtsp.host.c_str(), g_rtsp.port, g_rtsp.path.c_str());
        char session_hdr[512];
        snprintf(session_hdr, sizeof(session_hdr), "Session: %s\r\n", g_rtsp.session_id.c_str());
        rtsp_send_request(g_rtsp.rtsp_fd, "TEARDOWN", teardown_url, session_hdr, &cseq);
        close(g_rtsp.rtsp_fd);
        g_rtsp.rtsp_fd = -1;
    }
    
    if (g_rtsp.rtp_fd >= 0) {
        close(g_rtsp.rtp_fd);
        g_rtsp.rtp_fd = -1;
    }
    
    release_rtsp_shmem();
    
    g_rtsp.connected = false;
    ALOGI("RTSP stopped");
}

void rtsp_client::release_rtsp() {
    stop_rtsp();
    g_rtsp.url.clear();
    g_rtsp.host.clear();
    g_rtsp.path.clear();
    g_rtsp.session_id.clear();
    g_rtsp.control_url.clear();
    ALOGI("RTSP released");
}

bool rtsp_client::is_rtsp_connected() {
    return g_rtsp.connected && g_rtsp.playing;
}

int rtsp_client::get_rtsp_shmem_fd() {
    return g_rtsp.shmem_fd;
}

int rtsp_client::get_rtsp_width() {
    return g_rtsp.width;
}

int rtsp_client::get_rtsp_height() {
    return g_rtsp.height;
}
