import os
import uuid
import queue
import threading
import time
import gradio as gr
import pandas as pd
import soundfile as sf
import numpy as np
import pyarrow.dataset as ds
import pyarrow.fs as pafs
from huggingface_hub import HfApi, hf_hub_download, HfFileSystem

# --- CONFIGURATION ---
REPO_ID = "nitmztesting/MiZonalv2.9"
REPO_TYPE = "dataset"
METADATA_FILE = "metadata.csv"

TOKEN = os.environ.get("HF_TOKEN")
APP_PASSWORD = os.environ.get("APP_PASSWORD")

api = HfApi(token=TOKEN)

# --- ADVANCED QUEUE SYSTEM CONFIGURATION ---
QUEUE_TARGET_SIZE = 20  # Pre-loads and maintains exactly 10 files in the background per user
session_queues = {}     # Maps session_id -> queue.Queue()
session_workers = {}    # Maps session_id -> threading.Thread()
session_caches = {}     # Maps session_id -> {task_id: {"wav": path, "text": text}}
global_lock = threading.Lock() # Prevents duplicate worker creation steps

# --- GLOBAL BATCH WRITER CONFIGURATION ---
write_queue = queue.Queue()
BATCH_COMMIT_INTERVAL = 300  # Flushes local modifications to HF every 300 seconds

# --- CONCURRENCY-SAFE FETCH ---
def fetch_latest_tracker():
    """Always fetches the absolute freshest copy of metadata.csv from the Hub."""
    try:
        csv_path = hf_hub_download(
            repo_id=REPO_ID, 
            filename=METADATA_FILE, 
            repo_type=REPO_TYPE, 
            token=TOKEN,
            force_download=True  # Guarantees no cross-space conflicts
        )
        local_df = pd.read_csv(csv_path)
        local_df["id"] = local_df["id"].astype(str)
        local_df["status"] = local_df["status"].fillna("Pending")
        local_df["validation_count"] = pd.to_numeric(local_df["validation_count"]).fillna(0).astype(int)
        local_df["validated_by"] = local_df["validated_by"].astype(object)
        return local_df
    except Exception as e:
        print(f"CRITICAL ERROR loading metadata.csv: {e}")
        return pd.DataFrame()

# --- INITIALIZE REMOTE PARQUET STREAMING ---
print("Connecting to remote Parquet dataset shards...")
try:
    if not TOKEN:
        raise ValueError("HF_TOKEN is missing from environment variables.")
    
    fs = HfFileSystem(token=TOKEN)
    pa_fs = pafs.PyFileSystem(pafs.FSSpecHandler(fs))
    data_dir = f"datasets/{REPO_ID}/data"
    
    pa_dataset = ds.dataset(data_dir, filesystem=pa_fs, format="parquet")
    print(f"Successfully linked remote Parquet shards at: {data_dir}")
except Exception as e:
    print(f"⚠️ Failed to establish Parquet stream connection: {str(e)}")
    pa_dataset = None

# --- CORE PARQUET STREAM EXTRACTION LOGIC ---
def fetch_task_by_id(task_id, current_df):
    try:
        row = current_df[current_df["id"] == str(task_id)].iloc[0]
        text_content = row["text"]
        
        if pd.isna(text_content):
            text_content = ""
            
        if pa_dataset is None:
            return None, "Error: Remote dataset unavailable. Check logs.", ""
            
        scanner = pa_dataset.scanner(filter=ds.field("id") == str(task_id), columns=["audio"])
        table = scanner.to_table()
        
        if table.num_rows == 0:
            return None, f"Error: ID '{task_id}' not found.", ""
            
        row_dict = table.to_pylist()[0]
        audio_struct = row_dict.get("audio")
        
        if not audio_struct:
            return None, "Error: 'audio' data block missing.", ""
            
        local_wav_path = f"current_{task_id}.wav"
        
        if "bytes" in audio_struct and audio_struct["bytes"] is not None:
            with open(local_wav_path, "wb") as f:
                f.write(audio_struct["bytes"])
        elif "array" in audio_struct:
            arr = np.array(audio_struct["array"], dtype=np.float32)
            sr = audio_struct.get("sampling_rate", 16000)
            sf.write(local_wav_path, arr, sr)
            
        return local_wav_path, str(text_content), task_id
        
    except Exception as e:
        return None, f"Error processing stream: {str(e)}", ""

# --- SINGLE FACTORY FETCH JOB ---
def fetch_single_random_task():
    """Hunts down exactly one random pending item from the storage layer."""
    current_df = fetch_latest_tracker()
    
    if current_df.empty: 
        return None, "⚠️ Database failed to load.", "", None
        
    pending = current_df[current_df["status"] == "Pending"]
    if pending.empty:
        return None, "All pending files have been validated! 🎉", "", None
        
    random_task_id = pending.sample(n=1).iloc[0]["id"]
    wav, txt, tid = fetch_task_by_id(random_task_id, current_df)
    return wav, txt, tid, None

