File size: 32,167 Bytes
ae0b416
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
import os
import json
import requests
import subprocess
import shutil
import time
import sys
import threading
from typing import Dict, List, Optional, Any
from fastapi import FastAPI, HTTPException
from fastapi.responses import JSONResponse
import uvicorn
import torch
import librosa
from transformers import AutoModelForSpeechSeq2Seq, AutoProcessor, pipeline

# Fix Unicode encoding for Windows
if sys.platform == 'win32':
    import io
    sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding='utf-8')

# Initialize FastAPI app
app = FastAPI(title="Audio Transcriber", description="Audio transcription and upload service")

# ==== CONFIGURATION ====
# The new backend URL for state management and transcription upload
# It is now read from an environment variable, falling back to the default if not set.
BACKEND_URL = os.environ.get("BACKEND_URL", "https://samfredoly-acp.hf.space")
# The original Hugging Face repo IDs are still needed for fetching the audio files
# and the reference file list, as the backend only handles transcription storage.
SOURCE_REPO_ID = "Samfredoly/BG_Vid"  # Fetch audio files from here
TARGET_REPO_ID = "samfred2/A_Text"  # Target repo ID is now a constant for the backend
REFERENCE_REPO_ID = "Fred808/BG3"  # Reference repo to match audio filenames

# Path Configuration
DOWNLOAD_FOLDER = "downloads_audio"
TRANSCRIPTIONS_FOLDER = "transcriptions"
LOCAL_STATE_FOLDER = ".state_audio"

os.makedirs(DOWNLOAD_FOLDER, exist_ok=True)
os.makedirs(TRANSCRIPTIONS_FOLDER, exist_ok=True)
os.makedirs(LOCAL_STATE_FOLDER, exist_ok=True)

# Whisper Model Setup (using transformers)
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
TORCH_DTYPE = torch.float16 if torch.cuda.is_available() else torch.float32
WHISPER_MODEL_ID = f"openai/whisper-small"

# Global model cache
_whisper_model = None
_whisper_processor = None
_whisper_pipeline = None

def get_whisper_pipeline():
    """Get or initialize the Whisper pipeline."""
    global _whisper_model, _whisper_processor, _whisper_pipeline
    
    if _whisper_pipeline is not None:
        return _whisper_pipeline
    
    try:
        log_message(f"Loading Whisper model {WHISPER_MODEL_ID}...", "INFO")
        
        model = AutoModelForSpeechSeq2Seq.from_pretrained(
            WHISPER_MODEL_ID,
            torch_dtype=TORCH_DTYPE,
            low_cpu_mem_usage=True,
            use_safetensors=True
        )
        model = model.to(DEVICE)
        
        processor = AutoProcessor.from_pretrained(WHISPER_MODEL_ID)
        
        _whisper_pipeline = pipeline(
            "automatic-speech-recognition",
            model=model,
            tokenizer=processor.tokenizer,
            feature_extractor=processor.feature_extractor,
            torch_dtype=TORCH_DTYPE,
            device=DEVICE
        )
        
        log_message(f"βœ… Whisper model loaded successfully on {DEVICE.upper()}", "INFO")
        return _whisper_pipeline
        
    except Exception as e:
        log_message(f"❌ Failed to load Whisper model: {str(e)}", "ERROR")
        raise

# State Files
FAILED_FILES_LOG = "failed_audio_files.log"
HF_STATE_FILE = "processing_audio_state.json" # This is the filename the backend uses

# Processing Parameters
PROCESSING_DELAY = 2
MAX_RETRIES = 3
MIN_FREE_SPACE_GB = 1
WHISPER_MODEL = "small"  # Whisper model size

# NOTE: The Hugging Face API is still required for listing files in SOURCE_REPO_ID and REFERENCE_REPO_ID
from huggingface_hub import HfApi, hf_hub_url
HF_TOKEN = os.environ.get("HF_TOKEN", "")
hf_api = HfApi(token=HF_TOKEN)

