from transformers import pipeline, AutoTokenizer, VisionEncoderDecoderModel, AutoProcessor
import torch
from PIL import Image
from io import BytesIO
import base64

# Chuyển ảnh thành base64 (tùy chọn nếu bạn cần hiển thị hoặc xuất)
def pil_to_base64(image: Image.Image, format="PNG") -> str:
    buffered = BytesIO()
    image.save(buffered, format=format)
    return base64.b64encode(buffered.getvalue()).decode("utf-8")

def parse_to_json(result_text):
    """
    Nếu output là các dòng 'key: value', parse thành dict.
    Nếu không, gói nguyên text vào trường 'text'.
    """
    data = {}
    lines = [line.strip() for line in result_text.splitlines() if line.strip()]
    for line in lines:
        if ":" in line:
            key, val = line.split(":", 1)
            data[key.strip()] = val.strip()
        else:
            # Nếu không tách được, gom vào list chung
            data.setdefault("text", []).append(line)
    # Nếu chỉ có list 'text', chuyển về chuỗi
    if set(data.keys()) == {"text"}:
        data = {"text": "\n".join(data["text"])}
    return data

# class TrOCRModel:
#     def __init__(self, model_id="microsoft/trocr-base-printed", cache_dir=None, device=None):
#         self.model_id = model_id
#         self.cache_dir = cache_dir
#         self.device = device

#         self.processor = TrOCRProcessor.from_pretrained(self.model_id, cache_dir=self.cache_dir)
#         self.model = VisionEncoderDecoderModel.from_pretrained(self.model_id, cache_dir=self.cache_dir)
#         self.model.to(self.device)

#     def predict(self, image: Image.Image) -> str:
#         if image is None:
#             raise ValueError("No image provided")

#         image = image.convert("RGB")
#         pixel_values = self.processor(images=image, return_tensors="pt").pixel_values.to(self.device)

#         with torch.no_grad():
#             generated_ids = self.model.generate(pixel_values)
#             generated_text = self.processor.batch_decode(generated_ids, skip_special_tokens=True)[0]

#         return generated_text

class TrOCRModel:
    def __init__(self, model_id="microsoft/trocr-base-printed", cache_dir=None, device=None):
        self.pipe = pipeline("image-to-text", model=model_id, device=device)

    def predict(self, image: Image.Image) -> str:
        if image is None:
            raise ValueError("No image provided")

        image = image.convert("RGB")
        result = self.pipe(image)
        return result[0]['generated_text'] if result else ""
    
class EraXModel:
    def __init__(self, model_id="erax-ai/EraX-VL-2B-V1.5", cache_dir=None, device=None):
        self.pipe = pipeline("image-to-text", model=model_id, device=device)

    def predict(self, image: Image.Image) -> str:
        if image is None:
            raise ValueError("No image provided")
        
        decoded_image_text = pil_to_base64(image)
        base64_data = f"data:image;base64,{decoded_image_text}"
        messages = [
            {
                "role": "user",
                "content": [
                    {
                        "type": "image",
                        "image": base64_data,
                    },
                    {
                        "type": "text",
                        "text": "Trích xuất thông tin nội dung từ hình ảnh được cung cấp."
                    },
                ],
            }
        ]

        result = self.pipe(image)[0]['generated_texts']
        return parse_to_json(result) if result else {}