import re
import requests
import sys
import os

# Header required for downloading audio files
HEADERS = {
    "X-Access": "please don't steal these files"
}

# Function to fetch HTML from jpdb.io
def fetch_html(query):
    url = f"https://jpdb.io/search?q={query}"
    response = requests.get(url)
    response.raise_for_status()
    response.encoding = response.apparent_encoding  # Ensure correct encoding
    return response.text

# Function to extract spellings and audio IDs
def extract_spellings(html):
    pattern = re.compile(
        r'<ruby class="v(?: nf)?">(.*?)</ruby>.*?data-audio="m1/([a-f0-9]+)"',
        re.DOTALL
    )
    matches = pattern.findall(html)

    spellings = []
    for ruby_content, audio_id in matches:
        # Extract kanji and furigana parts separately
        kanji_parts = re.findall(r'([^<rt>]+)(?:<rt>(.*?)</rt>)?', ruby_content)

        # Build kanji and reading properly
        kanji_result = ""
        reading_result = ""

        for kanji, furigana in kanji_parts:
            kanji_result += kanji.strip()  # Keep kanji as is
            if furigana:
                reading_result += furigana.strip()  # Add furigana if present
            else:
                reading_result += kanji.strip()  # If no furigana, use kanji (kana case)

        filename = f"{kanji_result}-{reading_result}-{audio_id}.ogg"
        spellings.append((filename, audio_id))

    return spellings

# Function to download and decode audio files
def download_and_decode(audio_id, filename):
    url = f"https://jpdb.io/static/v/m1/{audio_id}"
    response = requests.get(url, headers=HEADERS)

    if response.status_code != 200:
        print(f"Failed to download {filename} ({response.status_code})")
        return

    # Save the raw encoded file first (optional)
    encoded_path = f"encoded_{audio_id}.ogg"
    with open(encoded_path, "wb") as f:
        f.write(response.content)

    # Decode the file
    with open(encoded_path, "rb") as f:
        data = bytearray(f.read())

    # XOR the first 4 bytes
    key = [0x06, 0x23, 0x54, 0x0f]
    for i in range(4):
        data[i] ^= key[i]

    # Save the decoded file
    with open(filename, "wb") as f:
        f.write(data)

    # Clean up the encoded file
    os.remove(encoded_path)
    print(f"Saved: {filename}")

def main():
    if len(sys.argv) != 2:
        print("Usage: python jpdb.py <query>")
        sys.exit(1)

    query = sys.argv[1]
    html = fetch_html(query)
    spellings = extract_spellings(html)

    for filename, audio_id in spellings:
        download_and_decode(audio_id, filename)

if __name__ == "__main__":
    main()