# Global State
processing_status = {
    "is_running": False,
    "current_file": None,
    "total_files": 0,
    "processed_files": 0,
    "failed_files": 0,
    "transcribed_files": 0,
    "last_update": None,
    "logs": []
}

def log_message(message: str, level: str = "INFO"):
    """Log messages with timestamp"""
    timestamp = time.strftime("%Y-%m-%d %H:%M:%S")
    log_entry = f"[{timestamp}] {level}: {message}"
    print(log_entry)
    processing_status["logs"].append(log_entry)
    processing_status["last_update"] = timestamp
    if len(processing_status["logs"]) > 100:
        processing_status["logs"] = processing_status["logs"][-100:]

def log_failed_file(filename: str, error: str):
    """Log failed files to persistent file"""
    with open(FAILED_FILES_LOG, "a") as f:
        f.write(f"{time.strftime('%Y-%m-%d %H:%M:%S')} - {filename}: {error}\n")

def get_disk_usage(path: str) -> Dict[str, float]:
    """Get disk usage statistics in GB"""
    statvfs = os.statvfs(path)
    total = statvfs.f_frsize * statvfs.f_blocks / (1024**3)
    free = statvfs.f_frsize * statvfs.f_bavail / (1024**3)
    used = total - free
    return {"total": total, "free": free, "used": used}

def check_disk_space(path: str = ".") -> bool:
    """Check if there's enough disk space"""
    disk_info = get_disk_usage(path)
    if disk_info["free"] < MIN_FREE_SPACE_GB:
        log_message(f'⚠️ Low disk space: {disk_info["free"]:.2f}GB free, {disk_info["used"]:.2f}GB used')
        return False
    return True

def cleanup_temp_files():
    """Clean up temporary files to free space"""
    log_message("🧹 Cleaning up temporary files...", "INFO")
    
    current_file = processing_status.get("current_file")
    for file in os.listdir(DOWNLOAD_FOLDER):
        if file != current_file and file.endswith((".wav", ".mp3")):
            try:
                os.remove(os.path.join(DOWNLOAD_FOLDER, file))
                log_message(f"πŸ—‘οΈ Removed old download: {file}", "INFO")
            except:
                pass

# Helper function to save state locally
def save_json_state(file_path: str, data: Dict[str, Any]):
    """Save state to JSON file"""
    with open(file_path, "w") as f:
        json.dump(data, f, indent=2)

# --- NEW API FUNCTIONS FOR STATE MANAGEMENT AND UPLOAD ---

def download_state_from_api() -> Dict[str, Any]:
    """Downloads the state file from the backend API."""
    url = f"{BACKEND_URL}/state/"
    default_state = {"next_download_index": 0, "file_states": {}}
    
    try:
        response = requests.get(url, timeout=10)
        response.raise_for_status()
        
        # The API returns {"state": {...}}
        state_data = response.json().get("state", default_state)
        
        # Ensure the structure is correct (migration logic from original load_json_state)
        if "file_states" not in state_data or not isinstance(state_data["file_states"], dict):
            state_data["file_states"] = {}
        if "next_download_index" not in state_data:
            state_data["next_download_index"] = 0
        
        log_message(f"βœ… Downloaded state: next_download_index={state_data['next_download_index']}, processed_files={len([f for f,s in state_data['file_states'].items() if s=='processed'])}", "INFO")
        return state_data
        
    except requests.exceptions.RequestException as e:
        log_message(f"⚠️ Failed to download state from API ({url}): {str(e)}. Starting from default state.", "WARNING")
        return default_state

