package main

import (
	"bytes"
	"errors"
	"log"
	"os"
	"strings"
	"sync"
	"time"

	"github.com/Eyevinn/mp4ff/mp4"
	"github.com/bluenviron/gortsplib/v4"
	"github.com/bluenviron/gortsplib/v4/pkg/base"
	"github.com/bluenviron/gortsplib/v4/pkg/description"
	"github.com/bluenviron/gortsplib/v4/pkg/format"
	"github.com/bluenviron/gortsplib/v4/pkg/format/rtph264"
	"github.com/gen2brain/aac-go"
	"github.com/pion/rtp"
	"github.com/zaf/g711"
)

var (
	currentVideoFile      *os.File
	currentAudioFile      *os.File
	currentMP4File        *os.File
	wg                    sync.WaitGroup
	fileMu                sync.Mutex
	formaH264Mu           sync.Mutex
	formaH264             *format.H264
	currentSegmentStarted bool // Флаг указывает записан ли уже IDR (с SPS/PPS) в текущий сегмент
)

func main() {
	// Парсинг URL RTSP
	u, err := base.ParseURL("rtsp://intercom-video-1.insit.ru/dp-gfahybswjoendkcpbaqvgkeizruudbhsfdr")
	if err != nil {
		log.Fatalln("Ошибка парсинга URL:", err)
	}

	// Инициализация клиента и подключение к камере
	client := gortsplib.Client{}
	err = client.Start(u.Scheme, u.Host)
	if err != nil {
		log.Fatalln("Ошибка соединения:", err)
	}
	defer client.Close()

	// Получение описания
	desc, _, err := client.Describe(u)
	if err != nil {
		log.Fatalln("Ошибка получения описания:", err)
	}

	// Установка настроек
	err = client.SetupAll(desc.BaseURL, desc.Medias)
	if err != nil {
		log.Fatalln("Ошибка настройки:", err)
	}

	// Проверка форматов для видео- и аудио-потоков
	var h264Format *format.H264
	var g711Format *format.G711
	for _, media := range desc.Medias {
		switch media.Formats[0].(type) {
		case *format.H264:
			h264Format = media.Formats[0].(*format.H264)
		case *format.G711:
			g711Format = media.Formats[0].(*format.G711)
		}
	}
	if h264Format == nil || g711Format == nil {
		log.Fatalln("Форматы не найдены")
	}

	// Создание декодера для видео-потока H264
	decoderH264, err := h264Format.CreateDecoder()
	if err != nil {
		log.Fatalln("Ошибка создания декодера H264:", err)
	}

	// Обработка RTP-пакетов: для видео – сбор NAL-ов и запись в файл с префиксом, для аудио – запись PCM-данных.
	client.OnPacketRTPAny(func(medi *description.Media, forma format.Format, pkt *rtp.Packet) {
		switch f := forma.(type) {
		case *format.H264:
			formaH264Mu.Lock()
			formaH264 = f
			formaH264Mu.Unlock()

			nalus, err := decoderH264.Decode(pkt)
			if err != nil && !errors.Is(err, rtph264.ErrMorePacketsNeeded) {
				log.Printf("Ошибка декодирования H264: %v", err)
				return
			}

			fileMu.Lock()
			if currentVideoFile != nil {
				// Если сегмент ещё не начат, ожидается появление ключевого кадра (IDR).
				if !currentSegmentStarted {
					var isIDR bool
					for _, nalu := range nalus {
						if len(nalu) > 0 && (nalu[0]&0x1F) == 5 {
							isIDR = true
							break
						}
					}
					if !isIDR {
						// Пакет не записывается, если в нём нет ключевого кадра.
						fileMu.Unlock()
						return
					}
					// При получении IDR происходит вставка SPS и PPS (с префиксами) в начало сегмента.
					if len(f.SPS) > 0 {
						if _, err := currentVideoFile.Write([]byte{0x00, 0x00, 0x00, 0x01}); err != nil {
							log.Printf("Ошибка записи стартового кода SPS: %v", err)
						}
						if _, err := currentVideoFile.Write(f.SPS); err != nil {
							log.Printf("Ошибка записи SPS: %v", err)
						}
					}
					if len(f.PPS) > 0 {
						if _, err := currentVideoFile.Write([]byte{0x00, 0x00, 0x00, 0x01}); err != nil {
							log.Printf("Ошибка записи стартового кода PPS: %v", err)
						}
						if _, err := currentVideoFile.Write(f.PPS); err != nil {
							log.Printf("Ошибка записи PPS: %v", err)
						}
					}
					currentSegmentStarted = true
				}
				// Запись каждой NAL-единицы с префиксом
				for _, nalu := range nalus {
					if _, err := currentVideoFile.Write([]byte{0x00, 0x00, 0x00, 0x01}); err != nil {
						log.Printf("Ошибка записи стартового кода NALU: %v", err)
					}
					if _, err := currentVideoFile.Write(nalu); err != nil {
						log.Printf("Ошибка записи NALU: %v", err)
					}
				}
			}
			fileMu.Unlock()

		case *format.G711:
			if strings.Contains(forma.RTPMap(), "PCMA") {
				sampleG711a := g711.DecodeAlaw(pkt.Payload)
				fileMu.Lock()
				if currentAudioFile != nil {
					if _, err := currentAudioFile.Write(sampleG711a); err != nil {
						log.Printf("Ошибка записи аудио: %v", err)
					}
				}
				fileMu.Unlock()
			} else if strings.Contains(forma.RTPMap(), "PCMU") {
				sampleG711u := g711.DecodeUlaw(pkt.Payload)
				fileMu.Lock()
				if currentAudioFile != nil {
					if _, err := currentAudioFile.Write(sampleG711u); err != nil {
						log.Printf("Ошибка записи аудио: %v", err)
					}
				}
				fileMu.Unlock()
			} else {
				log.Println("Аудиокодек не идентифицирован")
			}
		}
	})

	// Запуск воспроизведения
	_, err = client.Play(nil)
	if err != nil {
		log.Fatalln("Ошибка запуска воспроизведения:", err)
	}

	// Параметр записи
	period := time.Hour

	// Ожидание начала следующего часа
	now := time.Now()
	nextSegment := now.Truncate(time.Hour).Add(time.Hour)
	waitDuration := nextSegment.Sub(now)
	log.Printf("Ожидание до начала записи: %v\n", waitDuration)
	time.Sleep(waitDuration)
	log.Println("Начало записи фрагмента")

	// Создание начальных файлов
	initialTimestamp := time.Now().Format("15-04_02-01-2006")
	videoFilename := initialTimestamp + ".h264"
	audioFilename := initialTimestamp + ".pcm"
	mp4Filename := initialTimestamp + ".mp4"

	fileVideo, err := os.Create(videoFilename)
	if err != nil {
		log.Fatalln("Ошибка создания видеофайла:", err)
	}
	fileAudio, err := os.Create(audioFilename)
	if err != nil {
		log.Fatalln("Ошибка создания аудиофайла:", err)
	}
	fileMP4, err := os.Create(mp4Filename)
	if err != nil {
		log.Fatalln("Ошибка создания mp4 файла:", err)
	}

	fileMu.Lock()
	currentVideoFile = fileVideo
	currentAudioFile = fileAudio
	currentMP4File = fileMP4

	// Начало нового сегмента – флаг сбрасывается и ожидается первый IDR
	currentSegmentStarted = false
	fileMu.Unlock()

	// Тикер для периодической ротации файлов
	ticker := time.NewTicker(period)
	defer ticker.Stop()

	// Горутина для смены файлов
	go func() {
		for range ticker.C {
			newTimestamp := time.Now().Format("15-04_02-01-2006")
			newVideoFilename := newTimestamp + ".h264"
			newAudioFilename := newTimestamp + ".pcm"
			newMP4Filename := newTimestamp + ".mp4"

			newVideoFile, err := os.Create(newVideoFilename)
			if err != nil {
				log.Printf("Ошибка создания видеофайла для фрагмента %s: %v", newTimestamp, err)
				continue
			}
			newAudioFile, err := os.Create(newAudioFilename)
			if err != nil {
				log.Printf("Ошибка создания аудиофайла для фрагмента %s: %v", newTimestamp, err)
				continue
			}
			newMP4File, err := os.Create(newMP4Filename)
			if err != nil {
				log.Printf("Ошибка создания mp4 файла для фрагмента %s: %v", newTimestamp, err)
			}

			// Переключение на новые файлы
			fileMu.Lock()
			oldVideoFile := currentVideoFile
			oldAudioFile := currentAudioFile
			oldMP4File := currentMP4File

			oldVideoFilename := videoFilename
			oldAudioFilename := audioFilename
			oldMP4Filename := mp4Filename

			currentVideoFile = newVideoFile
			currentAudioFile = newAudioFile
			currentMP4File = newMP4File

			// Новый сегмент ещё не начат – флаг сбрасывается и ожидается первый IDR
			currentSegmentStarted = false

			// Обновление имен для следующей итерации
			videoFilename = newVideoFilename
			audioFilename = newAudioFilename
			mp4Filename = newMP4Filename
			fileMu.Unlock()

			// Закрытие старых файлов
			oldVideoFile.Close()
			oldAudioFile.Close()
			oldMP4File.Close()

			log.Println("Созданы новые файлы для записи:", videoFilename, audioFilename, mp4Filename)

			// Горутина для объединения видео и аудио в контейнер MP4
			go func(oldVideoFilename, oldAudioFilename, oldMP4Filename string) {
				formaH264Mu.Lock()
				h264Params := formaH264
				formaH264Mu.Unlock()

				// Конвертация PCM в AAC
				aacFilename := strings.Replace(oldAudioFilename, ".pcm", ".aac", 1)
				if err := convertPCMtoAAC(oldAudioFilename, aacFilename); err != nil {
					log.Printf("Ошибка конвертации аудио: %v", err)
					return
				}

				log.Printf("Файл %s успешно конвертирован в %s\n", oldAudioFilename, aacFilename)

				// Создание MP4 файл
				if err := createMP4File(oldVideoFilename, aacFilename, oldMP4Filename, h264Params); err != nil {
					log.Printf("Ошибка создания MP4: %v", err)
				} else {
					log.Printf("Файлы %s, %s успешно объединены в контейнер MP4: %s\n",
						oldVideoFilename, aacFilename, oldMP4Filename)
				}

				// Удаление временных файлов
				os.Remove(oldVideoFilename)
				os.Remove(aacFilename)
				os.Remove(oldAudioFilename)

				log.Printf("Файлы %s, %s, %s успешно удалены\n", oldVideoFilename, aacFilename, oldAudioFilename)
			}(oldVideoFilename, oldAudioFilename, oldMP4Filename)
		}
	}()

	wg.Add(1)
	wg.Wait()
}

