import os
import base64
from pathlib import Path
from dotenv import load_dotenv
from openai import OpenAI

load_dotenv()

client = OpenAI(
    base_url="https://openrouter.ai/api/v1",
    api_key=os.environ["OPENROUTER_API_KEY"],
)

# --- CONFIGURATION ---
IMAGE_DIR = "./images"       # folder with your images
OUTPUT_HTML = "gallery.html" # output file
MODEL = "openai/gpt-4o"     # must be a vision-capable model

# --- THE PROMPT ---
# This is where the quality lives. Notice: structured output format,
# length constraint, style guidance — all in one prompt.

SYSTEM_PROMPT = """You are a scientific image archivist. Your task is to write
concise, objective descriptions of images for a research gallery.

For each image, provide exactly this JSON structure:
{
  "filename": "the filename",
  "title": "A short descriptive title (max 60 characters)",
  "description": "A factual, neutral description (2-3 sentences). Describe what
    is visible: subject, setting, key visual elements, colours, composition.
    Do not interpret or evaluate — describe what you see.",
  "category": "one of: microscopy, field_photo, diagram, equipment, gel_blot,
    graph, portrait, other"
}

Rules:
- Be factual, not poetic
- Describe the visual content, not what it "means"
- Max 80 words for the description
- If the image is unclear, say so honestly
"""

# --- IMAGE ENCODING ---
def encode_image(path):
    """Read an image file and encode it as base64."""
    with open(path, "rb") as f:
        return base64.b64encode(f.read()).decode("utf-8")

# --- PROCESS ONE IMAGE ---
def describe_image(image_path):
    """Send an image to the LLM and get back a structured description."""
    ext = image_path.suffix.lower()
    mime_map = {
        ".jpg": "image/jpeg", ".jpeg": "image/jpeg",
        ".png": "image/png", ".gif": "image/gif",
        ".webp": "image/webp",
    }
    mime_type = mime_map.get(ext, "image/jpeg")

    b64 = encode_image(image_path)

    response = client.chat.completions.create(
        model=MODEL,
        messages=[
            {"role": "system", "content": SYSTEM_PROMPT},
            {
                "role": "user",
                "content": [
                    {
                        "type": "text",
                        "text": f"Describe this image. Filename: {image_path.name}"
                    },
                    {
                        "type": "image_url",
                        "image_url": {
                            "url": f"data:{mime_type};base64,{b64}"
                        },
                    },
                ],
            },
        ],
        response_format={"type": "json_object"},
    )

    return response.choices[0].message.content

# --- BUILD HTML ---
def build_html(descriptions):
    """Generate an HTML table from the image descriptions."""
    rows = ""
    for desc in descriptions:
        rows += f"""<tr>
            <td><img src="{desc['filename']}" width="150" loading="lazy"></td>
            <td><a href="{desc['filename']}">{desc['filename']}</a></td>
            <td><strong>{desc['title']}</strong></td>
            <td>{desc['description']}</td>
            <td><code>{desc['category']}</code></td>
        </tr>\n"""

    return f"""<!DOCTYPE html>
<html lang="en">
<head><meta charset="UTF-8"><title>Image Gallery</title>
<style>
  body 
  table 
  th, td 
  th 
  img 
  tr:hover 
</style></head>
<body>
<h1>Image Gallery</h1>
<table>
<tr><th>Thumbnail</th><th>File</th><th>Title</th><th>Description</th><th>Category</th></tr>
{rows}
</table>
</body></html>"""

# --- MAIN ---
def main():
    import json

    image_extensions = {".jpg", ".jpeg", ".png", ".gif", ".webp", ".bmp"}
    images = sorted([
        p for p in Path(IMAGE_DIR).iterdir()
        if p.suffix.lower() in image_extensions
    ])

    if not images:
        print(f"No images found in {IMAGE_DIR}")
        return

    print(f"Found {len(images)} images. Processing...")

    results = []
    for i, img in enumerate(images, 1):
        print(f"  [{i}/{len(images)}] {img.name} ...", end=" ")
        try:
            raw = describe_image(img)
            desc = json.loads(raw)
            results.append(desc)
            print("OK")
        except Exception as e:
            print(f"FAILED: {e}")
            results.append({
                "filename": img.name,
                "title": "ERROR",
                "description": str(e),
                "category": "other",
            })

    html = build_html(results)
    with open(OUTPUT_HTML, "w") as f:
        f.write(html)

    print(f"\nDone! Open {OUTPUT_HTML} in your browser.")

if __name__ == "__main__":
    main()