import subprocess
import time
import os
from wit import Wit
import random
import json
from fastapi import FastAPI, Request, HTTPException
import logging
import threading
from mutagen.mp3 import MP3

logging.basicConfig(filename='WebServer.log', filemode='w', level=logging.DEBUG)

# System updated for Hugging Face container environment
platform = "Linux"

# --- CHANGED SECTION ---
# Fetch the Wit.ai API Key(s) from Hugging Face Secrets
# If using multiple keys, separate them with commas in the Secret (e.g., KEY1,KEY2)
env_keys = os.getenv("WIT_API_KEY", "")
api_keys = [key.strip() for key in env_keys.split(",")] if env_keys else []

if not api_keys or not api_keys[0]:
    print("WARNING: WIT_API_KEY Secret is not set in Hugging Face!")
# ----------------------

wit_api = {}
resp = {}

for api_key in api_keys:
	wit_api[api_key] = Wit(api_key)

# Environment setup using a safe container directory (/tmp)
base_dir = "/tmp/SpeechRecognition"

if platform == "Linux":
	try:
		if not os.path.exists(base_dir):
			os.makedirs(f"{base_dir}/Cache", exist_ok=True)
			os.chdir(base_dir)
		else:    
			os.chdir(base_dir)
	except Exception as e:
		print(f"Error setting up directory: {e}")

app = FastAPI()

# Replace your current @app.get('/') with this:

@app.get('/')
def index(request: Request):
    return {
        "status": "online", 
        "message": "Speech server is running. Send a .webm audio file via POST to /speech"
    }

@app.post('/speech')
async def process(request: Request):
	start_execution = time.time()

	# Request Processing
	data_length = str(round(float(int(request.headers.get('content-length', 0)) / 1024),2))
	random_hash = random.randint(10000000000,99999999999)

	if float(data_length) > 500:
		print('The incoming speech request has exceeded the 500KB limit established by the Developer, Size: ' + str(data_length) + "KB")
		return

	print("[" + str(random_hash) + "] Speech request received from " + request.client.host + " | " + " Size of incoming data: " + data_length + "KB.")
	logging.info("[" + str(random_hash) + "] Speech request received from " + request.client.host + " | " + " Size of incoming data: " + data_length + "KB.")

	webm = await request.body()
	logging.debug("[" + str(random_hash) + "] .webm obtained from the Body request.")

	# File Writing
	with open("Cache/" + str(random_hash) + ".webm", "wb") as file:
		file.write(webm)
		file.close()
	logging.debug("[" + str(random_hash) + "] .webm written to file " + str(random_hash) + ".webm")

	# File Conversion
	logging.debug("[" + str(random_hash) + "] Starting conversion .webm to .mp3")
	start_time = time.time()

	if platform == "Linux":
		convert = subprocess.Popen([
			"ffmpeg", "-i", f"{base_dir}/Cache/{random_hash}.webm",
			"-af", "silenceremove=stop_periods=-1:stop_duration=0.02:stop_threshold=-53dB",
			"-vn", "-ac", "1", "-b:a", "64k", f"{base_dir}/Cache/{random_hash}.mp3"
		], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
		convert.wait() 
		
	logging.debug("[" + str(random_hash) + "] Conversion finished, time " + str(round(float((time.time() - start_time) * 1000),2)) + " msec")
	
	# File Open
	with open("Cache/" + str(random_hash) + ".mp3", "rb") as file:
		audio_read = file.read()
		file.close()

	# Audio Length
	length = 0
	try:
		audio = MP3("Cache/" + str(random_hash) + ".mp3")
		length = audio.info.length
	except:
		print("Error reading Audio")

	start_req = time.time()
	
	# Failsafe if no API keys are loaded
	if not api_keys:
		return {"error": "API keys not configured on server."}
        
	random_api_key = random.choice(api_keys)
	
	# Req to Wit.Ai
	try:
		resp[random_hash] = wit_api[random_api_key].speech(audio_read, {'Content-Type': 'audio/mpeg3'})
		logging.debug("[" + str(random_hash) + "] .mp3 sent to Wit.Ai | Request to Wit.Ai completed in " + str(round(float((time.time() - start_req) * 1000),2)) + " msec | Audio Length: " + str(round(float(length),2)) + "s")
		print("[" + str(random_hash) + "] .mp3 sent to Wit.Ai | Request to Wit.Ai completed in " + str(round(float((time.time() - start_req) * 1000),2)) + " msec | Audio Length: " + str(round(float(length),2)) + "s")
		print("[" + str(random_hash) + "] Recognized text ", resp[random_hash])
	except Exception as e:
		logging.debug(e)
		logging.info("[" + str(random_hash) + "] Unable to make request, manual check required.")

	# Clean-up
	def RemoveEnvironment(random_hash):
		try:
			os.remove("Cache/" + str(random_hash) + ".webm")
			os.remove("Cache/" + str(random_hash) + ".mp3")
			logging.debug("[" + str(random_hash) + "] Cleaned up environment.")
		except:
			logging.debug("[" + str(random_hash) + "] Unable to clean environment.")

	threading.Thread(target=RemoveEnvironment, args=[random_hash]).start()

	print("[" + str(random_hash) + "] Speech request completed for " + request.client.host)
	logging.info("[" + str(random_hash) + "] Speech request completed for " + request.client.host)

	logging.info("[" + str(random_hash) + "] Request completed in " + str(round(float((time.time() - start_execution) * 1000),2)) + " msec")
	
	try:
		print("[" + str(random_hash) + "] Answer from Wit.Ai API: " + resp[random_hash]['text'])
		return resp[random_hash]['text']
	except KeyError:
		return {"error": "Could not recognize speech or invalid Wit.ai response."}