// Sample – единый тип для описания сэмпла (для видео и аудио)
type Sample struct {
	Data   []byte
	Offset uint64
	Size   uint32
	Dur    uint32
	IsSync bool
}

// convertPCMtoAAC конвертирует PCM-файл в AAC-файл.
func convertPCMtoAAC(aFilename, outputFile string) error {
	if outputFile == "" {
		outputFile = strings.Replace(aFilename, ".pcm", ".aac", 1)
	}

	audioFileAAC, err := os.Create(outputFile)
	if err != nil {
		log.Println("Ошибка создания файла AAC:", err)
		return err
	}
	defer audioFileAAC.Close()

	audioFilePCM, err := os.Open(aFilename)
	if err != nil {
		log.Println("Ошибка открытия PCM файла:", err)
		return err
	}
	defer audioFilePCM.Close()

	opts := &aac.Options{
		SampleRate:  8000,
		NumChannels: 1,
	}

	encoderAAC, err := aac.NewEncoder(audioFileAAC, opts)
	if err != nil {
		log.Println("Ошибка создания инкодера AAC:", err)
		return err
	}
	defer encoderAAC.Close()

	if err = encoderAAC.Encode(audioFilePCM); err != nil {
		log.Println("Ошибка инкодирования в AAC:", err)
		return err
	}

	if err = os.Remove(aFilename); err != nil {
		log.Println("Ошибка удаления PCM файла:", err)
	} else {
		log.Printf("PCM файл %s успешно удалён.", aFilename)
	}
	return nil
}