def upload_state_to_api(state: Dict[str, Any]) -> bool:
    """

    Saves the state locally and uploads it to the backend API's /upload/ endpoint.

    This simulates the original HF state upload for locking/unlocking.

    """
    local_path = os.path.join(LOCAL_STATE_FOLDER, HF_STATE_FILE)
    url = f"{BACKEND_URL}/upload/"
    
    try:
        # 1. Save the current state locally
        save_json_state(local_path, state)
        
        # 2. Upload the state file to the backend
        with open(local_path, "rb") as f:
            files = {'file': (HF_STATE_FILE, f, 'application/json')}
            
            response = requests.post(url, files=files, timeout=30)
            response.raise_for_status()
            
            log_message(f"βœ… Successfully uploaded state file to API: {HF_STATE_FILE}", "INFO")
            return True
            
    except requests.exceptions.HTTPError as e:
        if hasattr(e, 'response') and e.response.status_code == 409:
            log_message(f"⚠️ State file already exists on server (409 Conflict) - Treating as success.", "INFO")
            return True
        log_message(f"❌ Failed to upload state file to API ({url}): {str(e)}", "ERROR")
        return False
    except requests.exceptions.RequestException as e:
        log_message(f"❌ Failed to upload state file to API ({url}): {str(e)}", "ERROR")
        return False
    except Exception as e:
        log_message(f"❌ An unexpected error occurred during API state upload: {str(e)}", "ERROR")
        return False

def upload_transcription_to_api(json_output_path: str, matched_filename: str) -> bool:
    """Uploads the transcription JSON file to the backend API's /upload/ endpoint."""
    url = f"{BACKEND_URL}/upload/"
    
    try:
        with open(json_output_path, "rb") as f:
            files = {'file': (os.path.basename(json_output_path), f, 'application/json')}
            
            response = requests.post(url, files=files, timeout=30)
            response.raise_for_status()
            
            log_message(f"βœ… Successfully uploaded transcription to API: {os.path.basename(json_output_path)}", "INFO")
            return True
            
    except requests.exceptions.HTTPError as e:
        if hasattr(e, 'response') and e.response.status_code == 409:
            log_message(f"⚠️ File already exists on server (409 Conflict) - Treating as success.", "INFO")
            return True
        log_message(f"❌ Failed to upload transcription to API ({url}): {str(e)}", "ERROR")
        return False
    except requests.exceptions.RequestException as e:
        log_message(f"❌ Failed to upload transcription to API ({url}): {str(e)}", "ERROR")
        return False
    except Exception as e:
        log_message(f"❌ An unexpected error occurred during API upload: {str(e)}", "ERROR")
        return False

def lock_file_for_processing(wav_filename: str, state: Dict[str, Any]) -> bool:
    """Marks a file as 'processing' in the state file and uploads the lock via API."""
    log_message(f"πŸ”’ Attempting to lock file: {wav_filename} (Marking as 'processing')", "INFO")
    
    state["file_states"][wav_filename] = "processing"
    
    if upload_state_to_api(state):
        log_message(f"βœ… Successfully locked file: {wav_filename} via API state upload", "INFO")
        return True
    else:
        log_message(f"❌ Failed to upload lock for file: {wav_filename}. Aborting processing.", "ERROR")
        # Revert local state change if upload fails
        if wav_filename in state["file_states"]:
            del state["file_states"][wav_filename]
        return False

def unlock_file_as_processed(wav_filename: str, state: Dict[str, Any], next_index: int) -> bool:
    """Marks a file as 'processed', updates the index, and uploads the state via API."""
    log_message(f"πŸ”“ Attempting to unlock file: {wav_filename} (Marking as 'processed')", "INFO")
    
    state["file_states"][wav_filename] = "processed"
    state["next_download_index"] = next_index
    
    if upload_state_to_api(state):
        log_message(f"βœ… Successfully unlocked and marked as processed: {wav_filename} via API state upload", "INFO")
        return True
    else:
        log_message(f"❌ Failed to upload final state for file: {wav_filename}.", "ERROR")
        return False

# --- END NEW API FUNCTIONS ---

