main.py

from fastapi import FastAPI, Request, Form
from fastapi.responses import HTMLResponse, JSONResponse
from fastapi.staticfiles import StaticFiles
from fastapi.templating import Jinja2Templates
import os
import json
import markdown
from markdown.extensions import fenced_code
import math
from typing import List

app = FastAPI()

# Serve static files (images)
app.mount("/glitch_data", StaticFiles(directory="glitch_data"), name="glitch_data")

# Set up Jinja2 templates
templates = Jinja2Templates(directory="templates")

# Add these constants at the top of the file
BUG_TAGS = ["Contains a bug", "Does not contain a bug"]
ANALYSIS_TAGS = ["Multiple images", "Requires further analysis", "Bad Sample"]
ISSUE_TAGS = ["Floating object", "Texture issues", "Low quality model"]

def load_results():
    global ISSUE_TAGS
    with open("glitch_data/all_search_results.json", "r") as f:
        all_results = json.load(f)
    
    for game, results in all_results.items():
        for result in results:
            image_dir = os.path.dirname(result["local_image_file"])
            metadata_path = os.path.join(image_dir, "metadata.json")
            if os.path.exists(metadata_path):
                with open(metadata_path, "r") as f:
                    metadata = json.load(f)
                    result["tags"] = metadata.get("tags", {"bug": [], "analysis": [], "issue": []})
                    # Update ISSUE_TAGS with any custom tags found in the metadata
                    for tag in result["tags"]["issue"]:
                        if tag not in ISSUE_TAGS:
                            ISSUE_TAGS.append(tag)
            else:
                result["tags"] = {"bug": [], "analysis": [], "issue": []}
    
    return all_results

def save_metadata(image_path, tags):
    global ISSUE_TAGS
    image_dir = os.path.dirname(image_path)
    metadata_path = os.path.join(image_dir, "metadata.json")
    with open(metadata_path, "w") as f:
        json.dump({"tags": tags}, f, indent=2)
    # Update ISSUE_TAGS with any new tags
    for tag in tags["issue"]:
        if tag not in ISSUE_TAGS:
            ISSUE_TAGS.append(tag)

# Load results and update ISSUE_TAGS when the app starts
all_results = load_results()

ITEMS_PER_PAGE = 10

@app.get("/", response_class=HTMLResponse)
async def root(request: Request, game: str = None, page: int = 1, view: str = "list", highlight: str = None):
    if game is None:
        return templates.TemplateResponse("game_selection.html", {
            "request": request,
            "games": list(all_results.keys())
        })

    game_results = all_results[game]
    total_pages = math.ceil(len(game_results) / ITEMS_PER_PAGE)
    start_index = (page - 1) * ITEMS_PER_PAGE
    end_index = start_index + ITEMS_PER_PAGE
    
    if view == "list":
        current_results = game_results[start_index:end_index]
    else:
        current_results = game_results  # Load all results for gallery view

    for result in current_results:
        image_dir = os.path.dirname(result["local_image_file"])
        content_path = os.path.join(image_dir, "symmary_of_the_article.md")
        
        if os.path.exists(content_path):
            with open(content_path, "r") as f:
                content = f.read()
                result["extracted_content_html"] = markdown.markdown(content, extensions=[fenced_code.FencedCodeExtension()])
                result["extracted_content_raw"] = content
        else:
            result["extracted_content_html"] = ""
            result["extracted_content_raw"] = ""

        # Add a flag to identify the highlighted image
        result["highlight"] = (result["local_image_file"] == highlight)

    return templates.TemplateResponse("index.html", {
        "request": request,
        "results": current_results,
        "page": page,
        "total_pages": total_pages,
        "view": view,
        "all_results": game_results,
        "highlight": highlight,
        "ITEMS_PER_PAGE": ITEMS_PER_PAGE,
        "BUG_TAGS": BUG_TAGS,
        "ANALYSIS_TAGS": ANALYSIS_TAGS,
        "ISSUE_TAGS": ISSUE_TAGS,
        "game": game,
    })

@app.post("/tag")
async def tag_image(image_path: str = Form(...), tag: str = Form(...), active: bool = Form(...), tag_type: str = Form(...)):
    for game, results in all_results.items():
        for result in results:
            if result["local_image_file"] == image_path:
                if "tags" not in result:
                    result["tags"] = {"bug": [], "analysis": [], "issue": []}
                
                if active:
                    if tag not in result["tags"][tag_type]:
                        if tag_type == "bug":
                            result["tags"]["bug"] = []  # Remove conflicting bug tag
                        result["tags"][tag_type].append(tag)
                else:
                    if tag in result["tags"][tag_type]:
                        result["tags"][tag_type].remove(tag)
                
                save_metadata(result["local_image_file"], result["tags"])
                return JSONResponse({"status": "success", "tags": result["tags"]})
    
    return JSONResponse({"status": "error", "message": "Image not found"})

@app.post("/add_tag")
async def add_tag(tag: str = Form(...), image_path: str = Form(...)):
    if tag not in ISSUE_TAGS:
        ISSUE_TAGS.append(tag)
    
    # Find the correct result and update its tags
    for game, results in all_results.items():
        for result in results:
            if result["local_image_file"] == image_path:
                if tag not in result["tags"]["issue"]:
                    result["tags"]["issue"].append(tag)
                save_metadata(result["local_image_file"], result["tags"])
                break
    
    return JSONResponse({"status": "success", "tags": ISSUE_TAGS})

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)