# --- CONTINUOUS BACKGROUND BUFFER FILLER ---
def queue_maintenance_worker(session_id):
    """Loop running in an isolated background thread keeping the pipeline populated."""
    print(f"Background prefetch thread spawned for Session ID: {session_id}")
    q = session_queues[session_id]
    
    while True:
        try:
            if q.qsize() < QUEUE_TARGET_SIZE:
                task_data = fetch_single_random_task()
                q.put(task_data)
                print(f"[{session_id}] Cached a new task. Current buffer size: {q.qsize()}/{QUEUE_TARGET_SIZE}")
            else:
                time.sleep(1)
        except Exception as e:
            print(f"Error inside prefetch worker thread for {session_id}: {e}")
            time.sleep(2)

# --- EXTERNAL POP PIPELINE MANAGER ---
def get_next_task_from_buffer(session_id):
    """Instantly delivers a task from the session's background buffer."""
    with global_lock:
        if session_id not in session_queues:
            session_queues[session_id] = queue.Queue()
            t = threading.Thread(target=queue_maintenance_worker, args=(session_id,), daemon=True)
            session_workers[session_id] = t
            t.start()
            
    q = session_queues[session_id]
    
    if q.empty():
        print(f"⚠️ Buffer empty on user hit. Fetching immediate fallback synchronously...")
        return fetch_single_random_task()
        
    return q.get()

# --- BATCH COMMIT BACKGROUND LOOP ---
def batch_commit_worker():
    """Background daemon that gathers updates from RAM and writes them to HF in bulk."""
    print("🚀 Batch commit background worker daemon initialized.")
    while True:
        time.sleep(BATCH_COMMIT_INTERVAL)
        
        if write_queue.empty():
            continue
            
        print(f"🔔 Found elements inside write buffer. Initiating compilation for {write_queue.qsize()} items...")
        try:
            # Pull down the absolute freshest tracker state from the hub
            live_df = fetch_latest_tracker()
            if live_df.empty:
                print("⚠️ Batch skip: Could not retrieve base tracker file from repository.")
                continue
                
            updates_processed = 0
            
            # Drain everything currently in the queue safely
            while not write_queue.empty():
                try:
                    item = write_queue.get_nowait()
                    idx_list = live_df.index[live_df["id"] == item["id"]].tolist()
                    
                    if idx_list:
                        idx = idx_list[0]
                        live_df.at[idx, "text"] = item["text"]
                        live_df.at[idx, "status"] = item["status"]
                        live_df.at[idx, "validated_by"] = item["validated_by"]
                        live_df.at[idx, "validation_count"] += 1
                        updates_processed += 1
                        
                    write_queue.task_done()
                except queue.Empty:
                    break
                    
            if updates_processed == 0:
                continue
                
            # Save out compiled states and issue a single master API upload call
            live_df.to_csv(METADATA_FILE, index=False)
            
            api.upload_file(
                path_or_fileobj=METADATA_FILE,
                path_in_repo=METADATA_FILE,
                repo_id=REPO_ID,
                repo_type=REPO_TYPE,
                commit_message=f"🤖 ANDREWBAWITLUNG Automated batch update of {updates_processed} task metrics."
            )
            print(f"✅ Successfully pushed batch update containing {updates_processed} items to HF.")
            
        except Exception as e:
            print(f"❌ Critical failure occurred in the batch commit loop execution: {e}")

# Start the automated database syncing daemon
updater_thread = threading.Thread(target=batch_commit_worker, daemon=True)
updater_thread.start()

# --- MEMORY QUEUE INGESTION ---
def execute_action(action, updated_text, task_id, username):
    """Instead of triggering API hits immediately, stash the event in the RAM buffer."""
    write_queue.put({
        "id": str(task_id),
        "text": updated_text,
        "status": action,
        "validated_by": username
    })
    print(f"📦 Modification appended locally for {task_id} via user {username}. Memory pool depth: {write_queue.qsize()}")

def process_action(action_type, updated_text, task_id, history, session_id, request: gr.Request):
    if not task_id:
        next_wav, next_txt, next_id, _ = get_next_task_from_buffer(session_id)
        return next_wav, next_txt, next_id, history, "No active target file available."
        
    username = request.username if request and request.username else "Unknown_Intern"
    execute_action(action_type, updated_text, task_id, username)
    
    if task_id not in history:
        history.append(task_id)
        
    # --- FAST CACHING LOGIC ---
    if session_id not in session_caches:
        session_caches[session_id] = {}

    # Save the current state to RAM
    session_caches[session_id][task_id] = {
        "wav": f"current_{task_id}.wav",
        "text": updated_text
    }
    
    # Keep the cache lightweight (only store the last 5 items)
    if len(session_caches[session_id]) > 5:
        oldest = list(session_caches[session_id].keys())[0]
        del session_caches[session_id][oldest]
    # --------------------------
        
    next_wav, next_txt, next_id, _ = get_next_task_from_buffer(session_id)
    
    if action_type == "VALIDATED":
        msg = f"✅ '{task_id}' processed as VALIDATED locally. Changes will sync shortly."
    else:
        msg = f"🗑️ '{task_id}' processed as Discarded locally. Changes will sync shortly."
        
    return next_wav, next_txt, next_id, history, msg