def download_with_retry(url: str, dest_path: str, max_retries: int = 3) -> bool:
    """Download file with retry logic and disk space checking"""
    if not check_disk_space():
        cleanup_temp_files()
        if not check_disk_space():
            log_message("❌ Insufficient disk space even after cleanup", "ERROR")
            return False
    
    try:
        os.makedirs(os.path.dirname(dest_path), exist_ok=True)
    except Exception as e:
        log_message(f"❌ Failed to create directory for download path {os.path.dirname(dest_path)}: {str(e)}", "ERROR")
        return False
        
    # The original code used HF_TOKEN for authorization headers, which is only needed 
    # if the source repo is private. We keep it for compatibility.
    headers = {"Authorization": f"Bearer {HF_TOKEN}"}
    for attempt in range(max_retries):
        try:
            with requests.get(url, headers=headers, stream=True) as r:
                r.raise_for_status()
                
                with open(dest_path, "wb") as f:
                    for chunk in r.iter_content(chunk_size=8192):
                        if chunk:
                            f.write(chunk)
                
                log_message(f"βœ… Download successful: {os.path.basename(dest_path)}", "INFO")
                return True
        except requests.exceptions.RequestException as e:
            log_message(f"⚠️ Download attempt {attempt + 1}/{max_retries} failed for {url}: {str(e)}", "WARNING")
            if attempt < max_retries - 1:
                time.sleep(2 ** attempt) # Exponential backoff
            else:
                log_message(f"❌ Download failed after {max_retries} attempts for {url}", "ERROR")
                return False
        except Exception as e:
            log_message(f"❌ An unexpected error occurred during download: {str(e)}", "ERROR")
            return False
    return False

def get_reference_map(reference_repo_id: str) -> Dict[str, str]:
    """

    Downloads the reference file list from the Hugging Face repo and creates a map

    from audio filename (without extension) to the reference filename.

    """
    log_message(f"Fetching reference file list from {reference_repo_id}...", "INFO")
    
    # This is a placeholder for the actual logic to get the file list.
    # Assuming the reference repo contains a list of files that match the audio files.
    # In a real scenario, this would involve listing files in the repo.
    # For now, we'll assume a simple list of files can be retrieved.
    
    try:
        # Use HfApi to list files in the reference repo
        repo_files = hf_api.list_repo_files(repo_id=reference_repo_id, repo_type="dataset")
        
        reference_map = {}
        for file in repo_files:
            # Assuming the reference files are named like 'audio_file_name.txt'
            # and we want to map the audio file name (e.g., 'audio_file_name.wav') to it.
            base_name, ext = os.path.splitext(file)
            if ext.lower() in ['.txt', '.json']: # Only consider text/json files as reference
                # The key is the audio file name without extension
                reference_map[base_name] = file
                
        log_message(f"βœ… Successfully created reference map with {len(reference_map)} entries.", "INFO")
        return reference_map
        
    except Exception as e:
        log_message(f"❌ Failed to fetch reference map from Hugging Face: {str(e)}", "ERROR")
        return {}

def find_matching_filename(audio_filename: str, reference_map: Dict[str, str]) -> Optional[str]:
    """Finds the matching reference filename for a given audio filename."""
    base_name, _ = os.path.splitext(audio_filename)
    return reference_map.get(base_name)

