import pyktok as pyk
pyk.specify_browser("chrome")  # Specify the browser to use for accessing TikTok
import os
import moviepy.editor as mp
from openai import OpenAI

client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))
import time
from pathlib import Path
from openai import OpenAI
import os

import argparse

client = OpenAI(
    api_key=os.environ.get("OPENAI_API_KEY"),
)

# def download_video(url):
#     pyk.save_tiktok(url, save_video=True, filename="tiktok_video.mp4")


def download_video(url, outpath=Path('.')):
    pyk.save_tiktok(url, True, str(outpath / 'tiktok_video.mp4'))
    print('Downloaded video wait...', end='')
    time.sleep(5)
    print('Done')


def extract_audio(path):
    outpath = Path(path).parent / "tiktok_audio.mp3"
    print('Extracting audio...', flush=True)
    video = mp.VideoFileClip(path)
    video.audio.write_audiofile(str(outpath))
    print('Audio extracted', flush=True)
    return outpath
# def transcribe_audio():
#     model = whisper.load_model("base")
#     result = model.transcribe("tiktok_audio.mp3")
#     return result["text"]

def transcribe_audio(audio_file_path):
    print('sending for transcription...', flush=True)
    with open(audio_file_path, 'rb') as audio_file:
        transcription = client.audio.transcriptions.create(model="whisper-1", file=audio_file, 
                                                           response_format="text")
    print('Transcription done', flush=True)
    return transcription


def format_recipe(transcript):
    prompt = f"Please format the following recipe transcript into a nicely formatted recipe:\n\n{transcript}"
    response = client.completions.create(engine="text-davinci-003",
    prompt=prompt,
    max_tokens=500,
    n=1,
    stop=None,
    temperature=0.7)
    return response.choices[0].text.strip()

def main():
    # tiktok_url = input("Enter the TikTok video URL: ")
    # tiktok_url = "https://www.tiktok.com/t/ZTLjYBSpt/"
    tiktok_url = "https://www.tiktok.com/@emmaaaaaaam_/video/7348493781961886981"
    _, user, _, video_id = tiktok_url.strip('https://').strip('http://').strip('/').split("/")
    
    outpath = Path(video_id)
    outpath.mkdir(parents=True, exist_ok=True)
    
    download_video(tiktok_url, outpath)
    video_path = [p for p in outpath.parent.rglob('*.mp4') if video_id in p.name][0]
    video_path.rename(outpath / video_path.name)
    video_path = str(outpath / video_path.name)
    
    audio_path = extract_audio(video_path)
    transcript_path = outpath / 'transcript.txt'
    if transcript_path.exists():
        with transcript_path.open() as f:
            transcript = f.read()
    else:
        transcript = transcribe_audio(audio_path)
        with transcript_path.open('w') as f:
            f.write(transcript)

    formatted_recipe = format_recipe(transcript)
    
    print("Formatted Recipe:")
    print(formatted_recipe)
    
    # # Clean up temporary files
    # os.remove("tiktok_video.mp4")
    # os.remove("tiktok_audio.mp3")

if __name__ == "__main__":
    main()