import os
import json
import re
import gradio as gr
from google import genai
from google.genai import types
import asyncio
from datasets import load_dataset, DatasetDict, Dataset
from huggingface_hub import login
import datetime

# Authenticate with HF token
hf_token = os.getenv("HF_TOKEN")
login(token=hf_token)
dataset_name = "spriambada3/ehealth_transcribe"


def init_dataset():
    try:
        dataset = load_dataset(dataset_name)
    except Exception as e:
        print(e)
        dataset = DatasetDict(
            {"data": Dataset.from_dict({"logintime": [], "email": [], "wa": []})}
        )
    print("init dataset result ")
    print(dataset)
    return dataset


def add_user(dataset, email, wa):
    new_data = {
        "logintime": datetime.datetime.now(),
        "email": email,
        "wa": wa,
    }
    dataset["data"] = dataset["data"].add_item(new_data)
    dataset.push_to_hub(dataset_name)  # Save to HF Hub

    print("add data successful")


def audio_from_bytes(audio_file_path: str):
    """Converts an audio file into Gemini-compatible format."""
    try:
        with open(audio_file_path, "rb") as f:
            audio_data = f.read()

        mime_type = "audio/mp3"  # Adjust based on your audio type
        return types.Part.from_bytes(data=audio_data, mime_type=mime_type)

    except FileNotFoundError:
        return "Error: Audio file not found!"
    except Exception as e:
        return f"An error occurred: {e}"


def transcribe_and_summarize(audio_file, session):
    """Processes audio with Gemini API and returns a SOAP summary."""
    if audio_file is None:
        return "No audio file uploaded."

    # Ensure API Key is set
    GEMINI_API_KEY = os.getenv("GEMINI_API_KEY")
    if not GEMINI_API_KEY:
        return "Error: GEMINI_API_KEY environment variable is missing."
    asyncio.set_event_loop(asyncio.new_event_loop())

    client = genai.Client(api_key=GEMINI_API_KEY)
    model = "gemini-2.0-flash"

    # Prepare the request
    contents = [
        types.Content(
            role="user",
            parts=[
                types.Part.from_text(
                    text="""Anda adalah asisten medis yang membantu dokter dalam menyusun catatan medis dalam bentuk paragraf menggunakan bahasa Indonesia.
Buat ringkasan SOAP berdasarkan percakapan dokter dan pasien dalam format berikut:

    Subjective:
    ICD10:
    Objective:
    Assessment:
    Plan:

Identifikasi dan berikan saran dalam bahasa Indonesia tindakan logis selanjutnya dalam format:
    ICD10:
    Obat:
    Laboratorium:
    Radiologi:
"""
                ),
                audio_from_bytes(audio_file),
            ],
        )
    ]

    generate_content_config = types.GenerateContentConfig(
        temperature=0,
        top_p=0.95,
        top_k=40,
        max_output_tokens=8192,
        response_mime_type="text/plain",
    )

    # Process the audio
    response_text = ""
    for chunk in client.models.generate_content_stream(
        model=model,
        contents=contents,
        config=generate_content_config,
    ):
        response_text += chunk.text
    counter_display, session = click_button(session)
    return response_text, counter_display, session


DATA_FILE = "user_data.json"


def load_user_data():
    """Load user data from JSON file."""
    if os.path.exists(DATA_FILE):
        with open(DATA_FILE, "r") as file:
            return json.load(file)
    return {}  # Return empty dictionary if no data


def save_user_data(username, email):
    """Save user data to JSON file with default counter = 10 if new."""
    data = load_user_data()

    if username not in data:  # New user
        data[username] = {"email": email, "counter": 10}  # Set default counter to 10

    with open(DATA_FILE, "w") as file:
        json.dump(data, file, indent=4)
    wa = username
    dataset = init_dataset()
    add_user(dataset, email, wa)
    return data


def is_valid_wa(username):
    """Check if username is all numbers, at least 11 characters, and starts with 08 or 62."""
    return re.fullmatch(r"^(08|62)\d{9,}$", username) is not None


def is_valid_email(email):
    """Check if email format is valid."""
    return re.fullmatch(r"^[\w\.-]+@[\w\.-]+\.\w+$", email) is not None


