#!/usr/bin/python3
from typing import Annotated
from fastapi import Request, Depends, FastAPI, UploadFile, Response
from fastapi.responses import PlainTextResponse
from fastapi.security import HTTPBasic, HTTPBasicCredentials
from starlette.middleware.sessions import SessionMiddleware
from nicegui import ui, app, run, Client
from concurrent.futures import ThreadPoolExecutor
from logging import info, warning, error, debug
from base64 import b64encode, b64decode
from dotenv import load_dotenv
from time import sleep
from asyncio import run_coroutine_threadsafe
from websocket_client import WebSocketClient

import logging
import os
import io
import aiohttp
import asyncio
import json
import traceback

load_dotenv()

logging.basicConfig()
logging.getLogger().setLevel(logging.DEBUG)

# Folder to store recorded audio
app.add_static_files("/static", "static")
app.add_middleware(SessionMiddleware, secret_key=os.environ["SESSION_KEY"])

security = HTTPBasic()
main_loop = asyncio.get_event_loop()

session_timeout = aiohttp.ClientTimeout(total=None,sock_connect=5,sock_read=5)

client_to_session_id = {}
session_to_websocket = {}

def handle_disconnect(client : Client):
    info(f"Client disconnected: {client.id}")
    if client.id in client_to_session_id:
        session_id = client_to_session_id[client.id]
        if session_id in session_to_websocket:
            del session_to_websocket[session_id]
        del client_to_session_id[client.id]

app.on_disconnect(handle_disconnect)

def start_recording(ui_elements):
    if not refresh_ui_enabled(ui_elements):
        return
    ui.run_javascript('startRecording()')

def stop_recording(ui_elements):
    ui.run_javascript('stopRecording()')

def stop_playback(ui_elements):
    ui.run_javascript('stopPlayback()')
    enable_ui(ui_elements, 0)

def start_playback(ui_elements):
    ui.run_javascript('startPlayback()')

def show_audio_notification():
    ui.notify(f"Saved recording! Waiting for response...")

def is_authorized(credentials: Annotated[HTTPBasicCredentials, Depends(security)]):
    user = os.environ["USERNAME"]
    pw = os.environ["PASSWORD"]
    return credentials.username == user and credentials.password == pw

def refresh_ui_enabled(ui_elements):
    with ui_elements["main"]:
        if "websocket" in ui_elements and ui_elements["websocket"].enabled:
            enable_ui(ui_elements, 0)
            return True
        if "websocket" in ui_elements and ui_elements["websocket"].generating:
            enable_ui(ui_elements, 2)
            return False
        enable_ui(ui_elements, 1)
        stop_recording(ui_elements["id"])
    return False

def after_websocket_init(ui_elements, was_successful):
    if not was_successful:
        enable_ui(ui_elements, 1)
    else:
        enable_ui(ui_elements, 0)

def enable_ui(ui_elements, state = 0):
    info(f"Enabled UI {ui_elements['id']}: {state}")
    if state == 0:
        ui_elements["micbutton"].props(remove="disabled")
        ui_elements["text"].set_text("Hold the button to record audio")
    elif state == 1:
        ui_elements["micbutton"].props("disabled")
        ui_elements["text"].set_text("Server is offline :(")
    elif state == 2:
        ui_elements["micbutton"].props("disabled")
        ui_elements["text"].set_text("Sending input data...")
    else:
        error(f"Invalid state inside enable_ui: {state}")

def after_end_audio(ui_elements, output):
    debug("after_end_audio")
    if output["status"] != "ok":
        refresh_ui_enabled(ui_elements)
        return

