#!/usr/bin/env python
### -----------------------------------------------------------------------
### (test_BASE, Revised) version_1.07 ALPHA, app.py
### -----------------------------------------------------------------------

# -------------------------------------------------------------------------
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# -------------------------------------------------------------------------

import os
import re
import uuid
import time
import psutil
import pydub
import subprocess
from tqdm import tqdm

import tempfile
from fpdf import FPDF
from pathlib import Path

import numpy as np
import soundfile as sf
import torch
from transformers import pipeline 

from gpuinfo import GPUInfo

import gradio as gr


###############################################################################
# Configuration.
###############################################################################

#if not torch.cuda.is_available():
    #DESCRIPTION += "\n<p>⚠️Running on CPU, This may not work on CPU.</p>"

CACHE_EXAMPLES = torch.device('cuda') and os.getenv("CACHE_EXAMPLES", "0") == "1"
#CACHE_EXAMPLES = torch.cuda.is_available() and os.getenv("CACHE_EXAMPLES", "0") == "1"
#USE_TORCH_COMPILE = os.getenv("USE_TORCH_COMPILE", "0") == "1"
#ENABLE_CPU_OFFLOAD = os.getenv("ENABLE_CPU_OFFLOAD", "0") == "1"

device = torch.device('cuda')
#device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

def transcribe(file_upload, progress=gr.Progress(track_tqdm=True)):

    file = file_upload # microphone if microphone is not None else 
    start_time = time.time()
    
    # -- ex subrosa ad sf
    audio, sr = sf.load(file_upload, sr=None)
    duration = sf.get_duration(y=audio, sr=sr)
    
    # -- asr pipeline
    with torch.no_grad():
        pipe = pipeline("automatic-speech-recognition", model="NbAiLab/nb-whisper-large", chunk_length_s=30, device=device)
    
    # -- process audio in chunks of 30 seconds
    chunk_size = sr * 30  # 30 seconds
    text = ""
    for start in range(0, len(audio), chunk_size):
        end = min(start + chunk_size, len(audio))
        chunk = audio[start:end]
        
        # -- convert audio chunk to format for pipeline
        chunk_file = "temp_chunk.wav"
        sf.write(chunk_file, chunk, sr)
        
        
        # -- chnk ad transcriptrauma
        chunk_text = pipe(chunk_file)["text"]
        text += chunk_text + " "

    end_time = time.time()
    output_time = end_time - start_time
    
    # --word count
    word_count = len(text.split())

    # --memory metrics
    memory = psutil.virtual_memory()
    
    # --cpu metric
    #cpu_usage = psutil.cpu_percent(interval=1)

    # --gpu metric
    #gpu_utilization, gpu_memory = GPUInfo.gpu_usage()
    
    # --system info string
    system_info = f"""
    #Processing time: {output_time:.2f} seconds.
    #Number of words: {word_count}
    GPU Memory: {gpu_memory}%
    GPU Utilization: {gpu_utilization}%
    """
    
    return text.strip(), system_info


###############################################################################
# Interface.
###############################################################################

HEADER_INFO = """
    # SWITCHVOX ✨|🇳🇴 *Transkribering av lydfiler til norsk bokmål.*
""".strip()
LOGO = "https://cdn-lfs-us-1.huggingface.co/repos/fe/3b/fe3bd7c8beece8b087fddcc2278295e7f56c794c8dcf728189f4af8bddc585e1/5112f67899d65e9797a7a60d05f983cf2ceefbe2f7cba74eeca93a4e7061becc?response-content-disposition=inline%3B+filename*%3DUTF-8%27%27logo.png%3B+filename%3D%22logo.png%22%3B&response-content-type=image%2Fpng&Expires=1725531489&Policy=eyJTdGF0ZW1lbnQiOlt7IkNvbmRpdGlvbiI6eyJEYXRlTGVzc1RoYW4iOnsiQVdTOkVwb2NoVGltZSI6MTcyNTUzMTQ4OX19LCJSZXNvdXJjZSI6Imh0dHBzOi8vY2RuLWxmcy11cy0xLmh1Z2dpbmdmYWNlLmNvL3JlcG9zL2ZlLzNiL2ZlM2JkN2M4YmVlY2U4YjA4N2ZkZGNjMjI3ODI5NWU3ZjU2Yzc5NGM4ZGNmNzI4MTg5ZjRhZjhiZGRjNTg1ZTEvNTExMmY2Nzg5OWQ2NWU5Nzk3YTdhNjBkMDVmOTgzY2YyY2VlZmJlMmY3Y2JhNzRlZWNhOTNhNGU3MDYxYmVjYz9yZXNwb25zZS1jb250ZW50LWRpc3Bvc2l0aW9uPSomcmVzcG9uc2UtY29udGVudC10eXBlPSoifV19&Signature=vNqupMg9p-Wx9BvqVlK5BhhyOeS58YjZW2-bEP4OVd0azdA0AfW0QG8EGxZ1h7UwuJnMwlPGayA0P46Ob9DSRH48BGjH176UgbPBcasSAI43jb9PJO9qIznrv9orzPt3ZrqTll0d9cKayQ96iPWQond-G5xbl0bNYb9qLXh9w3Ww%7EELKIFU9KeDOvIKww9cHftCeVFqCFJC%7Etimk-eOHo9g4xVfAaVMFoVNeJOVVpTW-MzPb1EGccyN9-3WJaF9Nwg3fkb7FRazg8IYcAatS2PahLpfp-zJup7y-ywnPzb8jJPgN3TBu6-M7hE4OHVcRmxeXk3VDRgSFVfbmnrlc%7Ew__&Key-Pair-Id=K24J24Z295AEI9"
SIDEBAR_INFO = f"""
<div align="center">
    <img src="{LOGO}" style="width: 100%; height: auto;"/>
</div>
"""

