import gradio as gr
import openai
from PIL import Image
import base64
from io import BytesIO
import tempfile
import os

# Set your OpenAI API key
openai.api_key =  os.environ['KEY'] # 🔒 REPLACE with your real key or use environment variable

# Define the prompt (you can make this external if you want)
PROMPT_GEN = """
You are an expert in archival retroconversion working with GGAOF (Gouvernement général de l’Afrique occidentale française) collections. You are processing a scanned metadata page containing one or more archival entries.
Your task is to extract and encode the information into structured XML blocks using the ISAD(G) and “Annexe 1 – Consignes générales de rétroconversion” rules.
🔧 For EACH entry on the image, generate exactly one <c> block under a single <dsc> root.
---
📘 STRUCTURE TO FOLLOW:
<dsc>
  <c altrender="ligeo-simple-standardisadg">
    <did>
      <unitid>FM GGAOF [cote exactly as shown]</unitid>
      <unittitle>[Title text after colon, without punctuation]</unittitle>
      <unitdate>[Date or date range exactly as shown, using one of these formats: "aaaa", "mm aaaa", "jj mm aaaa", or "aaaa-aaaa"]</unitdate>
      [<physdesc>[Optional physical description, if present]</physdesc>]
    </did>
    <scopecontent>
      <p>[Each bullet point or indented line]</p>
      ...
    </scopecontent>
    <dao href="/AFRIQUE/GGAOF_FM/G_1a20/[unitid with no spaces or slashes]_M"/>
  </c>
  ...
</dsc>
---
📌 EXTRACTION RULES:
1. <unitid>
   - Copy the cote exactly as shown (e.g. “14 G/10”).
   - Prepend “FM GGAOF ” to it.
2. <unittitle>
   - Extract the title after the colon.
   - Remove any final punctuation.
   - Write it on a single line.
3. <physdesc> (optional)
   - Only include if there is material information (e.g., "1 chemise", "179 pages").
4. <unitdate>
   - Copy date range as printed.
   - Accepted formats: "aaaa", "mm aaaa", "jj mm aaaa", "aaaa-aaaa".
   - NO slashes allowed.
5. <scopecontent>
   - Each bullet point, dash-indented, or subline becomes one <p>.
   - Retain all original punctuation, accents, and line content.
6. <dao>
   - Construct href as: "/AFRIQUE/GGAOF_FM/G_1a20/" + [unitid with no spaces or slashes] + "_M"
     - Example: unitid = "14 G/8" → dao = "14G8_M"
---
⛔ Do NOT hallucinate. Only use what is visible in the image.  
⛔ Do NOT add explanations or commentary.  
✅ Return ONLY the complete XML block as shown above.
"""

def clean_gpt_xml(output: str) -> str:
    lines = output.strip().splitlines()
    if lines[0].strip().startswith("```") and lines[-1].strip() == "```":
        return "\n".join(lines[1:-1]).strip()
    return output.strip()


import os

def process_image(image_path: str, prompt_text: str):
    print("✅ Image path received for processing:", image_path)
    
    # Extract filename (without extension)
    base_name = os.path.splitext(os.path.basename(image_path))[0]
    output_filename = f"{base_name}.xml"

    # Load and convert image
    image = Image.open(image_path)
    buffered = BytesIO()
    image.save(buffered, format="PNG")
    img_base64 = base64.b64encode(buffered.getvalue()).decode()
    print("📦 Image encoded to base64.")

    # Send to OpenAI
    try:
        response = openai.ChatCompletion.create(
            model="gpt-4o",
            messages=[
                {
                    "role": "user",
                    "content": [
                        # {"type": "text", "text": PROMPT_GEN},
                        {"type": "text", "text": prompt_text},
                        {"type": "image_url", "image_url": {"url": f"data:image/png;base64,{img_base64}"}}
                    ],
                }
            ],
            temperature=0,
        )
        xml_output = response['choices'][0]['message']['content']
        xml_output = clean_gpt_xml(xml_output)
        print("✅ XML received from GPT-4o.")
    except Exception as e:
        print("❌ Error during OpenAI request:", e)
        return "Error: Could not process image.\n\n" + str(e), None

    # Save XML to a specific filename
    output_path = os.path.join(tempfile.gettempdir(), output_filename)
    with open(output_path, "w", encoding="utf-8") as f:
        f.write(xml_output)

    print(f"💾 XML saved as: {output_path}")
    return xml_output, output_path


# === Gradio Interface ===
with gr.Blocks() as demo:
    gr.Markdown("## 🧾 Metadata → XML EAD Application")
    with gr.Row():
        with gr.Column():
            # image_input = gr.Image(type="pil", label="Upload Image")
            image_input = gr.Image(type="filepath", label="Upload Image")
            prompt_input = gr.Textbox(label="📝 Prompt", value=PROMPT_GEN, lines=20)
            submit_btn = gr.Button("🪄 Generate XML")

            # === Load Sample Images ===

            example_images = [
                ["images/page_0009.jpg"],
                ["images/page_0011.jpg"],
                ["images/page_0013.jpg"]
            ]

            gr.Examples(
            examples=example_images,
            inputs=image_input,
            label="🖼️ Example Images"
            )

        with gr.Column():
            xml_output = gr.Textbox(label="📄 Generated XML", lines=25, interactive=False)
            download_btn = gr.File(label="⬇️ Download XML")

    # submit_btn.click(fn=process_image, inputs=[image_input], outputs=[xml_output, download_btn])
    submit_btn.click(fn=process_image, inputs=[image_input, prompt_input], outputs=[xml_output, download_btn])

# Launch the app
demo.launch(debug=True)