#include <stdio.h>
#include <stdlib.h>


#include <libavformat/avformat.h>
#include <libavcodec/avcodec.h>
#include <libavutil/pixdesc.h>
#include <libswscale/swscale.h>
#include <libavutil/imgutils.h>


#include "rknn_api.h"  // Заголовочный файл RKNN SDK

#define MODEL_PATH "./model.rknn"  	// Путь к вашей RKNN модели
#define INPUT_WIDTH 640               // Ширина входа модели
#define INPUT_HEIGHT 640              // Высота входа модели
#define NUM_CLASSES 10                // Количество классов вашей модели

// Структура для хранения контекста RKNN
typedef struct {
    rknn_context ctx;
    rknn_input_output_num io_num;
    rknn_tensor_attr* input_attrs;
    rknn_tensor_attr* output_attrs;
} RKNNContext;

// Инициализация RKNN модели
int init_rknn_model(RKNNContext* rknn_ctx, const char* model_path) {
    int ret;
    
    // 1. Загрузка модели
    FILE* fp = fopen(model_path, "rb");
    if (!fp) {
        fprintf(stderr, "Failed to open RKNN model\n");
        return -1;
    }
    fseek(fp, 0, SEEK_END);
    int model_size = ftell(fp);
    void* model_data = malloc(model_size);
    fseek(fp, 0, SEEK_SET);
    if (fread(model_data, 1, model_size, fp) != model_size) {
        fprintf(stderr, "Failed to read RKNN model\n");
        fclose(fp);
        free(model_data);
        return -1;
    }
    fclose(fp);
    
    // 2. Инициализация RKNN
    ret = rknn_init(&rknn_ctx->ctx, model_data, model_size, 0, NULL);
    free(model_data);
    if (ret < 0) {
        fprintf(stderr, "rknn_init failed: %d\n", ret);
        return -1;
    }
    
    // 3. Получение информации о входах/выходах
    ret = rknn_query(rknn_ctx->ctx, RKNN_QUERY_IN_OUT_NUM, &rknn_ctx->io_num, sizeof(rknn_ctx->io_num));
    if (ret != RKNN_SUCC) {
        fprintf(stderr, "rknn_query failed: %d\n", ret);
        return -1;
    }
    
    // 4. Получение атрибутов тензоров
    rknn_ctx->input_attrs = calloc(rknn_ctx->io_num.n_input, sizeof(rknn_tensor_attr));
    for (int i = 0; i < rknn_ctx->io_num.n_input; i++) {
        rknn_ctx->input_attrs[i].index = i;
        ret = rknn_query(rknn_ctx->ctx, RKNN_QUERY_INPUT_ATTR, &(rknn_ctx->input_attrs[i]), sizeof(rknn_tensor_attr));
        if (ret != RKNN_SUCC) {
            fprintf(stderr, "rknn_query input_attrs failed: %d\n", ret);
            return -1;
        }
    }
    
    rknn_ctx->output_attrs = calloc(rknn_ctx->io_num.n_output, sizeof(rknn_tensor_attr));
    for (int i = 0; i < rknn_ctx->io_num.n_output; i++) {
        rknn_ctx->output_attrs[i].index = i;
        ret = rknn_query(rknn_ctx->ctx, RKNN_QUERY_OUTPUT_ATTR, &(rknn_ctx->output_attrs[i]), sizeof(rknn_tensor_attr));
        if (ret != RKNN_SUCC) {
            fprintf(stderr, "rknn_query output_attrs failed: %d\n", ret);
            return -1;
        }
    }
    
    return 0;
}

// Освобождение ресурсов RKNN
void release_rknn_model(RKNNContext* rknn_ctx) {
    if (rknn_ctx->ctx) {
        rknn_destroy(rknn_ctx->ctx);
    }
    if (rknn_ctx->input_attrs) {
        free(rknn_ctx->input_attrs);
    }
    if (rknn_ctx->output_attrs) {
        free(rknn_ctx->output_attrs);
    }
}

