import gradio as gr
from huggingface_hub import InferenceClient
import PyPDF2
import docx
import io

client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")

def extract_text_from_pdf(pdf_file):
    pdf_reader = PyPDF2.PdfReader(io.BytesIO(pdf_file))
    text = ""
    for page in pdf_reader.pages:
        text += page.extract_text() + "\n"
    return text

def extract_text_from_docx(docx_file):
    doc = docx.Document(io.BytesIO(docx_file))
    return "\n".join([para.text for para in doc.paragraphs])

def parse_cv(file):
    if file is None:
        return "Please upload a CV file."
    
    file_ext = file.name.split(".")[-1].lower()
    file_bytes = file.read()
    
    if file_ext == "pdf":
        text = extract_text_from_pdf(file_bytes)
    elif file_ext == "docx":
        text = extract_text_from_docx(file_bytes)
    else:
        return "Unsupported file format. Please upload a PDF or DOCX file."
    
    prompt = f"Analyze the following CV and generate a professional summary and improvement suggestions:\n\n{text}"
    response = client.text_generation(prompt, max_tokens=512)
    return response

def respond(
    message,
    history: list[tuple[str, str]],
    system_message,
    max_tokens,
    temperature,
    top_p,
):
    messages = [{"role": "system", "content": system_message}]

    for val in history:
        if val[0]:
            messages.append({"role": "user", "content": val[0]})
        if val[1]:
            messages.append({"role": "assistant", "content": val[1]})

    messages.append({"role": "user", "content": message})
    response = ""

    for message in client.chat_completion(
        messages,
        max_tokens=max_tokens,
        stream=True,
        temperature=temperature,
        top_p=top_p,
    ):
        token = message.choices[0].delta.content
        response += token
        yield response

demo = gr.Blocks()

with demo:
    gr.Markdown("## AI-powered CV Analyzer and Chatbot")
    with gr.Tab("Chatbot"):
        chat_interface = gr.ChatInterface(
            respond,
            additional_inputs=[
                gr.Textbox(value="You are a friendly Chatbot.", label="System message"),
                gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
                gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
                gr.Slider(
                    minimum=0.1,
                    maximum=1.0,
                    value=0.95,
                    step=0.05,
                    label="Top-p (nucleus sampling)",
                ),
            ],
        )
    
    with gr.Tab("CV Analyzer"):
        gr.Markdown("### Upload your CV (PDF or DOCX) to receive a professional analysis.")
        file_input = gr.File(label="Upload CV", type="file")
        output_text = gr.Textbox(label="CV Analysis Report", lines=10)
        analyze_button = gr.Button("Analyze CV")
        
        analyze_button.click(parse_cv, inputs=file_input, outputs=output_text)

if __name__ == "__main__":
    demo.launch()