import os
from dotenv import load_dotenv
from groq import Groq
import PyPDF2
import gradio as gr
from io import BytesIO
from fpdf import FPDF
import textwrap

# Load API Key
load_dotenv()
GROQ_API_KEY = os.getenv("GROQ_API_KEY")
client = Groq(api_key=GROQ_API_KEY)

def extract_pdf_text(file):
    reader = PyPDF2.PdfReader(file)
    all_pages = [page.extract_text() for page in reader.pages if page.extract_text()]
    return all_pages

def translate_chunk(chunk_text):
    try:
        prompt = f"Translate the following English academic content into formal Urdu:\n\n{chunk_text}"
        response = client.chat.completions.create(
            model="llama3-8b-8192",
            messages=[
                {"role": "system", "content": "You are a professional English-to-Urdu translator."},
                {"role": "user", "content": prompt}
            ]
        )
        return response.choices[0].message.content.strip()
    except Exception as e:
        return f"⚠️ Translation failed: {str(e)}"

def save_to_pdf(translated_chunks):
    pdf = FPDF()
    pdf.add_page()
    pdf.set_font("Arial", size=12)
    for translated in translated_chunks:
        wrapped_lines = textwrap.wrap(translated, width=100)
        for line in wrapped_lines:
            pdf.cell(200, 10, txt=line, ln=True, align='R')
        pdf.ln(5)
    output_path = "/tmp/translated_book.pdf"
    pdf.output(output_path)
    return output_path

def handle_translation(pdf_file):
    pages = extract_pdf_text(pdf_file)
    translated_chunks = []

    for i, page_text in enumerate(pages):
        if not page_text.strip():
            continue
        print(f"📄 Translating page {i+1}/{len(pages)}...")
        urdu = translate_chunk(page_text[:3000])
        translated_chunks.append(f"صفحہ {i+1}\n" + urdu)

    output_pdf = save_to_pdf(translated_chunks)
    return "✅ ترجمہ مکمل ہو گیا! نیچے سے فائل ڈاؤن لوڈ کریں۔", output_pdf

iface = gr.Interface(
    fn=handle_translation,
    inputs=gr.File(label="📘 انگریزی PDF فائل اپلوڈ کریں"),
    outputs=[
        gr.Textbox(label="🔁 سٹیٹس"),
        gr.File(label="📥 ترجمہ شدہ اردو PDF فائل")
    ],
    title="📚 PDF چپٹر بائی چپٹر اردو ترجمہ",
    description="یہ ایپ Groq LLaMA3 کا استعمال کرتے ہوئے انگریزی PDF کو چپٹر بائی چپٹر اردو میں ترجمہ کرتی ہے۔"
)

iface.launch()