// Обновленный препроцессинг с учетом атрибутов модели
void preprocess_to_int8(AVFrame* frame, int8_t* output, int width, int height) {
    // Явная конвертация в INT8 с нормализацией
    for (int y = 0; y < height; y++) {
        for (int x = 0; x < width; x++) {
            // Получаем RGB компоненты
            uint8_t r = frame->data[0][y * frame->linesize[0] + x * 3];
            uint8_t g = frame->data[0][y * frame->linesize[0] + x * 3 + 1];
            uint8_t b = frame->data[0][y * frame->linesize[0] + x * 3 + 2];
            
            // Конвертация в INT8 (-128..127)
            output[y*width*3 + x*3]     = (int8_t)(r - 128);
            output[y*width*3 + x*3 + 1] = (int8_t)(g - 128);
            output[y*width*3 + x*3 + 2] = (int8_t)(b - 128);
        }
    }
}

int run_rknn_inference(RKNNContext* rknn_ctx, AVFrame* frame) {
    int ret;
    
    // 1. Проверка атрибутов входного тензора
    printf("Input tensor attributes:\n");
    printf("Dims: %d,%d,%d,%d\n", 
           rknn_ctx->input_attrs[0].dims[0],
           rknn_ctx->input_attrs[0].dims[1],
           rknn_ctx->input_attrs[0].dims[2],
           rknn_ctx->input_attrs[0].dims[3]);
    printf("Type: %s\n", get_type_string(rknn_ctx->input_attrs[0].type));
    printf("Size: %d\n", rknn_ctx->input_attrs[0].size);

    // 2. Подготовка входных данных
    rknn_input inputs[1];
    memset(&inputs[0], 0, sizeof(rknn_input));
    
    inputs[0].index = 0;
    inputs[0].type = RKNN_TENSOR_INT8;  // Важно! Модель ожидает INT8
    inputs[0].size = rknn_ctx->input_attrs[0].size;
    inputs[0].fmt = RKNN_TENSOR_NHWC;
    inputs[0].buf = malloc(inputs[0].size);
    
    if (!inputs[0].buf) {
        fprintf(stderr, "Memory allocation failed\n");
        return -1;
    }

    // 3. Препроцессинг с конвертацией в INT8
    for (int y = 0; y < 640; y++) {
        for (int x = 0; x < 640; x++) {
            for (int c = 0; c < 3; c++) {
                // Конвертация UINT8 (0-255) в INT8 (-128-127)
                int8_t* buf_int8 = (int8_t*)inputs[0].buf;
                buf_int8[y*640*3 + x*3 + c] = 
                    (int8_t)(frame->data[0][y*frame->linesize[0] + x*3 + c] - 128);
            }
        }
    }

    // 4. Установка входных данных
    ret = rknn_inputs_set(rknn_ctx->ctx, 1, inputs);
    if (ret < 0) {
        fprintf(stderr, "rknn_inputs_set failed: %d\n", ret);
        free(inputs[0].buf);
        return -1;
    }


    // 5. Выполнение inference
    ret = rknn_run(rknn_ctx->ctx, NULL);
    if (ret < 0) {
        fprintf(stderr, "rknn_run failed: %d\n", ret);
        free(inputs[0].buf);
        return -1;
    }

    // 6. Получение результатов
    rknn_output outputs[rknn_ctx->io_num.n_output];
    memset(outputs, 0, sizeof(outputs));
    
    for (int i = 0; i < rknn_ctx->io_num.n_output; i++) {
        outputs[i].want_float = 1;
        outputs[i].is_prealloc = 0;
    }
    
    ret = rknn_outputs_get(rknn_ctx->ctx, rknn_ctx->io_num.n_output, outputs, NULL);
    if (ret < 0) {
        fprintf(stderr, "rknn_outputs_get failed: %d\n", ret);
        free(inputs[0].buf);
        return -1;
    }
    
    // 7. Обработка результатов
    float* output_data = (float*)outputs[0].buf;
    int max_index = 0;
    float max_prob = output_data[0];
    
    for (int i = 1; i < NUM_CLASSES; i++) {
        if (output_data[i] > max_prob) {
            max_prob = output_data[i];
            max_index = i;
        }
    }
    
    printf("Detected class: %d with probability: %.2f%%\n", max_index, max_prob * 100);
    
    // 8. Освобождение ресурсов
    rknn_outputs_release(rknn_ctx->ctx, rknn_ctx->io_num.n_output, outputs);
    free(inputs[0].buf);
    
    return 0;
}

