import os
from pathlib import Path
from typing import Optional

from .types import TimeAlignedType


def load_text(audio: str | os.PathLike[str]) -> Optional[TimeAlignedType]:
    """Load time-aligned text data given a TIMIT audio file.

    :param audio: Path to the TIMIT audio file.
    :return: A TimeAlignedType object.
    """
    import string

    from sonusai.mixture import get_sample_rate

    file = Path(audio).with_suffix('.TXT')
    if not os.path.exists(file):
        return None

    sample_rate = get_sample_rate(str(audio))

    with open(file, mode='r', encoding='utf-8') as f:
        line = f.read()

    fields = line.strip().split()
    start = int(fields[0]) / sample_rate
    end = int(fields[1]) / sample_rate
    text = ' '.join(fields[2:]).lower().translate(str.maketrans('', '', string.punctuation))

    return TimeAlignedType(start, end, text)


def load_words(audio: str | os.PathLike[str]) -> Optional[list[TimeAlignedType]]:
    """Load time-aligned word data given a TIMIT audio file.

    :param audio: Path to the TIMIT audio file.
    :return: A list of TimeAlignedType objects.
    """

    return _load_ta(audio, 'words')


def load_phonemes(audio: str | os.PathLike[str]) -> Optional[list[TimeAlignedType]]:
    """Load time-aligned phonemes data given a TIMIT audio file.

    :param audio: Path to the TIMIT audio file.
    :return: A list of TimeAlignedType objects.
    """

    return _load_ta(audio, 'phonemes')


def _load_ta(audio: str | os.PathLike[str], tier: str) -> Optional[list[TimeAlignedType]]:
    from sonusai.mixture import get_sample_rate

    if tier == 'words':
        file = Path(audio).with_suffix('.WRD')
    elif tier == 'phonemes':
        file = Path(audio).with_suffix('.PHN')
    else:
        raise ValueError(f'Unknown tier: {tier}')

    if not os.path.exists(file):
        return None

    sample_rate = get_sample_rate(str(audio))

    entries: list[TimeAlignedType] = []
    first = True
    with open(file, mode='r', encoding='utf-8') as f:
        for line in f.readlines():
            fields = line.strip().split()
            start = int(fields[0]) / sample_rate
            end = int(fields[1]) / sample_rate
            text = ' '.join(fields[2:])

            if first:
                first = False
            else:
                if start < entries[-1].end:
                    start = entries[-1].end - (entries[- 1].end - start) // 2
                    entries[-1] = TimeAlignedType(text=entries[-1].text, start=entries[-1].start, end=start)

                if end <= start:
                    end = start + 1 / sample_rate

            entries.append(TimeAlignedType(text=text, start=start, end=end))

    return entries


def _years_between(record, born):
    try:
        rec_fields = [int(x) for x in record.split('/')]
        brn_fields = [int(x) for x in born.split('/')]
        return rec_fields[2] - brn_fields[2] - ((rec_fields[1], rec_fields[0]) < (brn_fields[1], brn_fields[0]))
    except ValueError:
        return '??'


def _decode_dialect(d: str) -> str:
    if d in ['DR1', '1']:
        return 'New England'
    if d in ['DR2', '2']:
        return 'Northern'
    if d in ['DR3', '3']:
        return 'North Midland'
    if d in ['DR4', '4']:
        return 'South Midland'
    if d in ['DR5', '5']:
        return 'Southern'
    if d in ['DR6', '6']:
        return 'New York City'
    if d in ['DR7', '7']:
        return 'Western'
    if d in ['DR8', '8']:
        return 'Army Brat'

    raise ValueError(f'Unrecognized dialect: {d}')


def load_speakers(input_dir: Path) -> dict:
    speakers = {}
    with open(input_dir / 'SPKRINFO.TXT') as file:
        for line in file:
            if not line.startswith(';'):
                fields = line.strip().split()
                speaker_id = fields[0]
                gender = fields[1]
                dialect = _decode_dialect(fields[2])
                age = _years_between(fields[4], fields[5])
                speakers[speaker_id] = {'gender': gender, 'dialect': dialect, 'age': age}
    return speakers