def get_next_file_to_process(source_repo_id: str, state: Dict[str, Any]) -> Optional[Dict[str, Any]]:
    """

    Determines the next file to process based on the current state and the file list

    from the source Hugging Face repository.

    """
    log_message(f"Determining next file to process from {source_repo_id}...", "INFO")
    
    try:
        # 1. Get the list of all files in the source repo
        repo_files = hf_api.list_repo_files(repo_id=source_repo_id, repo_type="dataset")
        
        # Filter for audio files (e.g., .wav, .mp3)
        audio_files = sorted([f for f in repo_files if f.lower().endswith(('.wav', '.mp3'))])
        
        processing_status["total_files"] = len(audio_files)
        
        if not audio_files:
            log_message("No audio files found in the source repository.", "INFO")
            return None
            
        # 2. Get the next index from the state
        next_index = state.get("next_download_index", 0)
        file_states = state.get("file_states", {})
        
        # 3. Skip forward past all processed and processing files starting from next_index
        # This ensures we don't repeatedly find files that have already been handled
        current_index = next_index
        while current_index < len(audio_files):
            filename = audio_files[current_index]
            status = file_states.get(filename, "unprocessed")
            
            # If this file is processed or currently processing, skip it
            if status in ["processed", "processing"]:
                current_index += 1
                continue
            
            # If this file failed, we can retry it, so return it
            if status == "failed":
                file_url = hf_hub_url(repo_id=source_repo_id, filename=filename, repo_type="dataset")
                log_message(f"Found failed file for retry at index {current_index}: {filename}", "INFO")
                return {
                    "filename": filename,
                    "url": file_url,
                    "index": current_index
                }
            
            # If this file is unprocessed, we found our next file
            file_url = hf_hub_url(repo_id=source_repo_id, filename=filename, repo_type="dataset")
            log_message(f"Found next file at index {current_index}: {filename}", "INFO")
            return {
                "filename": filename,
                "url": file_url,
                "index": current_index
            }
        
        log_message("All files have been processed or are locked. Checking for any failed files from the start.", "INFO")
        
        # 4. If we've processed all files from next_index to end, check from beginning for failed files
        for i in range(0, next_index):
            filename = audio_files[i]
            status = file_states.get(filename, "unprocessed")
            
            if status == "failed":
                file_url = hf_hub_url(repo_id=source_repo_id, filename=filename, repo_type="dataset")
                log_message(f"Found failed file for retry at index {i}: {filename}", "INFO")
                return {
                    "filename": filename,
                    "url": file_url,
                    "index": i
                }
        
        log_message("All files have been processed. Waiting for new files...", "INFO")
        return None
        
    except Exception as e:
        log_message(f"❌ Failed to get next file to process: {str(e)}", "ERROR")
        return None

def run_whisper_transcription(audio_path: str, output_dir: str, model: str) -> Optional[str]:
    """

    Runs Whisper transcription using the transformers library.

    Returns the path to the generated JSON file on success.

    No ffmpeg dependency required.

    """
    log_message(f"πŸŽ™οΈ Starting transcription for {os.path.basename(audio_path)} with model {model}...", "INFO")
    
    try:
        # Get the Whisper pipeline
        pipe = get_whisper_pipeline()
        
        # Load audio using librosa
        log_message(f"Loading audio file: {audio_path}", "INFO")
        audio_data, sample_rate = librosa.load(audio_path, sr=16000)
        
        # Run transcription
        log_message(f"Running transcription...", "INFO")
        result = pipe(
            audio_data,
            chunk_length_s=30,
            batch_size=8,
            return_timestamps=True
        )
        
        # Extract text and chunks
        transcription_text = result.get("text", "")
        chunks = result.get("chunks", [])
        
        log_message(f"βœ… Transcription successful: {len(transcription_text)} characters", "INFO")
        
        # Prepare output JSON structure
        output_json = {
            "text": transcription_text,
            "chunks": chunks,
            "language": result.get("language", "en")
        }
        
        # Save to JSON file
        base_name, _ = os.path.splitext(os.path.basename(audio_path))
        json_output_path = os.path.join(output_dir, f"{base_name}.json")
        
        with open(json_output_path, "w", encoding="utf-8") as f:
            json.dump(output_json, f, indent=2, ensure_ascii=False)
        
        log_message(f"βœ… Saved transcription to: {json_output_path}", "INFO")
        return json_output_path
        
    except Exception as e:
        log_message(f"❌ An error occurred during transcription: {str(e)}", "ERROR")
        import traceback
        log_message(f"Traceback: {traceback.format_exc()}", "ERROR")
        return None

