import gradio as gr
import subprocess
import json
import os
import stat
import requests
import datetime
import zipfile
import matplotlib.font_manager
from huggingface_hub import HfApi, HfFolder

# Hugging Face repo and token
HF_REPO = "ArrcttacsrjksX/Texttoimage"
HF_ENGINE_URL = "https://huggingface.co/ArrcttacsrjksX/Texttoimage/resolve/main/engine"
HF_TOKEN = os.getenv("HF_TOKEN")  # Lấy token từ biến môi trường
ENGINE_EXECUTABLE = "./engine"

def download_engine():
    """Download engine from Hugging Face if not available."""
    if not os.path.exists(ENGINE_EXECUTABLE):
        headers = {"Authorization": f"Bearer {HF_TOKEN}"}
        response = requests.get(HF_ENGINE_URL, headers=headers, stream=True)
        if response.status_code == 200:
            with open(ENGINE_EXECUTABLE, "wb") as f:
                for chunk in response.iter_content(chunk_size=8192):
                    f.write(chunk)
            os.chmod(ENGINE_EXECUTABLE, os.stat(ENGINE_EXECUTABLE).st_mode | stat.S_IXUSR)
        else:
            raise Exception("Failed to download engine")

def ensure_executable(file_path):
    if not os.access(file_path, os.X_OK):
        os.chmod(file_path, os.stat(file_path).st_mode | stat.S_IXUSR)

def extract_and_load_fonts(directory="fontfile", extract_to="extracted_fonts"):
    if not os.path.exists(extract_to):
        os.makedirs(extract_to)
    fonts = []
    for root, _, files in os.walk(directory):
        for file in files:
            if file.endswith(".zip"):
                zip_path = os.path.join(root, file)
                try:
                    with zipfile.ZipFile(zip_path, 'r') as zip_ref:
                        zip_ref.extractall(extract_to)
                except Exception as e:
                    print(f"Failed to extract {zip_path}: {e}")
    for root, _, files in os.walk(extract_to):
        for file in files:
            if file.endswith(".ttf") or file.endswith(".otf") or file.endswith(".shx"):
                fonts.append(os.path.join(root, file))
    return fonts

def get_system_fonts():
    return matplotlib.font_manager.findSystemFonts(fontpaths=None, fontext='ttf')

def get_available_fonts():
    system_fonts = get_system_fonts()
    extracted_fonts = extract_and_load_fonts()
    return sorted(set(system_fonts + extracted_fonts))

def upload_to_huggingface(file_path, text_content, timestamp_folder):
    """Upload image and text to Hugging Face repo."""
    api = HfApi()
    HfFolder.save_token(HF_TOKEN)
    repo_path = f"{HF_REPO}/{timestamp_folder}"
    api.upload_file(path_or_fileobj=file_path, path_in_repo=f"{repo_path}/image.png", repo_id=HF_REPO)
    with open("temp_text.txt", "w") as f:
        f.write(text_content)
    api.upload_file(path_or_fileobj="temp_text.txt", path_in_repo=f"{repo_path}/text.txt", repo_id=HF_REPO)
    os.remove("temp_text.txt")

def call_engine(file_input, input_text, font_size, width, height, bg_color, text_color, mode, font_name, align, line_spacing, image_format):
    download_engine()
    ensure_executable(ENGINE_EXECUTABLE)
    if file_input:
        input_text = read_file_content(file_input)
    input_data = {
        "input_text": input_text,
        "font_size": font_size,
        "width": width,
        "height": height,
        "bg_color": bg_color,
        "text_color": text_color,
        "mode": mode,
        "font_path": font_name,
        "align": align,
        "line_spacing": line_spacing,
        "image_format": image_format
    }
    result = subprocess.run([ENGINE_EXECUTABLE, json.dumps(input_data)], capture_output=True, text=True)
    if result.returncode != 0:
        raise Exception(f"Engine error: {result.stderr}")
    output_path = result.stdout.strip()
    timestamp_folder = datetime.datetime.now().strftime("%d-%m-%Y %H-%M-%S")
    upload_to_huggingface(output_path, input_text, timestamp_folder)
    return output_path

def read_file_content(file):
    """Read text from uploaded file."""
    try:
        return file.decode("utf-8") if isinstance(file, bytes) else file.read().decode("utf-8")
    except Exception as e:
        return f"Error reading file: {e}"

with gr.Blocks() as demo:
    gr.Markdown("# 🖼️ Text to Image Converter")
    available_fonts = get_available_fonts()
    default_font = available_fonts[0] if available_fonts else ""
    
    with gr.Row():
        input_text = gr.Textbox(label="Enter Text", placeholder="Type or paste text here...", lines=5)
        file_input = gr.File(label="Upload a Text File", type="binary")
    
    with gr.Row():
        font_size = gr.Slider(0, 510, value=30, label="Font Size")
        font_name = gr.Dropdown(choices=available_fonts, value=default_font, label="Font")
        align = gr.Radio(["Left", "Center", "Right"], label="Text Alignment", value="Center")
        width = gr.Slider(0, 9000, value=800, label="Image Width")
        height = gr.Slider(0, 9000, value=600, label="Image Height")
    
    with gr.Row():
        bg_color = gr.ColorPicker(label="Background Color", value="#FFFFFF")
        text_color = gr.ColorPicker(label="Text Color", value="#000000")
    
    with gr.Row():
        mode = gr.Radio(["Plain Text", "LaTeX Math"], label="Rendering Mode", value="Plain Text")
        image_format = gr.Radio(["PNG", "JPEG"], label="Image Format", value="PNG")
    
    line_spacing = gr.Slider(1.0, 1000.0, value=1.2, step=0.1, label="Line Spacing")
    output_image = gr.Image(label="Generated Image")
    
    convert_button = gr.Button("Convert Text to Image")
    
    convert_button.click(
        call_engine,
        inputs=[file_input, input_text, font_size, width, height, bg_color, text_color, mode, font_name, align, line_spacing, image_format],
        outputs=output_image
    )

demo.launch()