#define HAVE_SYS_AUXV_H 0
#define getauxval(x) 0UL

#ifndef __GLIBC__
__attribute__((weak)) unsigned long getauxval(unsigned long type) {
    (void)type;
    return 0;
}
#endif

// Обновите функцию сохранения
int save_frame_as_jpeg(AVFrame *pFrame, int width, int height, int iFrame) {
    // 1. Находим кодировщик (попробуйте оба варианта)
    const AVCodec *jpegCodec = avcodec_find_encoder(AV_CODEC_ID_MJPEG);
    if (!jpegCodec) {
        jpegCodec = avcodec_find_encoder_by_name("mjpeg");
        if (!jpegCodec) {
            fprintf(stderr, "JPEG encoder not found\n");
            return -1;
        }
    }

    // 2. Настройка контекста
    AVCodecContext *jpegContext = avcodec_alloc_context3(jpegCodec);
    if (!jpegContext) {
        fprintf(stderr, "Could not allocate JPEG codec context\n");
        return -1;
    }

    jpegContext->pix_fmt = pFrame->format; // Используем исходный формат
    jpegContext->width = width;
    jpegContext->height = height;
    jpegContext->time_base = (AVRational){1, 25};
    jpegContext->qmin = 10;  // Минимальное качество (0-51)
    jpegContext->qmax = 30;  // Максимальное качество

    // 3. Открываем кодировщик
    if (avcodec_open2(jpegContext, jpegCodec, NULL) < 0) {
        fprintf(stderr, "Could not open JPEG codec\n");
        avcodec_free_context(&jpegContext);
        return -1;
    }

    // 4. Кодируем кадр
    AVPacket *pkt = av_packet_alloc();
    if (!pkt) {
        fprintf(stderr, "Could not allocate packet\n");
        avcodec_free_context(&jpegContext);
        return -1;
    }

    int ret = avcodec_send_frame(jpegContext, pFrame);
    if (ret < 0) {
        fprintf(stderr, "Error sending frame to encoder\n");
        goto cleanup;
    }

    ret = avcodec_receive_packet(jpegContext, pkt);
    if (ret < 0) {
        fprintf(stderr, "Error receiving packet from encoder\n");
        goto cleanup;
    }

    // 5. Сохраняем файл
    char filename[64];
    snprintf(filename, sizeof(filename), "frame_%03d.jpg", iFrame);
    FILE *JPEGFile = fopen(filename, "wb");
    if (JPEGFile) {
        fwrite(pkt->data, 1, pkt->size, JPEGFile);
        fclose(JPEGFile);
        printf("Saved frame: %s (%d KB)\n", filename, pkt->size/1024);
    } else {
        perror("Failed to open output file");
    }

cleanup:
    av_packet_unref(pkt);
    av_packet_free(&pkt);
    avcodec_free_context(&jpegContext);
    return ret;
}


