package main

import (
	"context"
	"fmt"
	"log"
	"os"
	"os/signal"
	"sync"
	"syscall"
	"time"

	"github.com/gen2brain/malgo"
	"github.com/joho/godotenv"
	"github.com/lokutor-ai/lokutor-orchestrator/pkg/audio"
	"github.com/lokutor-ai/lokutor-orchestrator/pkg/orchestrator"
	llmProvider "github.com/lokutor-ai/lokutor-orchestrator/pkg/providers/llm"
	noiseFilter "github.com/lokutor-ai/lokutor-orchestrator/pkg/providers/noise"
	sttProvider "github.com/lokutor-ai/lokutor-orchestrator/pkg/providers/stt"
	ttsProvider "github.com/lokutor-ai/lokutor-orchestrator/pkg/providers/tts"
)

const (
	SampleRate = 44100
	Channels   = 1
)

func main() {

	if err := godotenv.Load(); err != nil {
		log.Println("Note: No .env file found, using system environment variables")
	}

	groqKey := os.Getenv("GROQ_API_KEY")
	openaiKey := os.Getenv("OPENAI_API_KEY")
	anthropicKey := os.Getenv("ANTHROPIC_API_KEY")
	googleKey := os.Getenv("GOOGLE_API_KEY")
	geminiKey := os.Getenv("GEMINI_API_KEY")
	deepgramKey := os.Getenv("DEEPGRAM_API_KEY")
	assemblyKey := os.Getenv("ASSEMBLYAI_API_KEY")
	lokutorKey := os.Getenv("LOKUTOR_API_KEY")

	sttProviderName := os.Getenv("STT_PROVIDER")
	if sttProviderName == "" {
		sttProviderName = "groq"
	}
	llmProviderName := os.Getenv("LLM_PROVIDER")
	if llmProviderName == "" {
		// Default to groq unless GEMINI_API_KEY is set, then prefer Google
		if geminiKey != "" {
			llmProviderName = "google"
		} else {
			llmProviderName = "groq"
		}
	}

	lang := orchestrator.Language(os.Getenv("AGENT_LANGUAGE"))
	if lang == "" {
		lang = orchestrator.LanguageEs
	}

	if lokutorKey == "" {
		log.Fatal("Error: LOKUTOR_API_KEY must be set.")
	}

	var stt orchestrator.STTProvider
	switch sttProviderName {
	case "openai":
		if openaiKey == "" {
			log.Fatal("Error: OPENAI_API_KEY must be set for openai STT")
		}
		stt = sttProvider.NewOpenAISTT(openaiKey, "whisper-1")
	case "deepgram":
		if deepgramKey == "" {
			log.Fatal("Error: DEEPGRAM_API_KEY must be set for deepgram STT")
		}
		stt = sttProvider.NewDeepgramSTT(deepgramKey)
	case "assemblyai":
		if assemblyKey == "" {
			log.Fatal("Error: ASSEMBLYAI_API_KEY must be set for assemblyai STT")
		}
		stt = sttProvider.NewAssemblyAISTT(assemblyKey)
	case "groq":
		fallthrough
	default:
		if groqKey == "" {
			log.Fatal("Error: GROQ_API_KEY must be set for groq STT")
		}
		groqModel := os.Getenv("GROQ_STT_MODEL")
		if groqModel == "" {
			groqModel = "whisper-large-v3-turbo"
		}
		stt = sttProvider.NewGroqSTT(groqKey, groqModel)
	}

	if s, ok := stt.(interface{ SetSampleRate(int) }); ok {
		s.SetSampleRate(SampleRate)
	}

	// Wrap STT with noise suppression filter if model is available
	noiseModelPath := os.Getenv("NOISE_FILTER_MODEL")
	if noiseModelPath == "" {
		// Default path relative to the agent's working directory (lokutor-orchestrator)
		noiseModelPath = "../assets/onnx/versa-1.0/noise_suppressor_v2.onnx"
	}
	log.Printf("Noise filter model path: %s", noiseModelPath)
	if _, err := os.Stat(noiseModelPath); err == nil {
		wrappedSTT, err := noiseFilter.NewSTTWrapper(stt, noiseModelPath, SampleRate)
		if err != nil {
			log.Printf("Warning: failed to load noise filter (%v), using unfiltered STT", err)
		} else {
			stt = wrappedSTT
			log.Printf("Noise filter ACTIVE: model=%s provider=%s", noiseModelPath, wrappedSTT.Name())
		}
	} else {
		log.Printf("Noise filter model NOT FOUND at %s (err: %v), continuing without noise filter", noiseModelPath, err)
	}

	var llm orchestrator.LLMProvider
	switch llmProviderName {
	case "openai":
		if openaiKey == "" {
			log.Fatal("Error: OPENAI_API_KEY must be set for openai LLM")
		}
		llm = llmProvider.NewOpenAILLM(openaiKey, "gpt-4o")
	case "anthropic":
		if anthropicKey == "" {
			log.Fatal("Error: ANTHROPIC_API_KEY must be set for anthropic LLM")
		}
		llm = llmProvider.NewAnthropicLLM(anthropicKey, "claude-3-5-sonnet-20241022")
	case "google":
		// Prefer GEMINI_API_KEY, fall back to GOOGLE_API_KEY
		if geminiKey == "" {
			geminiKey = googleKey
		}
		if geminiKey == "" {
			log.Fatal("Error: GEMINI_API_KEY or GOOGLE_API_KEY must be set for google LLM")
		}
		geminiModel := os.Getenv("GEMINI_MODEL")
		if geminiModel == "" {
			geminiModel = "gemini-2.5-flash"
		}
		llm = llmProvider.NewGoogleLLM(geminiKey, geminiModel)
	case "groq":
		fallthrough
	default:
		if groqKey == "" {
			log.Fatal("Error: GROQ_API_KEY must be set for groq LLM")
		}
		llm = llmProvider.NewGroqLLM(groqKey, "meta-llama/llama-4-scout-17b-16e-instruct")
	}

	config := orchestrator.DefaultConfig()
	config.Language = lang
	config.SilenceTimeout = 10 * time.Second

	firstSpeaker := os.Getenv("FIRST_SPEAKER")
	if firstSpeaker == "bot" {
		config.FirstSpeaker = orchestrator.FirstSpeakerBot
	} else if firstSpeaker == "user" {
		config.FirstSpeaker = orchestrator.FirstSpeakerUser
	}

	fmt.Printf("Configured: STT=%s | LLM=%s | TTS=Lokutor\n", sttProviderName, llmProviderName)
	
	// Show LLM model details
	if llmProviderName == "google" {
		geminiModel := os.Getenv("GEMINI_MODEL")
		if geminiModel == "" {
			geminiModel = "gemini-2.5-flash"
		}
		fmt.Printf("  → Using Gemini model: %s\n", geminiModel)
	} else if llmProviderName == "groq" {
		fmt.Printf("  → Using Groq model: meta-llama/llama-4-scout-17b-16e-instruct\n")
	} else if llmProviderName == "openai" {
		fmt.Printf("  → Using OpenAI model: gpt-4o\n")
	} else if llmProviderName == "anthropic" {
		fmt.Printf("  → Using Anthropic model: claude-3-5-sonnet-20241022\n")
	}
	
	fmt.Printf("VAD: ImprovedRMS | Sample Rate: %dHz | Language: %s\n", SampleRate, lang)
	fmt.Printf("Features: clientVAD=%v tokenTTS=%v specLLM=%v pacing=%v cache=%v sum=%v pool=%d\n",
		config.ClientVAD, config.TokenLevelTTS, config.SpeculativeLLM,
		config.AdaptivePacing, config.ResponseCaching, config.ContextSummarization,
		config.TTSConnectionPoolSize)
	fmt.Println("Voice Agent Started! Listening to microphone...")
	fmt.Println("Press Ctrl+C to exit")

	var tts *ttsProvider.LokutorTTS
	if config.TTSConnectionPoolSize > 1 {
		tts = ttsProvider.NewLokutorTTSPool(lokutorKey, config.TTSConnectionPoolSize)
	} else {
		tts = ttsProvider.NewLokutorTTS(lokutorKey)
	}

	// Advanced VAD with ZCR and Peak tracking for better noise rejection.
	vad := orchestrator.NewImprovedRMSVAD(config.BargeInVADThreshold, 200*time.Millisecond, SampleRate)
	// Wait for ~80ms of continuous energy before muting (8 * 10ms frame size).
	// This filters out tiny spikes (keyboard clicks, desk bumps) while remaining fast.
	vad.SetMinConfirmed(2)

	orch := orchestrator.NewWithVAD(stt, llm, tts, vad, config)

	// Register an example tool so the agent can check the time
	orch.RegisterTool("get_time", func(args string) (string, error) {
		return fmt.Sprintf("The current time is %s", time.Now().Format("15:04:05")), nil
	})


	session := orch.NewSessionWithDefaults("user_123")

	// Let the LLM know the get_time tool is available
	session.SetTools([]orchestrator.Tool{
		{
			Type: "function",
			Function: map[string]interface{}{
				"name":        "get_time",
				"description": "Get the current time in the user's timezone",
				"parameters": map[string]interface{}{
					"type":       "object",
					"properties": map[string]interface{}{},
				},
			},
		},
	})


	systemPrompt := "You are a helpful and concise voice assistant. " +
		"This is a real-time conversational phone call. " +
		"Use short sentences suitable for speech. " +
		"If the user is silent for a long time, you will receive '[USER_SILENCE_TIMEOUT]'. In that case, briefly check if they are still there or if they need anything else."
	if lang == orchestrator.LanguageEs {
		systemPrompt = "Eres un asistente de voz útil y conciso. " +
			"Esta es una llamada telefónica conversacional en tiempo real. " +
			"Usa frases cortas adecuadas para el habla. " +
			"Si el usuario está en silencio por mucho tiempo, recibirás '[USER_SILENCE_TIMEOUT]'. En ese caso, pregunta brevemente si siguen ahí o si necesitan algo más."
	}
	orch.SetSystemPrompt(session, systemPrompt)

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	stream := orch.NewManagedStream(ctx, session)
	defer stream.Close()

	mctx, err := malgo.InitContext(nil, malgo.ContextConfig{}, nil)
	if err != nil {
		log.Fatal(err)
	}
	defer mctx.Uninit()

	var playbackMu sync.Mutex
	var playbackBytes []byte
	var e2eLogged bool
	var preRolling bool = true
	var playbackPaused bool
	const preRollSize = 44100 * 2 * 60 / 1000 // 60ms pre-roll

	// Audio enhancer for better sound quality
	audioEnhancer := audio.NewProcessor(audio.Config{
		TargetLUFS:      -18,
		LowShelfGain:    1.5,
		HighShelfGain:   0,
		PresenceGain:    0.5,
		HarmonicMix:     0.05,
		ReverbMix:       0.02,
		CompressRatio:   2,
		CompressThreshold: -24,
	})

	chunkPool := sync.Pool{
		New: func() interface{} {
			return make([]byte, 8192)
		},
	}

	inputChan := make(chan []byte, 512)
	go func() {
		for chunk := range inputChan {
			_ = stream.Write(chunk)
			chunkPool.Put(chunk)
		}
	}()



	onSamples := func(pOutput, pInput []byte, frameCount uint32) {
		if pInput != nil {
			buf := chunkPool.Get().([]byte)
			buf = buf[:cap(buf)] // Reset slice to full capacity so copy doesn't truncate data
			n := copy(buf, pInput)
			// Must block channel write to guarantee every sample gets cleanly added
			// instead of dropping samples and permanently breaking byte-phase
			inputChan <- buf[:n]
		}

		if pOutput != nil {
			bytesToRead := int(frameCount) * 2
			playbackMu.Lock()

			if preRolling {
				if len(playbackBytes) < preRollSize {
					playbackMu.Unlock()
					for i := range pOutput {
						pOutput[i] = 0
					}
					return
				}
				preRolling = false
			}

			playedRealAudio := false

			if len(playbackBytes) == 0 || playbackPaused {
				for i := range pOutput {
					pOutput[i] = 0
				}
			} else {
				playedRealAudio = true
				if len(playbackBytes) < bytesToRead {
					copy(pOutput, playbackBytes)
					for i := len(playbackBytes); i < bytesToRead; i++ {
						pOutput[i] = 0
					}
					playbackBytes = nil
				} else {
					copy(pOutput, playbackBytes[:bytesToRead])
					playbackBytes = playbackBytes[bytesToRead:]
				}
			}

			if !preRolling {
				if playedRealAudio && !e2eLogged {
					bd := stream.GetLatencyBreakdown()
					if bd.UserToPlay > 0 {
						fmt.Printf("\r\033[K⏱️  [REPORT] Pipe_Ready: %dms | STT_Net: %dms | LLM_1st_Net: %dms | TTS_1st_Net: %dms\n", 
							bd.BotStartLatency, bd.STT, bd.LLM, bd.LLMToTTSFirstByte)
						probStr := ""
						if bd.NoSpeechProb > 0.01 {
							probStr = fmt.Sprintf(" | Prob: %.2f", bd.NoSpeechProb)
						}
						fmt.Printf("\r\033[K⏱️  [E2E] User_Stop -> Play: %dms | TTS_Total: %dms%s\n",
							bd.UserToPlay, bd.TTSTotal, probStr)
						e2eLogged = true
					}
				}
				}
			playbackMu.Unlock()
		}
	}

	deviceConfig := malgo.DefaultDeviceConfig(malgo.Duplex)
	deviceConfig.Capture.Format = malgo.FormatS16
	deviceConfig.Capture.Channels = 1
	deviceConfig.Playback.Format = malgo.FormatS16
	deviceConfig.Playback.Channels = 1
	deviceConfig.SampleRate = SampleRate
	deviceConfig.Alsa.NoMMap = 1

	device, err := malgo.InitDevice(mctx.Context, deviceConfig, malgo.DeviceCallbacks{
		Data: onSamples,
	})
	if err != nil {
		log.Fatal(err)
	}
	defer device.Uninit()

	if err := device.Start(); err != nil {
		log.Fatal(err)
	}

	go func() {
		currentGeneration := 0
		botHasAudio := false
		for event := range stream.Events() {
			switch event.Type {
			case orchestrator.UserSpeaking:
				playbackMu.Lock()
				playbackPaused = true
				playbackMu.Unlock()
				e2eLogged = false
				rms := stream.LastRMS()
				fmt.Printf("\r\033[K🎤 [USER] Speaking... (RMS: %.4f)\n", rms)

			case orchestrator.Interrupted:
				playbackMu.Lock()
				if botHasAudio {
					// Bot was speaking - this was just brief noise, don't kill playback
					fmt.Printf("\r\033[K🔄 [INTERRUPTED] Brief noise - resuming bot playback.\n")
				} else {
					playbackBytes = nil
					preRolling = true
					currentGeneration = event.Generation
					fmt.Printf("\r\033[K🛑 [INTERRUPTED] User started talking (gen: %d).\n", currentGeneration)
				}
				botHasAudio = false
				playbackPaused = false
				playbackMu.Unlock()
			case orchestrator.BotResumed:
				playbackMu.Lock()
				playbackPaused = false
				playbackMu.Unlock()
				fmt.Printf("\r\033[K🔄 [RESUMED] User input was noise.\n")



			case orchestrator.UserStopped:
				fmt.Printf("\r\033[K⌛ [STT] Processing...\n")
			case orchestrator.TranscriptFinal:
				fmt.Printf("\r\033[K📝 [TRANSCRIPT] %s\n", event.Data.(string))
				raw, proc := stream.ExportLastUserAudio()
				if raw != nil {
					ts := time.Now().Format("20060102-150405")
					rawPath := fmt.Sprintf("/tmp/lokutor_user_raw_%s.wav", ts)
					procPath := fmt.Sprintf("/tmp/lokutor_user_processed_%s.wav", ts)
					_ = os.WriteFile(rawPath, audio.NewWavBuffer(raw, SampleRate), 0644)
					_ = os.WriteFile(procPath, audio.NewWavBuffer(proc, SampleRate), 0644)
					fmt.Printf("\r\033[K💾 Saved user audio: %s (raw), %s (processed)\n", rawPath, procPath)
				}

			case orchestrator.BotThinking:
				playbackMu.Lock()
				playbackBytes = nil
				playbackPaused = true
				currentGeneration = event.Generation
				playbackMu.Unlock()
				fmt.Printf("\r\033[K🧠 [LLM] Thinking... (gen: %d)\n", currentGeneration)

			case orchestrator.ToolCall:
				if tc, ok := event.Data.(orchestrator.ToolCallEventData); ok {
					fmt.Printf("\r\033[K🛠️ [TOOL] Calling %s(%s)...\n", tc.Name, tc.Arguments)
				}

			case orchestrator.BotResponse:
				if resp, ok := event.Data.(string); ok {
					fmt.Printf("\r\033[K💬 [AGENT] %s\n", resp)
				}
			case orchestrator.BotSpeaking:
				playbackMu.Lock()
				playbackPaused = false
				playbackMu.Unlock()
				
				if !e2eLogged {
					bd := stream.GetLatencyBreakdown()
					if bd.UserToPlay > 0 {
						fmt.Printf("\r\033[K🔊 [TTS] Speaking... (gen: %d)\n", event.Generation)
						fmt.Printf("\r\033[K⏱️  [REPORT] Pipe_Ready: %dms | STT_Net: %dms | LLM_1st_Net: %dms | TTS_1st_Net: %dms\n", 
							bd.BotStartLatency, bd.STT, bd.LLM, bd.LLMToTTSFirstByte)
						fmt.Printf("\r\033[K⏱️  [E2E] User_Stop -> Play: %dms | TTS_Total: %dms | Prob: %.2f\n",
							bd.UserToPlay, bd.TTSTotal, bd.NoSpeechProb)
						e2eLogged = true
					}
				}
			case orchestrator.AudioChunk:
				if event.Generation < currentGeneration {
					continue
				}
				chunk := event.Data.([]byte)
				// Apply audio enhancement for better sound quality
				enhanced := audioEnhancer.Process(chunk, SampleRate, 1)
				playbackMu.Lock()
				playbackBytes = append(playbackBytes, enhanced...)
				botHasAudio = true
				playbackMu.Unlock()
				fmt.Printf("\r\033[K🔊 [AUDIO] received %d bytes, total buffered: %d\n", len(enhanced), len(playbackBytes))

			case orchestrator.ErrorEvent:
				fmt.Printf("\r\033[K❌ [ERROR] %v\n", event.Data)
			}
		}
	}()

	sig := make(chan os.Signal, 1)
	signal.Notify(sig, syscall.SIGINT, syscall.SIGTERM)
	sigReceived := <-sig
	fmt.Printf("\nShutting down (signal=%v)...\n", sigReceived)

	_ = device.Stop()
	stream.Close()
	time.Sleep(50 * time.Millisecond)
}