// createMP4File упаковывает видео (H.264) и аудио (AAC) файлы в MP4.
func createMP4File(videoFile, audioFile, outputFile string, formaH264 *format.H264) error {
	videoData, err := os.ReadFile(videoFile)
	if err != nil {
		return err
	}
	videoSamples := parseH264Samples(videoData)

	audioData, err := os.ReadFile(audioFile)
	if err != nil {
		return err
	}
	audioSamples := parseAACSamples(audioData)

	f := mp4.NewFile()

	videoTrack := createVideoTrack(formaH264.SPS, formaH264.PPS)
	audioTrack := createAudioTrack(8000, 1)

	addSamplesToVideoTrack(videoTrack, videoSamples)
	addSamplesToAudioTrack(audioTrack, audioSamples, 8000)

	moov := mp4.NewMoovBox()
	moov.AddChild(videoTrack)
	moov.AddChild(audioTrack)
	f.AddChild(moov, 0)

	mdat := mp4.MdatBox{}
	mdat.Data = append(videoData, audioData...)
	f.AddChild(&mdat, 0)

	outFile, err := os.Create(outputFile)
	if err != nil {
		return err
	}
	defer outFile.Close()
	return f.Encode(outFile)
}

// createVideoTrack создаёт видео-трек на основе SPS и PPS.
func createVideoTrack(sps, pps []byte) *mp4.TrakBox {
	sps = stripAnnexB(sps)
	pps = stripAnnexB(pps)

	if len(sps) < 4 {
		log.Fatalln("SPS слишком короткий или пустой")
	}
	if len(pps) < 1 {
		log.Fatalln("PPS слишком короткий или пустой")
	}

	trak := mp4.NewTrakBox()
	timescale := uint32(90000)

	mdhd := &mp4.MdhdBox{
		Timescale: timescale,
		Language:  convertLanguage("und"),
	}

	hdlr := &mp4.HdlrBox{
		HandlerType: "vide",
		Name:        "VideoHandler",
	}

	avc1 := mp4.NewVisualSampleEntryBox("avc1")
	avc1.Width = 1280
	avc1.Height = 720

	avcC, err := mp4.CreateAvcC([][]byte{sps}, [][]byte{pps}, true)
	if err != nil {
		log.Fatalf("ошибка создания avcC: %v", err)
	}
	avcC.AVCLevelIndication = 0x1f
	avc1.AddChild(avcC)

	stbl := mp4.NewStblBox()
	stsd := mp4.NewStsdBox()
	stsd.AddChild(avc1)
	stbl.AddChild(stsd)
	stbl.AddChild(&mp4.SttsBox{})
	stbl.AddChild(&mp4.StscBox{})
	stbl.AddChild(&mp4.StszBox{})
	stbl.AddChild(&mp4.StcoBox{})

	trak.Mdia = &mp4.MdiaBox{
		Mdhd: mdhd,
		Hdlr: hdlr,
		Minf: &mp4.MinfBox{
			Stbl: stbl,
		},
	}

	return trak
}

