main.py

import os
import io
import wave
import httpx
import json
from fastapi import FastAPI, Depends, HTTPException, File, UploadFile, Query, Header
from fastapi.responses import Response
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy import text
from database import get_db
from db_utils import check_user_preferences, resolve_semantic_location
from sqlalchemy.future import select

from dotenv import load_dotenv
load_dotenv()

from vad_utils import apply_vad_filter
from llm_utils import process_with_llm
from prompts import (
	SYSTEM_PROMPT,
	INTENTS_PROMPT,
	PREFERENCE_INTENTS_PROMPT,
	NAVIGATION_INTENTS_PROMPT,
	CLARIFICATIONS_PROMPT,
	TRIP_CHANGES_PROMPT,
	SEMANTICS_PROMPT
	)

# --- State management for user sessions ---
class SessionState:
	def __init__ (self):
		self.conversation_history = []
		self.current_location = None # GPS

		self.turn_count = 0 # Number of user turns taken in the current session
		self.current_route_params = {} # The active routing parameters (origin, destination, via, avoid)
		self.final_llm_response = None # The final JSON response from the LLM containing the parsed trip details
		self.route_created = False # Flag indicating route generation completion
		self.pending_preference = True # Tracks if user preference confirmation is needed and what type (starts as True)
		self.navigation_started = False # Indicates active navigation state

		self.semantic_context = {
			"origin_label": None, "destination_label": None, # Labels include work, home, school, etc.
			"origin_known": False, "destination_known": False,
			"origin_value": None, "destination_value": None # Actual coordinates/addresses
		}
	
	def clear_trip_context(self):
		self.current_route_params = {}
		self.final_llm_response = None
		self.route_created = False
		self.pending_preference = True
		self.navigation_started = False

		self.semantic_context = {
			"origin_label": None, "destination_label": None,
			"origin_known": False, "destination_known": False,
			"origin_value": None, "destination_value": None
		}

SESSIONS = {}

# --- ASR server handoff configuration ---
ASR_URL = "http://172.16.3.217:80/transcribe"
ASR_API_KEY = os.getenv("ASR_API_KEY")
if not ASR_API_KEY:
	raise ValueError("CRITICAL ERROR: ASR_API_KEY is missing from the environment variables.")

app = FastAPI(title="HeyRoute API")

@app.get("/health")
async def health_check():
	return {"status": "online", "message": "The HeyRoute FastAPI server is running!"}

@app.get("/db-check")
async def db_check(db: AsyncSession = Depends(get_db)):
	try:
		await db.execute(text("SELECT 1"))
		return {"database_status": "connected", "message": "Successfully communicating with postgresql"}
	except Exception as e:
		raise HTTPException(status_code=500, detail=f"Database connection failed: {str(e)}")

# --- Voice processing endpoint ---
# This endpoint handles the flow of audio file processing, namely:
# 1. App sends audio file to this endpoint (DONE)
# 2. This endpoint applies VAD to clean the audio (DONE)
# 3. This endpoint sends the cleaned audio to the ASR server for transcription (DONE)
# 4. This endpoint receives the transcription from the ASR server (DONE)
# 5. This endpoint sends the transcription to the LLM server for intent extraction (DONE)
# 6. This endpoint receives the intent from the LLM server (DONE)
# 7. This endpoint sends the transcription back to the LLM for further processing (destination, preferences, etc.) (TODO)
# 8. This endpoint receives the final JSON navigation payload from the LLM server (TODO)
# 9. This endpoint sends the JSON navigation payload back to the App (TODO)

