package handlers

import (
    "encoding/binary"
    "encoding/json"
    "fmt"
    "net"
    "regexp"
    "sync"

    "rtp-app/pkg/rtp"

    vosk "github.com/alphacep/vosk-api/go"
)

const (
    Rate           = 16000
    bufSize        = 2048
)

type VoskResult struct {
    Text string `json:"text"`
}

func HandleConnection(conn *net.UDPConn, wg *sync.WaitGroup) {
    model, err := vosk.NewModel("model-ru")
    if err != nil {
        fmt.Printf("Не удалось открыть модель Vosk: %v", err)
    }
    recognizer, err := vosk.NewRecognizer(model, float64(Rate))
    if err != nil {
        fmt.Printf("Не удалось запустить модель: %v", err)
    }
    defer wg.Done()
    buffer := make([]byte, bufSize)
    mu := sync.Mutex{}
    re := regexp.MustCompile(`[\.\:]`)
    for {
        n, addr, err := conn.ReadFromUDP(buffer)
        if err != nil {
            fmt.Println("Ошибка при чтении данных:", err)
            return
        }
        mu.Lock()
        CallID := fmt.Sprintf("%s-%d", re.ReplaceAllString(addr.IP.String(), "-"), addr.Port)
        mu.Unlock()
        wg.Add(1)
        processRTP(buffer[:n], CallID, wg, recognizer)
    }
}

func processRTP(data []byte, CallID string, wg *sync.WaitGroup, recognizer *vosk.VoskRecognizer) {
    packet := &rtp.Packet{}
    if err := packet.Unmarshal(data); err != nil {
        fmt.Printf("Ошибка при декодировании RTP пакета от %s: %v\n", CallID, err)
        return
    }
    if recognizer.AcceptWaveform(BeToLe(packet.Payload)) != 0 {
        resultStr := recognizer.Result()
        var result VoskResult
        if err := json.Unmarshal([]byte(resultStr), &result); err != nil {
            fmt.Printf("Ошибка демаршалинга результата: %v", err)
        }
        if result.Text != "" {
            fmt.Printf("Полная транскрипция: %s\t%s\n", CallID, result.Text)
        }
    } else {
        partialResultStr := recognizer.PartialResult()
        var partialResult VoskResult
        if err := json.Unmarshal([]byte(partialResultStr), &partialResult); err != nil {
            fmt.Printf("Ошибка демаршалинга частичного результата: %v", err)
            return
        }
        if partialResult.Text != "" {
            fmt.Printf("Частичная транскрипция: %s\t%s\n", CallID, partialResult.Text)
        }
    }
}

func BeToLe(input []byte) []byte {
    if len(input)%4 != 0 {
        panic("Длина входных данных должна быть кратна 4")
    }
    output := make([]byte, len(input))
    for i := 0; i < len(input); i += 4 {
        binary.LittleEndian.PutUint32(output[i:i+4], binary.BigEndian.Uint32(input[i:i+4]))
    }
    return output
}