import os
import sys
import time
from pathlib import Path
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: Path, api_key: str, max_retries: int = 3) -> str:
    """Transcribe a single audio file using the AI Builder Space API."""
    headers = {
        "Authorization": f"Bearer {api_key}"
    }

    for attempt in range(1, max_retries + 1):
        try:
            with open(file_path, "rb") as f:
                files = {
                    "audio_file": (file_path.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}"
                if attempt < max_retries and response.status_code in (429, 500, 502, 503, 504):
                    wait_time = 2 ** attempt
                    print(f"Warning: Request failed for {file_path.name} ({error_msg}). Retrying in {wait_time}s...")
                    time.sleep(wait_time)
                    continue
                raise RuntimeError(f"Transcription failed for {file_path.name}: {error_msg}")

        except requests.RequestException as e:
            if attempt < max_retries:
                wait_time = 2 ** attempt
                print(f"Warning: Network error for {file_path.name}: {e}. Retrying in {wait_time}s...")
                time.sleep(wait_time)
                continue
            raise RuntimeError(f"Network error while transcribing {file_path.name}: {e}")

    raise RuntimeError(f"Exceeded max retries for {file_path.name}")


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 = Path("memos")
    if not memos_dir.exists() or not memos_dir.is_dir():
        print(f"Error: Directory '{memos_dir}' does not exist.", file=sys.stderr)
        sys.exit(1)

    transcripts_dir = Path("transcripts")
    transcripts_dir.mkdir(parents=True, exist_ok=True)

    # Find audio files in memos/
    memo_files = sorted([
        f for f in memos_dir.iterdir()
        if f.is_file() and f.suffix.lower() in SUPPORTED_EXTENSIONS
    ])

    if not memo_files:
        print(f"No supported audio files found in '{memos_dir}'.")
        return

    print(f"Found {len(memo_files)} voice memo(s) to transcribe.")

    for idx, memo_file in enumerate(memo_files, start=1):
        target_file = transcripts_dir / f"{memo_file.stem}.txt"
        print(f"[{idx}/{len(memo_files)}] Transcribing {memo_file.name} -> {target_file.name}...")
        try:
            transcript_text = transcribe_file(memo_file, api_key)
            target_file.write_text(transcript_text, encoding="utf-8")
            print(f"Successfully saved transcript to {target_file}")
        except Exception as e:
            print(f"Error transcribing {memo_file.name}: {e}", file=sys.stderr)
            sys.exit(1)

    print("\nAll voice memos transcribed successfully.")


if __name__ == "__main__":
    main()