int main() {
    const char *url = "rtsp://192.168.1.11/live/0/MAIN";
    AVFormatContext *fmt_ctx = NULL;
    AVCodecContext *codec_ctx = NULL;
    const AVCodec *codec = NULL;
    AVPacket *pkt = av_packet_alloc();
    AVFrame *frame = av_frame_alloc(), *rgb_frame = av_frame_alloc();
    struct SwsContext *sws_ctx = NULL;
    int video_stream_idx = -1, frame_count = 0;
    RKNNContext rknn_ctx = {0};
    
    // 1. Инициализация RKNN модели
    if (init_rknn_model(&rknn_ctx, MODEL_PATH) != 0) {
        fprintf(stderr, "Failed to initialize RKNN model\n");
        return -1;
    }
    
    // 2. Инициализация FFmpeg (как в вашем оригинальном коде)
    avformat_network_init();
    
    AVDictionary *options = NULL;
    av_dict_set(&options, "rtsp_transport", "udp", 0);
    av_dict_set(&options, "stimeout", "5000000", 0);
    av_dict_set(&options, "buffer_size", "1024000", 0);

    printf("Opening input stream...\n");
    if ((avformat_open_input(&fmt_ctx, url, NULL, &options)) < 0) {
        fprintf(stderr, "Ошибка открытия RTSP потока\n");
        return -1;
    }
    av_dict_free(&options);

    if (avformat_find_stream_info(fmt_ctx, NULL) < 0) {
        fprintf(stderr, "Ошибка получения информации о потоке\n");
        return -1;
    }

    // Находим видеопоток
    for (int i = 0; i < fmt_ctx->nb_streams; i++) {
        if (fmt_ctx->streams[i]->codecpar->codec_type == AVMEDIA_TYPE_VIDEO) {
            video_stream_idx = i;
            break;
        }
    }
    
    if (video_stream_idx == -1) {
        fprintf(stderr, "Видеопоток не найден\n");
        return -1;
    }

    // Настраиваем декодер
    codec = avcodec_find_decoder_by_name("h264");
    if (!codec) {
        fprintf(stderr, "Не найден декодер H.264\n");
        return -1;
    }

    codec_ctx = avcodec_alloc_context3(codec);
    avcodec_parameters_to_context(codec_ctx, fmt_ctx->streams[video_stream_idx]->codecpar);

    AVDictionary *opts = NULL;
    av_dict_set(&opts, "threads", "1", 0);
    if (avcodec_open2(codec_ctx, codec, &opts) < 0) {
        fprintf(stderr, "Ошибка инициализации декодера\n");
        return -1;
    }
    av_dict_free(&opts);

    // Подготовка конвертера в RGB
    sws_ctx = sws_getContext(
        INPUT_WIDTH, INPUT_HEIGHT, codec_ctx->pix_fmt,
        INPUT_WIDTH, INPUT_HEIGHT, AV_PIX_FMT_RGB24,
        SWS_BILINEAR, NULL, NULL, NULL);
    
    if (!sws_ctx) {
        fprintf(stderr, "Ошибка инициализации конвертера\n");
        return -1;
    }

    // Выделяем память для RGB кадра
    rgb_frame->format = AV_PIX_FMT_RGB24;
    rgb_frame->width = INPUT_WIDTH;
    rgb_frame->height = INPUT_HEIGHT;
    av_frame_get_buffer(rgb_frame, 0);

    // Главный цикл декодирования
    while (av_read_frame(fmt_ctx, pkt) >= 0) {
        if (pkt->stream_index == video_stream_idx) {
            if (avcodec_send_packet(codec_ctx, pkt) == 0) {
                while (avcodec_receive_frame(codec_ctx, frame) == 0) {
                    // Конвертируем в RGB и изменяем размер
                    sws_scale(sws_ctx, 
                             (const uint8_t * const *)frame->data, frame->linesize,
                             0, frame->height,
                             rgb_frame->data, rgb_frame->linesize);
                    
                    // Обрабатываем каждый 10-й кадр
                    //if (frame_count % 10 == 0) {
											//	save_frame_as_jpeg(rgb_frame, codec_ctx->width, codec_ctx->height, frame_count/10);
											
											printf("RKNN run");
                        if (run_rknn_inference(&rknn_ctx, rgb_frame) != 0) {
                            fprintf(stderr, "Inference failed\n");
                        }
                    //}
                    frame_count++;
                }
            }
        }
        av_packet_unref(pkt);
    }

    // Очистка
    av_frame_free(&frame);
    av_frame_free(&rgb_frame);
    av_packet_free(&pkt);
    sws_freeContext(sws_ctx);
    avcodec_free_context(&codec_ctx);
    avformat_close_input(&fmt_ctx);
    avformat_network_deinit();
    release_rknn_model(&rknn_ctx);

    return 0;
}

E RKNN: rknn_run, input(images) tensor not attached to memory!
rknn_run failed: -5
Inference failed
RKNN runInput tensor attributes:
Dims: 1,640,640,3
Type: INT8
Size: 1228800