@app.post("/api/voice/vad")
async def process_voice_activity(
	audio_file: UploadFile = File(...), 
	download: bool = Query(False, description="Set to true to download the cleaned audio file"),
	x_user_id: str = Header(..., description="User ID for database querying"),
	x_session_id: str = Header(..., description="Session ID for the current routing trip"),
	x_current_lat: float = Header(None, description="Current latitude"),
	x_current_lng: float = Header(None, description="Current longitude"),
	db: AsyncSession = Depends(get_db)
	):

	print(f"Processing audio file: {audio_file.filename} for user: {x_user_id}, session: {x_session_id}")
	
	try:
		# Read the uploaded audio file
		audio_bytes = await audio_file.read()

		# Run the raw audio through the VAD filter utility
		clean_audio_bytes = apply_vad_filter(audio_bytes)

		# Convert the cleaned audio bytes to a WAV format for ASR processing
		wav_io = io.BytesIO()
		with wave.open(wav_io, 'wb') as wav_file:
			wav_file.setnchannels(1)  # Mono
			wav_file.setsampwidth(2)  # 16-bit samples
			wav_file.setframerate(16000)  # 16 kHz sample rate
			wav_file.writeframes(clean_audio_bytes)

		# Reset the BytesIO stream position to the beginning for reading
		wav_io.seek(0)

		# If the user requested to download the cleaned audio, return it as a response (this skips the ASR step)
		if download:
			return Response(
				content = wav_io.getvalue(),
				media_type = "audio/wav",
				headers = {"Content-Disposition": f"attachment; filename=clean_{audio_file.filename}.wav"}
			)
		
		# Send the cleaned audio to the ASR server for transcription
		async with httpx.AsyncClient(timeout=30.0) as client:
			files = {'file': (f"clean_{audio_file.filename}.wav", wav_io, "audio/wav")}
			headers = {'X-API-Key': ASR_API_KEY}

			asr_response = await client.post(ASR_URL, files=files, headers=headers)

			# Check if the ASR server responded with an error
			if asr_response.status_code != 200:
				raise HTTPException(status_code = asr_response.status_code,
									detail = f"ASR Server Error: {asr_response.text}"
				)
			
			# Extract the transcription data from the ASR server's response
			transcription_data = asr_response.json()
			user_text = transcription_data.get("text", "").strip()
		
		# ----------------------------------------------------------------------------------------------
		# LLM INTERACTION SECTION BELOW ----------------------------------------------------------------
		# ----------------------------------------------------------------------------------------------

		# --- LLM Intent Extraction ---
		# Initialize or load the user session state
		if x_session_id not in SESSIONS:
			SESSIONS[x_session_id] = SessionState()
		state = SESSIONS[x_session_id]

		if x_current_lat is not None and x_current_lng is not None:
			state.current_location = {"lat": x_current_lat, "lng": x_current_lng}
			print(f"[DEBUG] Current location for session {x_session_id}: {state.current_location}")

		# Add the ASR transcription to the conversation history
		state.conversation_history.append({"role": "user", "content": user_text})

		# Selecting the intent mode
		intent_mode = "SETUP"
		if state.pending_preference != True and state.pending_preference != False and state.final_llm_response:
			intent_mode = "PREFERENCE_CONFIRMATION"
		elif state.navigation_started:
			intent_mode = "NAVIGATION"

		# Build the intent prompt based on intent mode
		if intent_mode == "PREFERENCE_CONFIRMATION":
			intent_prompt = [
				{"role": "system", "content": PREFERENCE_INTENTS_PROMPT},
				{"role": "user", "content": (
					"The conversation so far:\n" +
					"\n".join(f"{m['role']}: {m['content']}" for m in state.conversation_history)
				)}
			]
		elif intent_mode == "NAVIGATION":
			intent_prompt = [
				{"role": "system", "content": NAVIGATION_INTENTS_PROMPT},
				{"role": "user", "content": (
					"The user is currently in navigation mode.\nThe conversation so far:\n" +
					"\n".join(f"{m['role']}: {m['content']}" for m in state.conversation_history)
				)}
			]
		else: # Intent mode = SETUP
			intent_prompt = [
				{"role": "system", "content": INTENTS_PROMPT},
				{"role": "system", "content": f"Known semantic places: {state.semantic_context}"},
				{"role": "user", "content": (
					"The conversation so far:\n" +
					"\n".join(f"{m['role']}: {m['content']}" for m in state.conversation_history)
				)}
			]

		# Process the intent
		detected_intents = await process_with_llm(intent_prompt)
		print(f"[DEBUG] Intent mode: {intent_mode}, Detected intents for session {x_session_id}: {detected_intents}")

		# --- Conversational routing tree ---
		response_text = ""
		final_travel_json = state.final_llm_response

		# 1. Clarifications
		if detected_intents.get("clarifications"):
			clarification_prompt = [
                {"role": "system", "content": SYSTEM_PROMPT},
                {"role": "system", "content": CLARIFICATIONS_PROMPT},
                {"role": "user", "content": (
                    "The conversation so far:\n" +
                    "\n".join(f"{m['role']}: {m['content']}" for m in state.conversation_history)
                )}
            ]
			response_text = await process_with_llm(clarification_prompt)
			state.conversation_history.append({"role": "assistant", "content": response_text})

		# 2. Cancellations
		elif detected_intents.get("cancellation"):
			response_text = "Okay, I've cancelled your trip. Let me know if you'd like to start a new one later."
			state.conversation_history.append({"role": "assistant", "content": response_text})
			state.clear_trip_context()
		
		# 3. Routing (w/ interruptions)
		elif detected_intents.get("generate_routes") or detected_intents.get("trip_changes"):

			# Selecting which prompt to use (TRIP_CHANGES_PROMPT for updates, SYSTEM_PROMPT otherwise)
			if detected_intents.get("trip_changes"):
				extraction_prompt = [
					{"role": "system", "content": SYSTEM_PROMPT},
                    {"role": "system", "content": TRIP_CHANGES_PROMPT},
                    {"role": "user", "content": (
                        "The conversation so far:\n" +
                        "\n".join(f"{m['role']}: {m['content']}" for m in state.conversation_history) +
                        "\n\nPlease output the final travel JSON now."
                    )}
				]
				
			else:
				extraction_prompt = [
					{"role": "system", "content": SYSTEM_PROMPT},
					{"role": "user", "content": (
						"The conversation so far is:\n" +
						"\n".join(f"{m['role']}: {m['content']}" for m in state.conversation_history) +
						"\n\nPlease output the final travel JSON now."
						)}
				]

			# Get the travel JSON
			final_travel_json = await process_with_llm(extraction_prompt)
			state.final_llm_response = final_travel_json
			
			# --- Database preference querying ---
			resolved_destination_coords = None
			history_avoid = None
			history_familiar = None
			history_option = None
			destination_name = final_travel_json.get("destination", "").strip().lower()

			# Semantic location resolution is performed first to resolve keywords like "home", "work", etc.
			resolved_destination_coords = await resolve_semantic_location(
				db = db,
				user_id = x_user_id,
				destination_label = destination_name
			)
			if resolved_destination_coords:
				print(f"[DEBUG] Semantic location matched, coordinates: {resolved_destination_coords}")
			
			history_avoid, history_familiar, history_option = await check_user_preferences(
				db = db,
				user_id = x_user_id,
				destination = final_travel_json.get("destination")
			)
			print(f"[DEBUG] DB preferences - Avoid: {history_avoid}, Option: {history_option}, Familiar: {history_familiar}")

			# --- Interruption logic if preferences are detected ---
			avoid_input = final_travel_json.get("avoid", [])

			if state.pending_preference == True:
				if history_option:
					state.pending_preference = "route_option"
					response_text = f"You usually take the {history_option} route for this trip. Do you want me to use the {history_option} route?"
					state.conversation_history.append({"role": "assistant", "content": response_text})
                
				elif not avoid_input and history_avoid:
					state.pending_preference = "avoid"
					response_text = f"You usually avoid {history_avoid} on this trip. Do you want me to avoid it again?"
					state.conversation_history.append({"role": "assistant", "content": response_text})
                
				elif history_familiar:
					state.pending_preference = "familiarity"
					response_text = f"You usually take the route {history_familiar}. Do you want me to use it again?"
					state.conversation_history.append({"role": "assistant", "content": response_text})

				else:
					state.pending_preference = False
                    # TODO: Mapbox functions
					response_text = "I have your destination. Let's get your route ready!" 
					state.conversation_history.append({"role": "assistant", "content": response_text})
            
			else:
                # TODO: Mapbox functions
				response_text = "Updating your route now!"
				state.conversation_history.append({"role": "assistant", "content": response_text})


		# 4. Preference resolution: YES
		elif detected_intents.get("preference_remembering"):
			destination_name = state.final_llm_response.get("destination", "").strip()
			history_avoid, history_option, history_familiar = await check_user_preferences(
                db=db, user_id=x_user_id, destination=destination_name
            )
			# Apply the specific preference based on what was asked (aka the value of pending_preference)
			if state.pending_preference == "route_option" and history_option:
				state.final_llm_response["option"] = history_option
				response_text = f"Got it, applying the {history_option} route."

			elif state.pending_preference == "avoid" and history_avoid:
				if "avoid" not in state.final_llm_response or not isinstance(state.final_llm_response["avoid"], list):
					state.final_llm_response["avoid"] = []
				state.final_llm_response["avoid"].append(history_avoid)
				response_text = f"Got it, I will avoid {history_avoid}."

			elif state.pending_preference == "familiarity" and history_familiar:
				if "via" not in state.final_llm_response or not isinstance(state.final_llm_response["via"], list):
					state.final_llm_response["via"] = []
				state.final_llm_response["via"].append(history_familiar)
				response_text = f"Got it, routing you through {history_familiar}."
			
			else:
				response_text = "Got it. Updating your route parameters now."

			state.pending_preference = False

			# TODO: Mapbox functions
            # TODO: Later: save this preference acceptance to the DB (Task 4.6).

			response_text += " Getting your route ready!"
			state.conversation_history.append({"role": "assistant", "content": response_text})

		# 5. Preference resolution: NO
		elif detected_intents.get("no_preference_remembering"):
			state.pending_preference = False

			response_text = "No problem, I'll ignore that. Getting your standard route ready!"
			state.conversation_history.append({"role": "assistant", "content": response_text})
			
			# TODO: Mapbox functions
            # TODO: Later: save this preference rejection to the DB (Task 4.6).

		# 6. Start navigation
		elif detected_intents.get("start_nav"): 
			state.navigation_started = True
            
			response_text = "Starting navigation. Drive safely!"
			state.conversation_history.append({"role": "assistant", "content": response_text})
            
            # TODO: Save final trip to the DB

		# 7. Request alternate routes
		elif detected_intents.get("request_alternates"):
			response_text = "Let me look for some other ways to get there."
			state.conversation_history.append({"role": "assistant", "content": response_text})
            
            # TODO: Fetch alternate routes with format_alternates_response

		# 8. Select specific route
		elif detected_intents.get("select_route"): 
			# Note: In the future, use a quick LLM extraction here to figure out WHICH route they picked
			response_text = "Got it, switching to that route now."
			state.conversation_history.append({"role": "assistant", "content": response_text})
            
            # TODO: Update the active Mapbox route

		# 9. Start new trip (mid-navigation)
		elif detected_intents.get("start_new_trip"):
			# Note: Not working well yet, still needs a flag so it knows if the user is responding to this and not preferences
			response_text = "You are currently navigating. Should I cancel this trip and start a new one?"
			state.conversation_history.append({"role": "assistant", "content": response_text})

		# 10. Fallback
		else:
			response_text = "I'm sorry, I don't know that one."
			state.conversation_history.append({"role": "assistant", "content": response_text})
		
		# Response
		return {
            "transcription": user_text,
            "intents": detected_intents,
            "heyroute": response_text,
            "travel_data": final_travel_json,
            "history": state.conversation_history,
            "turn_number": state.turn_count,
            "message": "Turn processed successfully."
        }

		# EVERYTHING BELOW IS THE OLD CODE PRE-INTERRUPT
		# if detected_intents.get("generate_routes") == True:
		# 	extraction_prompt = [
		# 		{"role": "system", "content": SYSTEM_PROMPT},
		# 		{"role": "user", "content": (
		# 			"The conversation so far is:\n" +
		# 			"\n".join(f"{m['role']}: {m['content']}" for m in state.conversation_history) +
		# 			"\n\nPlease output the final travel JSON now."
		# 			)}
		# 	]
		# 	final_travel_json = await process_with_llm(extraction_prompt)
		# 	print(f"[DEBUG] Extracted travel data for session {x_session_id}: {final_travel_json}")

		# 	state.final_llm_response = final_travel_json

		# # --- Database preference querying ---
		# resolved_destination_coords = None
		# history_avoid = None
		# history_familiar = None
		# history_option = None

		# if final_travel_json:
		# 	destination_name = final_travel_json.get("destination", "").strip().lower()

		# 	# Semantic location resolution is performed first to resolve keywords like "home", "work", etc.
		# 	resolved_destination_coords = await resolve_semantic_location(
		# 		db = db,
		# 		user_id = x_user_id,
		# 		destination_label = destination_name
		# 	)

		# 	if resolved_destination_coords:
		# 		print(f"[DEBUG] Semantic location matched, coordinates: {resolved_destination_coords}")

		# 	history_avoid, history_familiar, history_option = await check_user_preferences(
		# 		db = db,
		# 		user_id = x_user_id,
		# 		destination = final_travel_json.get("destination")
		# 	)
		# 	print(f"[DEBUG] DB preferences - Avoid: {history_avoid}, Option: {history_option}, Familiar: {history_familiar}")
		
		# return {
		# 	"transcription": user_text,
		# 	"intents": detected_intents,
		# 	"travel_data": final_travel_json,
		# 	"database_layer":{
		# 		"semantic_resolved_coords": resolved_destination_coords,
		# 		"historical_preferences": {
		# 			"most_avoided_road": history_avoid,
		# 			"most_familiar_road": history_familiar,
		# 			"preferred_route_option": history_option
		# 		}
		# 	},
		# 	"message": "Audio processed and transcribed, intents extracted, initial travel data generated, and database context retrieved."
		# }

	except httpx.RequestError as exc:
		raise HTTPException(status_code=500, detail=f"Unable to connect to the ASR server: {str(exc)}")
	except Exception as e:
		raise HTTPException(status_code=500, detail=f"Unable to process the audio file: {str(e)}")