@app.post('/api/v1/send')
async def end_audio(credentials: Annotated[HTTPBasicCredentials, Depends(security)], request: Request):
    if not is_authorized(credentials):
        return Response({"status": "Unauthorized"}, status_code=401)
    try:
        session_id = request.session["id"]
        if session_id not in session_to_websocket:
            return {"status": "nok"}
        websocket = session_to_websocket[session_id]
        with websocket.ui_elements["main"]:
            if not refresh_ui_enabled(websocket.ui_elements):
                return {"status": "nok"}
        
        debug("Handle audio")
        content = await request.body()
        bytes = io.BytesIO(content)

        await websocket.push_msg_to_queue({"method": "end_audio", "audio_content": bytes}, after_end_audio)

        with websocket.ui_elements["main"]:
            enable_ui(websocket.ui_elements, 2)
            stop_recording(session_id)
            show_audio_notification()

        return {"status": "ok"}
    except Exception as ex:
        error(traceback.format_exc())
        return {"status": "nok"}
    
def after_handle_audio(ui_elements, output):
    debug("after_handle_audio")
    if output["status"] != "ok":
        refresh_ui_enabled(ui_elements)
        return
    debug("Successfully handled audio")

def after_handle_upstream(ui_elements, response):
    if "audio" in response:
        with ui_elements["main"]:
            ui_elements["audio"].set_source(f"data:audio/webm;base64,{response['audio']}")
            start_playback(ui_elements["id"])
    pass
    
@app.post('/api/v1/upload')
async def handle_audio(credentials: Annotated[HTTPBasicCredentials, Depends(security)], request: Request):
    global main_containers, main_loop
    if not is_authorized(credentials):
        return Response({"status": "Unauthorized"}, status_code=401)
    try:
        session_id = request.session["id"]
        if session_id not in session_to_websocket:
            return {"status": "nok"}
        websocket = session_to_websocket[session_id]
        with websocket.ui_elements["main"]:
            if not refresh_ui_enabled(websocket.ui_elements):
                return {"status": "nok"}

        debug("Handle audio")
        content = await request.body()
        bytes = io.BytesIO(content)

        await websocket.push_msg_to_queue({"method": "audio", "audio_content": bytes}, after_handle_audio)

        return {"status": "ok"}
    except Exception as ex:
        error(traceback.format_exc())
        return {"status": "nok"}

## Initialize page

@ui.page("/")
async def main_page(credentials: Annotated[HTTPBasicCredentials, Depends(security)], request: Request, client : Client) -> None:
    global main_containers, mic_buttons, text_containers
    if not is_authorized(credentials):
        return Response("Get outta here män :(", status_code=401)
    #await ui.context.client.connected()
    session_id = request.session["id"]
    ui.add_head_html(f"<script type='text/javascript' src='/static/record.js'></script>") 
    ui.add_head_html(f"<link type='text/tailwindcss' rel='stylesheet' href='/static/style.css'>") 

    container = ui.column().classes('w-full h-[calc(100vh-2rem)] items-center justify-center')

    with container:
        ui_elements = {}
        ui_elements["id"] = session_id
        ui_elements["main"] = container
        label = ui.label('Checking if the server is online...').classes('text-xl mb-20 whitespace-nowrap text-[3vw]')
        ui_elements["text"]  = label
        with ui.element().classes('w-[50vw] max-h-[75%] aspect-square rounded-full bg-grey text-white flex items-center justify-center whitespace-nowrap text-[3vw] active:scale-95 transition-transform').props("id=recordbutton disabled") as button:
            ui.image('/static/microphone.png').classes('h-[75%] w-[75%] object-contain')
        button.on('mousedown', lambda: start_recording(ui_elements))
        button.on('mouseup', lambda: stop_recording(ui_elements))
        button.on('touchstart', lambda: start_recording(ui_elements))
        button.on('touchend', lambda: stop_recording(ui_elements))
        ui_elements["micbutton"] = button

        audio = ui.audio(src="").classes("disabled opacity-0").props("id=audioout")
        audio.on('ended', lambda: stop_playback(ui_elements))
        ui_elements["audio"] = audio

        ui_elements["websocket"] = WebSocketClient(ui_elements, after_handle_upstream)
        ui_elements["websocket"].start_thread(after_websocket_init)
        client_to_session_id[client.id] = session_id
        session_to_websocket[session_id] = ui_elements["websocket"]
ui.run(show=False)