import os
import streamlit as st
from together import Together
from PIL import Image
import moviepy.editor as mp
from gtts import gTTS
import tempfile
import base64
from io import BytesIO
import numpy as np

# Set up Together.ai API key
together_api_key = os.environ.get("TOGETHER_API_KEY",)
client = Together(api_key=together_api_key)

def generate_interview_transcript(role, experience, additional_details, interview_type):
    prompt = f"""As a professional interview transcript writer, create a {interview_type} mock interview script for a {experience} {role} candidate. 
    Include interviewer and candidate responses using gender-neutral names. 
    The interview should be realistic, professional, and concise.
    Incorporate the following additional details about the candidate: {additional_details}
    
    Format the transcript as follows:
    Interviewer (Alex): [Question]
    Candidate (Sam): [Response]
    
    Create a short but relevant interview with 3-5 question-answer pairs."""

    response = client.chat.completions.create(
        model="meta-llama/Llama-3-8b-chat-hf",
        messages=[
            {"role": "system", "content": "You are a professional interview transcript writer."},
            {"role": "user", "content": prompt}
        ],
    )
    return response.choices[0].message.content

def generate_audio(script, lang='en'):
    lines = script.split('\n')
    full_text = " ".join([line.split(':', 1)[1].strip() for line in lines if ':' in line])

    with tempfile.NamedTemporaryFile(suffix=".mp3", delete=False) as temp:
        audio_file = temp.name
        tts = gTTS(text=full_text, lang=lang, slow=False)
        tts.save(audio_file)
    return audio_file

def generate_image(prompt):
    response = client.images.generate(
        prompt=prompt,
        model="stabilityai/stable-diffusion-2-1",
        steps=30,
        width=512,
        height=512,
        n=1
    )
    image_data = base64.b64decode(response.data[0].b64_json)
    return Image.open(BytesIO(image_data))

def create_video_with_images_and_audio(script, audio_file, interviewer_img, interviewee_img, output_video):
    audio_clip = mp.AudioFileClip(audio_file)
    total_audio_duration = audio_clip.duration
    lines = [line for line in script.split('\n') if ':' in line]

    clips = []
    start_time = 0

    for line in lines:
        speaker, text = line.split(':', 1)
        text = text.strip()
        duration = len(text.split()) * 0.5  # Estimate duration based on word count

        if 'Interviewer' in speaker:
            img = np.array(interviewer_img)
        else:
            img = np.array(interviewee_img)

        clip = mp.ImageClip(img).set_duration(duration)
        clip = clip.set_start(start_time)
        clips.append(clip)

        start_time += duration

    video = mp.CompositeVideoClip(clips, size=(512, 512))
    final_clip = video.set_audio(audio_clip)
    final_clip.write_videofile(output_video, codec='libx264', fps=24)

def create_mock_interview(role, experience, additional_details, interview_type):
    with st.spinner("Generating interview transcript..."):
        transcript = generate_interview_transcript(role, experience, additional_details, interview_type)
    st.success("Interview transcript generated!")
    st.text_area("Generated Transcript", transcript, height=200)

    with st.spinner("Generating audio..."):
        audio_file = generate_audio(transcript)
    st.success("Audio generated!")

    with st.spinner("Generating images for interviewer and interviewee..."):
        interviewer_img = generate_image("A professional interviewer sitting in a modern office setting, facing the camera, upper body shot")
        interviewee_img = generate_image(f"A confident{experience} job candidate sitting for an interview, facing the camera, upper body shot")
    st.success("Images generated!")

    with st.spinner("Creating final video..."):
        with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as temp:
            output_video = temp.name
            create_video_with_images_and_audio(transcript, audio_file, interviewer_img, interviewee_img, output_video)
    st.success("Video created successfully!")
    return output_video

# Streamlit UI
st.set_page_config(page_title="AIInterviewCoach", page_icon="🎥", layout="wide")

st.title("🎬 AIInterviewCoach")
st.subheader("Generate Custom Mock Interview Videos with AI")

col1, col2 = st.columns(2)
with col1:
    role = st.text_input("🎨 Job Role", "Data Analyst")
    experience = st.selectbox("🌟 Experience Level", ["Entry-level", "Mid-level", "Senior", "Executive"])
with col2:
    interview_type = st.selectbox("🎭 Interview Scenario", ["Standard", "Behavioral", "Technical", "Case Study"])
    additional_details = st.text_area("✨ Additional Details", "Proficient in SQL and Python")

if st.button("🎥 Generate Mock Interview"):
    st.write("Generating interview based on the following parameters:")
    st.write(f"- Role: {role}")
    st.write(f"- Experience Level: {experience}")
    st.write(f"- Interview Type: {interview_type}")
    st.write(f"- Additional Details: {additional_details}")
    
    output_video = create_mock_interview(role, experience, additional_details, interview_type)
    st.success("🌟 Your mock interview is ready!")
    st.video(output_video)
    with open(output_video, 'rb') as f:
        st.download_button(
            label="📥 Download Video",
            data=f,
            file_name="my_mock_interview.mp4",
            mime="video/mp4"
        )
    os.unlink(output_video)  # Delete the temporary file

st.sidebar.title("🚀 About AIInterviewCoach")
st.sidebar.info("""
AIInterviewCoach creates realistic mock interview videos.

🔧 Features:
- Custom scenarios
- Role-specific questions
- AI-generated visuals
- Professional voice-overs

🌈 Created by: Dhruv Tibarewal
""")

st.sidebar.title("💡 Tips")
st.sidebar.success("""
- Be specific about your role
- Share unique skills
- Try different scenarios
- Practice, iterate, succeed! 🌟
""")