def handle_validate(updated_text, task_id, history, session_id, request: gr.Request):
    return process_action("VALIDATED", updated_text, task_id, history, session_id, request)

def handle_discard(updated_text, task_id, history, session_id, request: gr.Request):
    return process_action("Discarded", updated_text, task_id, history, session_id, request)

def go_back(history, current_task, session_id):
    if not history:
        return None, gr.update(), current_task, history, "⚠️ No historical operations logged."
        
    last_item = str(history.pop())
    
    # --- FAST CACHE RETRIEVAL ---
    if session_id in session_caches and last_item in session_caches[session_id]:
        cached = session_caches[session_id][last_item]
        if os.path.exists(cached["wav"]):
            return cached["wav"], cached["text"], last_item, history, f"⚡ Instant Reload: {last_item}."
    # ----------------------------

    # Fallback to slow network call if it wasn't in the cache
    local_wav_path = f"current_{last_item}.wav"
    
    if os.path.exists(local_wav_path):
        current_df = fetch_latest_tracker()
        matching_rows = current_df[current_df["id"] == last_item]
        if not matching_rows.empty:
            text_content = matching_rows.iloc[0]["text"]
            if pd.isna(text_content):
                text_content = ""
            return local_wav_path, str(text_content), last_item, history, f"↩️ Network Fallback Reload: {last_item}."
            
    n_wav, n_txt, n_id, _ = get_next_task_from_buffer(session_id)
    return n_wav, n_txt, n_id, history, f"↩️ Network Fallback Reloaded: {n_id}."

# --- CUSTOM AUTH LOGIC ---
def intern_auth(username, password):
    return password == APP_PASSWORD

# --- GRADIO INTERFACE LAYOUT ---
with gr.Blocks(title="ASR Speech Validation Dashboard") as demo:
    gr.Markdown("# 🎙️ Primary Speech Validation Workflow")
    
    session_id = gr.State(lambda: str(uuid.uuid4()))
    current_file_state = gr.State("")
    session_history = gr.State([])
    
    with gr.Row():
        with gr.Column(scale=2):
            audio_player = gr.Audio(
                label="Listen to Audio", 
                type="filepath", 
                interactive=False, 
                autoplay=True
            )
            
            text_editor = gr.Textbox(label="Transcript (Edit if necessary)", lines=4)
            
            with gr.Row():
                back_btn = gr.Button("↩️ Back", variant="secondary")
                discard_btn = gr.Button("✂️ Needs Trim (Discard)", variant="stop")
                validate_btn = gr.Button("✅ Approve & Validate", variant="primary")
                
            status_msg = gr.Markdown("")
            
        with gr.Column(scale=1):
            gr.Markdown("### 🛠️ Instructions")
            gr.Markdown(
                "1. **Listen:** Play the audio clip. Check for background noise, cut-offs, or bad quality.\n"
                "2. **Read/Edit:** Verify that the transcript exactly matches the audio. Fix any typos directly in the text box.\n"
                "3. **Approve:** If the audio is clean and the transcript is accurate, click `Approve & Validate`.\n"
                "4. **Discard:** If the audio has leading/trailing silence or needs to be cut, click `Needs Trim`. This sends it to the Cutter Space for another intern to fix."
            )
            
    demo.load(fn=get_next_task_from_buffer, inputs=[session_id], outputs=[audio_player, text_editor, current_file_state, status_msg])
    
    validate_btn.click(
        fn=handle_validate, 
        inputs=[text_editor, current_file_state, session_history, session_id], 
        outputs=[audio_player, text_editor, current_file_state, session_history, status_msg]
    )
    
    discard_btn.click(
        fn=handle_discard, 
        inputs=[text_editor, current_file_state, session_history, session_id], 
        outputs=[audio_player, text_editor, current_file_state, session_history, status_msg]
    )
    
    back_btn.click(
        fn=go_back, 
        inputs=[session_history, current_file_state, session_id], 
        outputs=[audio_player, text_editor, current_file_state, session_history, status_msg]
    )

if __name__ == "__main__":
    demo.launch(
        auth=intern_auth, 
        auth_message="Please enter your ENROLLMENT NUMBER as the Username, and the shared team password."
    )