// stripAnnexB удаляет стартовый код Annex-B из NAL-единицы.
func stripAnnexB(nalu []byte) []byte {
	if len(nalu) >= 4 && bytes.Equal(nalu[:4], []byte{0x00, 0x00, 0x00, 0x01}) {
		return nalu[4:]
	}
	if len(nalu) >= 3 && bytes.Equal(nalu[:3], []byte{0x00, 0x00, 0x01}) {
		return nalu[3:]
	}
	return nalu
}

// createAudioTrack создаёт аудио-трек.
func createAudioTrack(sampleRate, channels int) *mp4.TrakBox {
	trak := mp4.NewTrakBox()
	timescale := uint32(sampleRate)

	mdhd := &mp4.MdhdBox{
		Timescale: timescale,
		Language:  convertLanguage("und"),
	}

	hdlr := &mp4.HdlrBox{
		HandlerType: "soun",
		Name:        "SoundHandler",
	}

	mp4a := mp4.NewAudioSampleEntryBox("mp4a")
	mp4a.ChannelCount = uint16(channels)
	mp4a.SampleRate = uint16(sampleRate)
	mp4a.SampleSize = 16

	asc := getAACConfig(sampleRate, channels)
	esds := createESDSBox(asc)
	mp4a.AddChild(esds)

	stbl := mp4.NewStblBox()
	stsd := mp4.NewStsdBox()
	stsd.AddChild(mp4a)
	stbl.AddChild(stsd)
	stbl.AddChild(&mp4.SttsBox{})
	stbl.AddChild(&mp4.StscBox{})
	stbl.AddChild(&mp4.StszBox{})
	stbl.AddChild(&mp4.StcoBox{})

	trak.Mdia = &mp4.MdiaBox{
		Mdhd: mdhd,
		Hdlr: hdlr,
		Minf: &mp4.MinfBox{
			Stbl: stbl,
		},
	}

	return trak
}