def save_to_pdf(text, summary):
    pdf = FPDF()
    pdf.add_page()
    pdf.set_font("Arial", size=12)

    if text:
        pdf.multi_cell(0, 10, "Transkribert Tekst:\n" + text)

    pdf.ln(10)  # Paragraph metric

    if summary:
        pdf.multi_cell(0, 10, "Summary:\n" + summary)

    pdf_output_path = "transcription_.pdf"
    pdf.output(pdf_output_path)
    return pdf_output_path

css = """
#transcription_output textarea {
    background-color: #000000;  /* black */
    color: #00FF00 !important;  /* text color */
    font-size: 18px;  /* font size */
}

#system_info_box textarea {
    background-color: #ffe0b3;  /* orange */
    color: black !important;  /* text color */
    font-size: 16px;  /* font size */
    font-weight: bold;  /* bold font */
}
"""

iface = gr.Blocks(css=css)

with iface:

    gr.HTML(SIDEBAR_INFO)
    gr.Markdown(HEADER_INFO)

    with gr.Row():
        gr.Markdown('''
        ##### 🔊 Last opp lydfila 
        ##### ☕️ Trykk på "Transkriber" knappen og vent på svar
        ##### ⚡️ Går rimelig bra kjapt med Norwegian NB-Whisper Large..
        ##### 😅 Planlegger tilleggs-funksjoner senere
        
        ''')
        #microphone = gr.Audio(label="Microphone", sources="microphone", type="filepath")
        upload = gr.Audio(label="Upload audio", sources="upload", type="filepath")
        transcribe_btn = gr.Button("Transkriber")

    with gr.Row():   
        with gr.Column(scale=3):
            text_output = gr.Textbox(label="Transkribert Tekst", elem_id="transcription_output")
        with gr.Column(scale=1):
            system_info = gr.Textbox(label="Antall sekunder, ord:", elem_id="system_info_box")
        
    
    with gr.Tabs():
        with gr.TabItem("Download PDF"):
            pdf_text_only = gr.Button("Last ned pdf med resultat")
            pdf_output = gr.File(label="/.pdf")

            pdf_text_only.click(fn=lambda text: save_to_pdf(text, ""), inputs=[text_output], outputs=[pdf_output])

    with gr.Row():
        gr.Markdown('''
        <div align="center">
            <a href="https://opensource.com/resources/what-open-source">
                <img src="https://badgen.net/badge/Open%20Source%20%3F/Yes%21/blue?icon=github" alt="Open Source? Yes!">
        </a>
        <span style="display:inline-block; width: 20px;"></span> 
        <a href="https://opensource.org/licenses/Apache-2.0">
            <img src="https://img.shields.io/badge/License-Apache_2.0-blue.svg" alt="License: Apache 2.0">
        </a>
        </div>
        ''')
    transcribe_btn.click(
        fn=transcribe, 
        inputs=[upload], # microphone 
        outputs=[text_output, system_info]
    )    
    
    #transcribe_btn.click(fn=transcribe, inputs=[microphone, upload], outputs=[text_output, system_info])



iface.launch(share=True,debug=True)