import os
import urllib.request
from time import sleep
from typing import Dict, List, Generator

import gradio as gr
from openai import OpenAI
from dotenv import load_dotenv

load_dotenv()

class MockInterviewer:

    def __init__(self) -> None:
        self._client = OpenAI(api_key=os.environ['OPENAI_API_KEY'])
        self._assistant_id_cache: Dict[str, str] = {}
        self.clear_thread()

    def interface_chat(self, message: Dict, history: List[List], job_role: str, company: str) -> Generator:
        yield self._chat(message.get('text'), job_role, company)

    def clear_thread(self) -> None:
        print('Initializing new thread')
        self._thread = self._client.beta.threads.create()
    
    def transcript(self, audio: str, job_role: str, company: str) -> str:
        with open(audio, 'rb') as audio_file:
            transcriptions = self._client.audio.transcriptions.create(
                model='whisper-1',
                file=audio_file,
            )
        os.remove(audio)
        response = self._chat(transcriptions.text, job_role, company)
        return [(transcriptions.text, response)]
    
    def _chat(self, message: str, job_role: str, company: str) -> str:
        print('Started chat')
        self._validate_fields(job_role, company)
        assistant_id = self._init_assistant(job_role, company)
        return self._send_message(message, assistant_id)

    def _validate_fields(self, job_role: str, company: str) -> None:
        if not job_role and not company:
            raise gr.Error('Job Role and Company are required fields.')
        if not job_role:
            raise gr.Error('Job Role is a required field.')
        if not company:
            raise gr.Error('Company is a required field.')

    def _send_message(self, message: str, assistant_id: str) -> str:
        self._client.beta.threads.messages.create(thread_id=self._thread.id, role='user', content=message)
        print('Message created')
        run = self._client.beta.threads.runs.create(thread_id=self._thread.id, assistant_id=assistant_id)
        print('Run created')

        # Check if the Run requires action (function call)
        while True:
            run_status = self._client.beta.threads.runs.retrieve(thread_id=self._thread.id, run_id=run.id)
            print(f'Run status: {run_status.status}')
            if run_status.status == 'completed':
                break

            sleep(1)  # Wait for a second before checking again

        # Retrieve and return the latest message from the assistant
        messages = self._client.beta.threads.messages.list(thread_id=self._thread.id)
        response = messages.data[0].content[0].text.value
        print(f'Assistant response: {response}')
        return response

    def _create_files(self, company: str) -> List[str]:
        if company.lower() == 'amazon':
            url = 'https://www.aboutamazon.com/about-us/leadership-principles'
            filename = 'leadership_principles.html'
        else:
            return []
        
        filename, headers = urllib.request.urlretrieve(url, filename)
        with open(filename, 'rb') as file:
            assistant_file = self._client.files.create(file=file, purpose='assistants')
            file_ids = [assistant_file.id]
        os.remove(filename)
        return file_ids

    def _init_assistant(self, job_role: str, company: str) -> str:
        cache_key = self._create_cache_key(job_role, company)
        if cache_key in self._assistant_id_cache:
            print(f'Fetched from cache for key {cache_key}')
            return self._assistant_id_cache.get(cache_key)
        else:
            print(f'Initializing new assistant for key {cache_key}')
            file_ids = self._create_files(company)

            assistant = self._client.beta.assistants.create(
                name='Mock Interviewer',
                instructions=f'You are an AI mock interviewer for {job_role} roles at {company}. Please make it obvious this is your purpose. If you have been provided a file, use it as an interview guide.',
                model='gpt-4-0125-preview',
                tools=[
                    {
                        'type': 'retrieval'  # This adds the knowledge base as a tool
                    }
                ],
                file_ids=file_ids)
            
            self._assistant_id_cache[cache_key] = assistant.id
            return assistant.id

    def _create_cache_key(self, job_role: str, company: str) -> str:
        return f'{job_role.lower()}+{company.lower()}'

# Creating the Gradio interface
with gr.Blocks() as demo:
    mock_interviewer = MockInterviewer()

    with gr.Row():
        job_role = gr.Textbox(label='Job Role', placeholder='Product Manager')
        company = gr.Textbox(label='Company', placeholder='Amazon')

    chat_interface = gr.ChatInterface(
        fn=mock_interviewer.interface_chat,
        additional_inputs=[job_role, company],
        title='I am your AI mock interviewer',
        description='Make your selections above to configure me.',
        multimodal=True,
        retry_btn=None,
        undo_btn=None).queue()

    chat_interface.load(mock_interviewer.clear_thread)
    chat_interface.clear_btn.click(mock_interviewer.clear_thread)

    audio = gr.Audio(sources=['microphone'], type='filepath', editable=False)
    audio.stop_recording(fn=mock_interviewer.transcript,
                         inputs=[audio, job_role, company],
                         outputs=[chat_interface.chatbot],
                         api_name=False)

if __name__ == '__main__':
    demo.launch().queue()