import os
from fastapi import FastAPI, Request
from fastapi.responses import JSONResponse
from pydantic import BaseModel
from huggingface_hub import InferenceClient
import gradio as gr
import re

# Initialize FastAPI app
app = FastAPI()

# Initialize Hugging Face Inference Client
clientHFInference = InferenceClient()

# Pydantic model for API input
class InfographicRequest(BaseModel):
    description: str

# Load prompt template from environment variable
SYSTEM_INSTRUCT = os.getenv("SYSTEM_INSTRUCTOR", "Generate a high-quality infographic based on the description provided.")
PROMPT_TEMPLATE = os.getenv("PROMPT_TEMPLATE", "Create an infographic with the following details: {description}")

async def extract_code_blocks(markdown_text):
    """
    Extracts code blocks from the given Markdown text.
    """
    code_block_pattern = re.compile(r'```.*?\n(.*?)```', re.DOTALL)
    code_blocks = code_block_pattern.findall(markdown_text)
    return code_blocks

@app.post("/generate")
async def generate_infographic(request: InfographicRequest):
    description = request.description
    prompt = PROMPT_TEMPLATE.format(description=description)
    
    response = clientHFInference.text_to_image(
        model="stabilityai/stable-diffusion-xl",  # Using an advanced image-based AI model
        inputs=prompt
    )
    
    generated_image_url = response.get("image_url", None)
    
    if generated_image_url:
        return JSONResponse(content={"image_url": generated_image_url})
    else:
        return JSONResponse(content={"error": "No infographic generated"}, status_code=500)

# Gradio UI for Hugging Face Spaces
def generate_infographic_ui(description):
    response = generate_infographic(InfographicRequest(description=description))
    return response["image_url"] if "image_url" in response else "Error generating infographic"

demo = gr.Interface(
    fn=generate_infographic_ui,
    inputs="text",
    outputs="image",
    title="AI Infographic Generator",
    description="Enter a description, and the AI will generate a high-quality infographic."
)

demo.launch(share=True)