def process_audio_file(audio_path: str, reference_map: Dict[str, str], output_filename: str) -> bool:
    """

    Transcribes the audio file, renames the output JSON to match the reference,

    and uploads the result to the API.

    """
    
    # 1. Run transcription
    json_output_path = run_whisper_transcription(audio_path, TRANSCRIPTIONS_FOLDER, WHISPER_MODEL)
    
    if not json_output_path:
        return False
        
    # 2. Rename the JSON file to the matched filename
    # The output_filename already includes the correct extension (e.g., .txt or .json)
    # We assume the reference map provides the full target filename.
    
    # The whisper output is a JSON file named after the audio file.
    # We need to rename it to the target filename (which should be a JSON file for the backend).
    
    # The output_filename is the matched filename from the reference map (e.g., 'audio_file_name.txt')
    # The backend expects a JSON file. Let's assume the matched filename should be used as the base
    # but with a .json extension for the upload.
    
    # Let's stick to the original logic: the backend expects a JSON file with the name
    # of the audio file (or the matched reference file) with a .json extension.
    
    # Since the whisper output is already a JSON file, we just need to rename it
    # to the desired final name.
    
    # The output_filename passed here is the base name of the audio file or the matched reference file.
    # If it's a reference file name (e.g., 'file.txt'), we should probably use 'file.json'.
    
    # For simplicity and to match the backend's expectation (which handles JSON),
    # we will rename the whisper output JSON to the base name of the audio file
    # and ensure it has a .json extension.
    
    base_name, _ = os.path.splitext(output_filename)
    final_json_filename = f"{base_name}.json"
    final_json_path = os.path.join(TRANSCRIPTIONS_FOLDER, final_json_filename)
    
    try:
        if json_output_path != final_json_path:
            shutil.move(json_output_path, final_json_path)
            log_message(f"βœ… Renamed transcription to: {final_json_filename}", "INFO")
    except Exception as e:
        log_message(f"❌ Failed to rename transcription file: {str(e)}", "ERROR")
        return False
        
    # 3. Upload transcription to API
    if upload_transcription_to_api(final_json_path, final_json_filename):
        processing_status["transcribed_files"] += 1
        # Clean up the local transcription file after successful upload
        try:
            os.remove(final_json_path)
            log_message(f"πŸ—‘οΈ Cleaned up local transcription file: {final_json_path}", "INFO")
        except Exception as e:
            log_message(f"❌ Failed to clean up transcription file: {str(e)}", "ERROR")
        return True
    else:
        log_message(f"❌ Failed to upload transcription to API: {final_json_filename}", "ERROR")
        return False

