import gradio as gr
from llama_cpp import Llama
from paddleocr import PaddleOCR
from PIL import Image

# Load GGUF model
llm = Llama(
    model_path="./deepseek-v3-0324.Q4_K_M.gguf",  # Make sure this file is in your repo
    n_ctx=2048,
    n_threads=8,
    n_gpu_layers=20  # Set to 0 if you are on CPU-only
)

# OCR Function
def ocr_inference(img, lang):
    ocr = PaddleOCR(use_angle_cls=True, lang=lang, use_gpu=False)
    result = ocr.ocr(img, cls=True)[0]
    txts = [line[1][0] for line in result]
    return " ".join(txts)

# Step 1: Convert text to base form words
def text_inference(text, language):
    prompt = (
        f"Given the following {language} text, convert each word into its base form. "
        f"Remove all duplicates. Return the base form words as a comma-separated list.\n\n"
        f"Text:\n{text}"
    )
    response = llm(prompt, max_tokens=256, stop=["</s>"])
    output_text = response["choices"][0]["text"].strip()
    words = [w.strip() for w in output_text.split(",") if w.strip()]
    return words

# Step 2: Generate flashcards for those words
def make_flashcards(words, language):
    prompt = (
        f"For each {language} word in the list, write a flashcard in this format:\n"
        f"word - definition - example sentence - translated sentence.\n\n"
        f"Words:\n{', '.join(words)}"
    )
    response = llm(prompt, max_tokens=512, stop=["</s>"])
    return response["choices"][0]["text"].strip()

# Wrapper logic to handle OCR or text
def flashcard_pipeline(text, image, language):
    if image:
        text = ocr_inference(image, language)
    if not text:
        return "", "Please provide either text or an image."
    words = text_inference(text, language)
    flashcards = make_flashcards(words, language)
    return "\n".join(words), flashcards

# Gradio UI
demo = gr.Interface(
    fn=flashcard_pipeline,
    inputs=[
        gr.Textbox(label="Input Text (leave empty to use image)", lines=4, placeholder="Type or paste sentence here..."),
        gr.Image(label="Upload Image for OCR (optional)", type="filepath"),
        gr.Dropdown(choices=["korean", "japan", "french", "ch"], label="Language (for OCR and LLM)")
    ],
    outputs=[
        gr.Textbox(label="Base Form Words"),
        gr.Textbox(label="Flashcards"),
    ],
    title="Language Flashcard Generator (with OCR + DeepSeek GGUF)",
    description="Either input text or upload an image. The app will extract words, lemmatize them, and generate flashcards."
)

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