import os
import torch
from PIL import Image
import gradio as gr
from datasets import load_dataset
from unsloth import FastLanguageModel
from transformers import AutoProcessor, TrainingArguments, Trainer
import subprocess
import sys

# Installe unsloth dynamiquement depuis GitHub s’il n’est pas déjà là
try:
    import unsloth
except ImportError:
    subprocess.check_call([sys.executable, "-m", "pip", "install", "git+https://github.com/unslothai/unsloth.git"])
    import unsloth


# === CONFIG ===
CKPT = "unsloth/Llama-3.2-11B-Vision-Instruct"
MODEL_SAVE_PATH = "./llama3-handwriting-ocr"
DATASET_NAME = "om440/partial-rimes-handwritten-dataset"
MAX_TOKENS = 256


# === Load or fine-tune the model with LoRA ===
def train_with_lora():
    print("🚀 Loading dataset...")
    dataset = load_dataset(DATASET_NAME)

    processor = AutoProcessor.from_pretrained(CKPT)

    def preprocess(example):
        image = example["image"].convert("RGB")
        prompt = (
            "Output ONLY the raw text as it appears in the image, nothing else.\n"
            "You have an image containing both handwritten and printed text...\n"
            "Transcribe EXACTLY all visible text..."
        )
        messages = [
            {
                "role": "user",
                "content": [
                    {"type": "text", "text": prompt},
                    {"type": "image"}
                ]
            }
        ]
        chat_input = processor.apply_chat_template(messages, add_generation_prompt=True)
        model_inputs = processor(
            text=chat_input,
            images=[image],
            return_tensors="pt"
        )
        labels = processor.tokenizer(
            example["text"],
            return_tensors="pt",
            padding="max_length",
            truncation=True,
            max_length=MAX_TOKENS
        ).input_ids[0]
        model_inputs["labels"] = labels
        return model_inputs

    print("🔄 Preprocessing dataset...")
    tokenized_dataset = dataset["train"].map(lambda x: preprocess(x, processor), remove_columns=dataset["train"].column_names)

    print("📦 Loading base model with LoRA...")
    model, _ = FastLanguageModel.from_pretrained(
        model_name=CKPT,
        max_seq_length=MAX_TOKENS,
        dtype=None,  # auto-select bf16/fp16
        load_in_4bit=True  # Memory-efficient
    )

    model = FastLanguageModel.get_peft_model(
        model,
        r=8,
        lora_alpha=16,
        lora_dropout=0.05,
        bias="none",
        use_gradient_checkpointing=True,
        random_state=42,
        use_rslora=False,
        loftq_config=None
    )

    print("🏋️ Starting fine-tuning with LoRA...")
    training_args = TrainingArguments(
        output_dir=MODEL_SAVE_PATH,
        per_device_train_batch_size=1,
        gradient_accumulation_steps=2,
        num_train_epochs=2,
        save_steps=100,
        logging_steps=10,
        learning_rate=2e-4,
        fp16=True,
        report_to="none"
    )

    trainer = Trainer(
        model=model,
        tokenizer=processor.tokenizer,
        args=training_args,
        train_dataset=tokenized_dataset
    )

    trainer.train()

    print("💾 Saving fine-tuned LoRA adapter...")
    model.save_pretrained(MODEL_SAVE_PATH)
    processor.save_pretrained(MODEL_SAVE_PATH)


# === Inference ===
def load_model():
    from peft import PeftModel
    processor = AutoProcessor.from_pretrained(MODEL_SAVE_PATH)
    model, _ = FastLanguageModel.from_pretrained(
        model_name=CKPT,
        max_seq_length=MAX_TOKENS,
        dtype=None,
        load_in_4bit=True
    )
    model = PeftModel.from_pretrained(model, MODEL_SAVE_PATH)
    model.eval()
    return model, processor


def extract_text(image_path):
    model, processor = load_model()
    image = Image.open(image_path).convert("RGB")

    prompt = (
        "Output ONLY the raw text as it appears in the image, nothing else."
        "You have an image containing both handwritten and printed text in French and/or English, and also punctuation and underscores.\n"
        "Your task: transcribe EXACTLY all visible text, preserving all characters, accents, punctuation, spacing, and line breaks.\n"
        "Include tables and forms clearly if present.\n"
        "Do NOT add any explanations, comments, summaries, or extra text.\n"
        "Check the output first to not duplicate results."
        "Preserve the original reading order, including line breaks and the natural layout of tables or forms. Output the text exactly as it appears visually, maintaining the structure."
        "Don't indicate blank space."
        "Don't separate handwritten and printed text."
        "DO NOT confuse between '.' a point and '|' a border."
        "Extract only the raw text and do not add any comment."
        "Extract only the data available."
    )

    messages = [
        {
            "role": "user",
            "content": [
                {"type": "text", "text": prompt},
                {"type": "image"}
            ]
        }
    ]

    chat_input = processor.apply_chat_template(messages, add_generation_prompt=True)
    inputs = processor(
        text=chat_input,
        images=[image],
        return_tensors="pt"
    ).to("cuda")

    with torch.no_grad():
        outputs = model.generate(**inputs, max_new_tokens=300)

    result = processor.decode(outputs[0], skip_special_tokens=True)

    if "assistant" in result.lower():
        result = result[result.lower().find("assistant") + len("assistant"):].strip()

    return result


# === Gradio UI ===
def launch_gradio():
    print("🌐 Launching Gradio interface...")
    demo = gr.Interface(
        fn=extract_text,
        inputs=gr.Image(type="filepath", label="Upload Image"),
        outputs=gr.Textbox(label="Extracted Text"),
        title="Fine-tuned Handwritten Text Extractor",
        description="Upload an image with printed or handwritten text and extract it using a fine-tuned LLaMA-3 Vision model."
    )
    demo.launch(share=True)


# === Run ===
if __name__ == "__main__":
    if not os.path.exists(MODEL_SAVE_PATH):
        train_with_lora()
    launch_gradio()