def login(username, email, session):
    """Handles user login or registration with validation."""

    if not is_valid_wa(username):
        return (
            "❌ Invalid WA! Nomor WA minimal 12 digits dimulai '08' atau '62'.",
            session,
            gr.update(visible=True),
            gr.update(visible=False),
        )

    if not is_valid_email(email):
        return (
            "❌ Invalid email format! Please enter a valid email.",
            session,
            gr.update(visible=True),
            gr.update(visible=False),
        )

    data = save_user_data(username, email)  # Save or retrieve user data

    session["username"] = username
    session["counter"] = data[username]["counter"]

    return "", session, gr.update(visible=False), gr.update(visible=True)


def click_button(session):
    """Decrease counter on button click."""
    err = f"⚠️ Quota habis. Silahkan mengunjungi https://ehealth.co.id atau WA 6285777779926 untuk menambah kuota"
    if session["counter"] > 0:
        session["counter"] -= 1

        # Update the user data in JSON file
        data = load_user_data()
        data[session["username"]]["counter"] = session["counter"]

        with open(DATA_FILE, "w") as file:
            json.dump(data, file, indent=4)

        if session["counter"] == 0:
            return (err, session)
        return f"Quota: {session['counter']}", session
    else:
        return (err, session)


# Gradio Interface
with gr.Blocks(theme=gr.themes.Default()) as demo:
    session = gr.State({"username": None, "counter": 0})  # Manage session state
    # Login Section
    login_block = gr.Column(visible=True)
    with login_block:
        gr.HTML(
            """
        <div style="text-align: center;">
            <a href="https://youtube.com/shorts/DmiVhj9ROag?si=j5Opmjny3kNdLlrf" target="_blank" style="font-size: 20px;">
                <strong>Klik disini untuk Demo Video <img src="https://upload.wikimedia.org/wikipedia/commons/b/b8/YouTube_Logo_2017.svg" alt="YouTube" width="100"><br></strong>
            </a>
        </div>
        """
        )
        email_input = gr.Textbox(label="Email")
        username_input = gr.Textbox(label="WA", type="password")  # Hide input
        login_button = gr.Button("🔑 Login / Register")
        gr.Markdown(
            """### dengan login, saya menyetujui ketentuan penggunaan data perusahaan https://eHealth.co.id dan tidak akan menuntut eHealth.co.id dalam uji coba gratis AI Transkripsi Medis ini
                        
            seluruh data yang saya sediakan adalah data yang benar dan tidak melanggar hukum
            
            saya memahami bahwa tidak ada data suara maupun tulisan medis yang akan disimpan oleh eHealth.co.id, namun perusahaan tidak dapat menjamin perlakuan data penyedia model AI (OpenAI, DeepSeek, Google, Mistral, dll.) 

            ### setelah quota habis, saya dapat menambah quota dengan mengunjungi https://ehealth.co.id atau WA 6285777779926"""
        )
        output_text = gr.Textbox(label="Status", interactive=False)

    # Main User Interface (After Login)
    user_block = gr.Column(visible=False)
    with user_block:
        counter_display = gr.Textbox(label="Status Message", interactive=False)

        gr.Interface(
            fn=transcribe_and_summarize,
            inputs=[gr.Audio(type="filepath", sources="microphone"), session],
            outputs=["text", counter_display, session],
            description="Halo, pastikan HP/Laptop memiliki microphone untuk merekam percakapan dokter-pasien menjadi rekam medis SOAP. Akun berlangganan https://ehealth.co.id dapat terintegrasi SATUSEHAT & BPJS secara otomatis",
            allow_flagging="never",
            theme="light",
        )
        use_case_description = gr.Markdown(
            """
        Selain Rekam Medis Pasien, dokumen lain yang dapat digitalisasi:

        - Surgery Notes atau Catatan Tindakan Lain 

        - Inform Concern (Dokter dan Keluarga Pasien/Pasien)

        - Counseling

        - Nursing reports

        - Clinical documentation

        - Continue Care Document (untuk RS dokumentasi pemberian obat, infus, dll)."""
        )

        # Login button action
        login_button.click(
            login,
            [username_input, email_input, session],
            [output_text, session, login_block, user_block],
            trigger_mode="once",
        )

demo.launch(allowed_paths=["./images/eHwhite.png", "eHwhite.png", "./images/pp.png"])