def main_processing_loop():
    """The main loop that continuously checks for and processes new audio files."""
    global processing_status
    
    if processing_status["is_running"]:
        log_message("Processing loop is already running.", "WARNING")
        return
        
    processing_status["is_running"] = True
    log_message("πŸš€ Audio transcription processing loop started.", "INFO")
    
    # 1. Get the reference map once
    reference_map = get_reference_map(REFERENCE_REPO_ID)
    if not reference_map:
        log_message("❌ Could not get reference map. Stopping loop.", "CRITICAL")
        processing_status["is_running"] = False
        return
        
    try:
        while processing_status["is_running"]:
            time.sleep(PROCESSING_DELAY)
            
            # 1. Download FRESH state from the API at the start of each iteration
            # This ensures we respect the next_download_index that other workers may have set
            current_state = download_state_from_api()
            next_file_info = get_next_file_to_process(SOURCE_REPO_ID, current_state)
            
            if next_file_info is None:
                log_message("πŸ’€ No new audio files to process. Sleeping for a while...", "INFO")
                time.sleep(PROCESSING_DELAY * 5)
                continue
                
            target_file = next_file_info['filename']
            audio_url = next_file_info['url']
            target_index = next_file_info['index']
            
            processing_status["current_file"] = target_file
            success = False
            matched_filename = None
            
            try:
                # 2. Lock file by updating state on the API
                # IMPORTANT: Update next_download_index when locking to prevent other workers from picking same file
                old_index = current_state["next_download_index"]
                current_state["next_download_index"] = target_index + 1
                log_message(f"πŸ“ Incrementing next_download_index from {old_index} to {current_state['next_download_index']}", "INFO")
                
                if not lock_file_for_processing(target_file, current_state):
                    log_message(f"❌ Failed to lock file {target_file}. Skipping.", "ERROR")
                    time.sleep(PROCESSING_DELAY)
                    continue
                    
                local_wav_path = os.path.join(DOWNLOAD_FOLDER, os.path.basename(target_file))
                log_message(f"⬇️ Downloading audio file: {target_file}", "INFO")
                
                if download_with_retry(audio_url, local_wav_path):
                    
                    # Extract base filename for matching
                    base_filename = os.path.basename(target_file)
                    matched_filename = find_matching_filename(base_filename, reference_map)
                    
                    # Use matched filename if found, otherwise use original filename
                    output_filename = matched_filename if matched_filename else base_filename
                    
                    # 3. Process and Upload transcription to API
                    if process_audio_file(local_wav_path, reference_map, output_filename):
                        success = True
                        log_message(f"βœ… Finished processing: {target_file}", "INFO")
                    else:
                        log_message(f"❌ Processing failed for: {target_file}", "ERROR")
                else:
                    log_message(f"❌ Download failed for: {target_file}", "ERROR")
                    
            except Exception as e:
                log_message(f"πŸ”₯ An unhandled error occurred while processing {target_file}: {str(e)}", "ERROR")
                log_failed_file(target_file, str(e))
                
            finally:
                # 4. Unlock/Mark as processed by updating state on the API
                # IMPORTANT: Keep the incremented next_download_index from locking
                
                if success:
                    # Mark as processed and keep the incremented index, then upload state
                    unlock_file_as_processed(target_file, current_state, current_state["next_download_index"])
                    processing_status["processed_files"] += 1
                else:
                    # Mark as failed but keep the incremented index so next worker can proceed
                    log_message(f"⚠️ File {target_file} failed. Marking as 'failed' and updating state.", "WARNING")
                    current_state["file_states"][target_file] = "failed"
                    # Keep the incremented next_download_index - don't change it
                    upload_state_to_api(current_state)
                
                # Clean up the downloaded audio file regardless of success
                try:
                    if os.path.exists(local_wav_path):
                        os.remove(local_wav_path)
                        log_message(f"πŸ—‘οΈ Cleaned up local audio file: {local_wav_path}", "INFO")
                except Exception as e:
                    log_message(f"❌ Failed to clean up audio file: {str(e)}", "ERROR")
                
                processing_status["current_file"] = None
                time.sleep(PROCESSING_DELAY)
                
    except Exception as e:
        log_message(f"πŸ”₯ Critical error in main processing loop: {str(e)}", "CRITICAL")
        
    finally:
        processing_status["is_running"] = False
        log_message("πŸ›‘ Audio transcription processing loop stopped.", "INFO")

# --- FastAPI Endpoints (Unchanged) ---

# Add to configuration section
AUTO_START_PROCESSING = os.environ.get("AUTO_START_PROCESSING", "true").lower() == "true"

@app.on_event("startup")
async def startup_event():
    """Conditionally start processing based on environment variable."""
    if AUTO_START_PROCESSING:
        log_message("πŸš€ AUTO_START_PROCESSING enabled - Starting processing loop...", "INFO")
        thread = threading.Thread(target=main_processing_loop, daemon=True)
        thread.start()
        log_message("βœ… Background processing thread started", "INFO")
    else:
        log_message("⏸️ AUTO_START_PROCESSING disabled - Use /start endpoint to begin", "INFO")

@app.get("/")
async def root():
    """Root endpoint to check service status."""
    return {"message": "Audio Transcriber Service is running", "status": processing_status}

@app.get("/status")
async def get_status():
    """Get the current processing status."""
    return processing_status

@app.post("/start")
async def start_processing():
    """Start the background processing loop."""
    if processing_status["is_running"]:
        return JSONResponse(status_code=200, content={"message": "Processing already running."})
    
    thread = threading.Thread(target=main_processing_loop)
    thread.start()
    return JSONResponse(status_code=200, content={"message": "Processing started in background."})

@app.post("/stop")
async def stop_processing():
    """Stop the background processing loop."""
    if not processing_status["is_running"]:
        return JSONResponse(status_code=200, content={"message": "Processing is not running."})
    
    processing_status["is_running"] = False
    return JSONResponse(status_code=200, content={"message": "Processing stop requested. Will stop after current file."})

# --- Main Execution ---

if __name__ == "__main__":
    # This block is for local testing and won't be used in the final sandbox execution
    # but is good practice for a runnable script.
    uvicorn.run(app, host="0.0.0.0", port=8000)