// createESDSBox создаёт ESDS-бокс для AAC.
func createESDSBox(asc []byte) *mp4.EsdsBox {
	return &mp4.EsdsBox{
		ESDescriptor: mp4.ESDescriptor{
			DecConfigDescriptor: &mp4.DecoderConfigDescriptor{
				ObjectType: 0x40,
				StreamType: 0x05,
				DecSpecificInfo: &mp4.DecSpecificInfoDescriptor{
					DecConfig: asc,
				},
			},
		},
	}
}

func getAACConfig(sampleRate, channels int) []byte {
	// Пример конфигурации для 8000 Гц, моно
	return []byte{0x12, 0x10}
}

// addSamplesToVideoTrack заполняет таблицы сэмплов видео-трека.
func addSamplesToVideoTrack(trak *mp4.TrakBox, samples []Sample) {
	stbl := trak.Mdia.Minf.Stbl

	stts := mp4.SttsBox{}
	for _, s := range samples {
		stts.SampleCount = []uint32{1}
		stts.SampleTimeDelta = []uint32{s.Dur}
	}

	stsz := mp4.StszBox{}
	for _, s := range samples {
		stsz.SampleSize = []uint32{s.Size}
	}

	stss := mp4.StssBox{}
	for i, s := range samples {
		if s.IsSync {
			stss.SampleNumber = []uint32{uint32(i + 1)}
		}
	}

	stbl.AddChild(&stts)
	stbl.AddChild(&stsz)
	stbl.AddChild(&stss)
}

// addSamplesToAudioTrack заполняет таблицы сэмплов аудио-трека.
func addSamplesToAudioTrack(trak *mp4.TrakBox, samples []Sample, sampleRate int) {
	stbl := trak.Mdia.Minf.Stbl

	stts := mp4.SttsBox{}
	for _, s := range samples {
		stts.SampleCount = []uint32{1}
		stts.SampleTimeDelta = []uint32{s.Dur}
	}

	stsz := mp4.StszBox{}
	for _, s := range samples {
		stsz.SampleSize = []uint32{uint32(len(s.Data))}
	}

	stbl.AddChild(&stts)
	stbl.AddChild(&stsz)
}

// parseH264Samples парсит H.264 поток на сэмплы.
func parseH264Samples(data []byte) []Sample {
	var samples []Sample
	start := 0
	for i := 0; i < len(data)-3; i++ {
		if bytes.Equal(data[i:i+4], []byte{0x00, 0x00, 0x00, 0x01}) {
			if start != 0 {
				sampleData := data[start:i]
				samples = append(samples, Sample{
					Size:   uint32(len(sampleData)),
					Dur:    3600,
					IsSync: isKeyFrame(sampleData),
				})
			}
			start = i + 4
		}
	}
	if start < len(data) {
		sampleData := data[start:]
		samples = append(samples, Sample{
			Size:   uint32(len(sampleData)),
			Dur:    3600,
			IsSync: isKeyFrame(sampleData),
		})
	}
	return samples
}

// isKeyFrame определяет, является ли NALU ключевым (IDR).
func isKeyFrame(nalu []byte) bool {
	if len(nalu) == 0 {
		return false
	}
	nalType := nalu[0] & 0x1F
	return nalType == 5
}

// parseAACSamples парсит AAC данные и возвращает сэмплы.
func parseAACSamples(data []byte) []Sample {
	var samples []Sample
	frameSize := 1024 // примерный размер AAC-фрейма
	for i := 0; i < len(data); i += frameSize {
		end := i + frameSize
		if end > len(data) {
			end = len(data)
		}
		samples = append(samples, Sample{
			Data: data[i:end],
			Dur:  1024,
			Size: uint32(end - i),
		})
	}
	return samples
}

// convertLanguage преобразует строку языка в 16-битное значение.
func convertLanguage(lang string) uint16 {
	if len(lang) != 3 {
		return 0x55C4
	}
	b1 := lang[0] - 0x60
	b2 := lang[1] - 0x60
	b3 := lang[2] - 0x60
	return uint16((b1 << 10) | (b2 << 5) | b3)
}