import os
import sys
import glob
import time
import requests

API_URL = "https://space.ai-builders.com/backend/v1/audio/grok-transcription"
SUPPORTED_EXTENSIONS = {".wav", ".mp3", ".m4a", ".aac", ".ogg", ".flac", ".mp4", ".mkv", ".webm"}

def transcribe_file(file_path: str, api_key: str, max_retries: int = 3) -> str:
    """Transcribe an audio file using AI Builder Space grok-transcription API."""
    headers = {
        "Authorization": f"Bearer {api_key}"
    }
    
    file_name = os.path.basename(file_path)
    
    for attempt in range(1, max_retries + 1):
        try:
            with open(file_path, "rb") as f:
                files = {
                    "audio_file": (file_name, f)
                }
                data = {
                    "simple": "true"
                }
                response = requests.post(API_URL, headers=headers, files=files, data=data, timeout=60)
            
            if response.status_code == 200:
                result = response.json()
                return result.get("text", "")
            else:
                error_msg = f"HTTP {response.status_code}: {response.text}"
                print(f"  [Attempt {attempt}/{max_retries}] Request failed for {file_name}: {error_msg}", file=sys.stderr)
        except Exception as e:
            print(f"  [Attempt {attempt}/{max_retries}] Exception during request for {file_name}: {e}", file=sys.stderr)
        
        if attempt < max_retries:
            time.sleep(2 * attempt)
            
    raise RuntimeError(f"Failed to transcribe {file_name} after {max_retries} attempts.")

def main():
    api_key = os.environ.get("AI_BUILDER_TOKEN")
    if not api_key:
        print("Error: AI_BUILDER_TOKEN environment variable is not set.", file=sys.stderr)
        sys.exit(1)
        
    memos_dir = "memos"
    transcripts_dir = "transcripts"
    
    if not os.path.isdir(memos_dir):
        print(f"Error: Directory '{memos_dir}' does not exist.", file=sys.stderr)
        sys.exit(1)
        
    os.makedirs(transcripts_dir, exist_ok=True)
    
    # List audio files
    audio_files = []
    for entry in os.listdir(memos_dir):
        if entry.startswith("."):
            continue
        ext = os.path.splitext(entry)[1].lower()
        if ext in SUPPORTED_EXTENSIONS:
            audio_files.append(entry)
            
    audio_files.sort()
    
    if not audio_files:
        print(f"No supported audio files found in '{memos_dir}'.")
        return

    print(f"Found {len(audio_files)} audio file(s) in '{memos_dir}'. Starting transcription...")
    
    success_count = 0
    for filename in audio_files:
        input_path = os.path.join(memos_dir, filename)
        stem = os.path.splitext(filename)[0]
        output_path = os.path.join(transcripts_dir, f"{stem}.txt")
        
        print(f"Transcribing {filename}...")
        try:
            transcript_text = transcribe_file(input_path, api_key)
            with open(output_path, "w", encoding="utf-8") as out_f:
                out_f.write(transcript_text)
            print(f"Saved transcript to {output_path} ({len(transcript_text)} characters)")
            success_count += 1
        except Exception as e:
            print(f"Error processing {filename}: {e}", file=sys.stderr)
            
    print(f"Transcription finished: {success_count}/{len(audio_files)} succeeded.")

if __name__ == "__main__":
    main()