"""Transcribe downloaded recordings via Whisper and extract the chosen company.

Reads from:  data/recordings/
Writes to:   data/processed/transcripts.csv

Filename format: {session_code}_{participant_code}_{phase}.{wav|webm}
  phase is one of: baseline, prelim, final

Response extraction (prelim / final only): 'A' / 'B' (regex first, GPT-4o-mini
fallback). The baseline clip carries no decision, so its response is left blank.

Already-transcribed files are skipped on re-runs so the script is safe to resume.
"""

import csv
import os
import re
import sys
from pathlib import Path

import httpx
from dotenv import load_dotenv

load_dotenv(Path.home() / ".config" / "otree" / ".env")

DATA_DIR       = Path(__file__).parent.parent.parent / 'data'
RECORDINGS_DIR = DATA_DIR / 'recordings'
OUT_CSV        = DATA_DIR / 'processed' / 'transcripts.csv'
WHISPER_KEY    = os.environ.get('WHISPER_KEY')

PHASES = ('baseline', 'prelim', 'final')

FIELDNAMES = ['session_code', 'participant_code', 'phase',
              'filename', 'transcript', 'response', 'response_method']

if not WHISPER_KEY:
    print('Error: WHISPER_KEY not set in .env')
    sys.exit(1)

if not RECORDINGS_DIR.exists():
    print(f'Recordings directory not found: {RECORDINGS_DIR}')
    sys.exit(1)

# ── Resume support ────────────────────────────────────────────────────────────
done = set()
if OUT_CSV.exists():
    with open(OUT_CSV, newline='') as f:
        done = {row['filename'] for row in csv.DictReader(f)}
    print(f'Resuming — {len(done)} file(s) already processed.')

files = sorted(
    f for f in RECORDINGS_DIR.iterdir()
    if f.suffix in {'.wav', '.webm'} and f.name not in done
)

if not files:
    print('Nothing to transcribe.')
    sys.exit(0)

print(f'Processing {len(files)} file(s)...\n')


# ── Response extraction ───────────────────────────────────────────────────────

_A_RE = re.compile(r'\bcompany\s*a\b|\boption\s*a\b|\bthe\s+a\b|\ba\.?$', re.I)
_B_RE = re.compile(r'\bcompany\s*b\b|\boption\s*b\b|\bthe\s+b\b|\bb\.?$', re.I)


def extract_response_regex(transcript):
    """Return 'A', 'B', or None if ambiguous / absent."""
    a = bool(_A_RE.search(transcript))
    b = bool(_B_RE.search(transcript))
    if a and not b:
        return 'A'
    if b and not a:
        return 'B'
    return None


def extract_response_gpt(transcript):
    """Ask GPT-4o-mini which company was chosen. Returns 'A', 'B', or 'unclear'."""
    system = ('You extract a decision from a transcript in which a participant says '
              'which of two shipping companies ("Company A" or "Company B") they '
              'would choose for their online order. Reply with only "A", "B", or "unclear".')
    valid = {'A', 'B'}
    try:
        resp = httpx.post(
            'https://api.openai.com/v1/chat/completions',
            headers={'Authorization': f'Bearer {WHISPER_KEY}'},
            json={
                'model': 'gpt-4o-mini',
                'messages': [
                    {'role': 'system', 'content': system},
                    {'role': 'user', 'content': transcript},
                ],
                'max_tokens': 3,
                'temperature': 0,
            },
            timeout=30,
        )
        resp.raise_for_status()
        answer = resp.json()['choices'][0]['message']['content'].strip().upper()
        return answer if answer in valid else 'unclear'
    except Exception as e:
        print(f'    GPT error: {e}')
        return 'unclear'


# ── Main loop ─────────────────────────────────────────────────────────────────
rows = []
gpt_calls = 0

for f in files:
    parts = f.stem.split('_')  # e.g. abc12def_gh3ij456_prelim
    session_code = parts[0] if len(parts) > 0 else ''
    participant_code = parts[1] if len(parts) > 1 else ''
    phase = parts[2] if len(parts) > 2 else ''
    mimetype = 'audio/wav' if f.suffix == '.wav' else 'audio/webm'

    try:
        with open(f, 'rb') as audio:
            resp = httpx.post(
                'https://api.openai.com/v1/audio/transcriptions',
                headers={'Authorization': f'Bearer {WHISPER_KEY}'},
                files={'file': (f.name, audio, mimetype)},
                data={'model': 'whisper-1'},
                timeout=60,
            )
        resp.raise_for_status()
        transcript = resp.json().get('text', '').replace('\n', ' ').replace('\r', ' ').strip()
    except Exception as e:
        transcript = ''
        print(f'  {f.name} — Whisper ERROR: {e}')

    if phase in ('prelim', 'final'):
        response = extract_response_regex(transcript)
        if response is not None:
            response_method = 'regex'
        else:
            response = extract_response_gpt(transcript)
            response_method = 'gpt'
            gpt_calls += 1
    else:
        response, response_method = '', ''

    print(f'  {f.name}')
    print(f'  → "{transcript[:100]}"')
    print(f'  → response: {response} ({response_method})\n')

    rows.append({
        'session_code':     session_code,
        'participant_code': participant_code,
        'phase':            phase,
        'filename':         f.name,
        'transcript':       transcript,
        'response':         response,
        'response_method':  response_method,
    })

# ── Write CSV ─────────────────────────────────────────────────────────────────
OUT_CSV.parent.mkdir(parents=True, exist_ok=True)
write_header = not OUT_CSV.exists()
with open(OUT_CSV, 'a', newline='', encoding='utf-8') as f:
    writer = csv.DictWriter(f, fieldnames=FIELDNAMES)
    if write_header:
        writer.writeheader()
    writer.writerows(rows)

print(f'Done. {len(rows)} file(s) processed ({gpt_calls} GPT fallback(s)).')
print(f'Saved to {OUT_CSV}')