import os
import base64
import re
from flask import Flask, render_template, request, redirect, url_for, flash
from openai import OpenAI
from PIL import Image
from io import BytesIO

# Initialize Flask app
app = Flask(__name__)
app.secret_key = 'supersecretkey'

# Set upload folder and ensure it exists with correct permissions
UPLOAD_FOLDER = 'static/uploads'
os.makedirs(UPLOAD_FOLDER, mode=0o777, exist_ok=True)  # Ensure folder exists with correct permissions
app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER

# Initialize OpenAI client
api_key = os.getenv("OPENAI_API_KEY")
if not api_key:
    raise ValueError("❌ Error: OPENAI_API_KEY is not set! Check Hugging Face Secrets.")
client = OpenAI(api_key=api_key)

# Function to convert image to Base64 PNG string
def convert_to_png_base64(image_path):
    try:
        with Image.open(image_path) as img:
            img = img.convert("RGBA")
            buffer = BytesIO()
            img.save(buffer, format="PNG")
            return base64.b64encode(buffer.getvalue()).decode("utf-8")
    except Exception as e:
        print(f"Error: {e}")
        return None

# Function to format response into clean HTML
def format_response(message_content):
    # Replace ### Heading with <h3>Heading</h3>
    formatted = re.sub(r'^### (.*?)$', r'<h3>\1</h3>', message_content, flags=re.M)
    # Replace **bold** with <strong>bold</strong>
    formatted = re.sub(r'\*\*(.*?)\*\*', r'<strong>\1</strong>', formatted)
    # Replace "- " with list items
    formatted = re.sub(r'^\- (.*?)$', r'<li>\1</li>', formatted, flags=re.M)
    # Wrap lists with <ul>
    if '<li>' in formatted:
        formatted = f"<ul>{formatted}</ul>"
    return formatted
    
@app.route('/', methods=['GET', 'POST'])
def index():
    if request.method == 'POST':
        if 'image' in request.files:
            file = request.files['image']
            if file.filename != '':
                filepath = os.path.join(app.config['UPLOAD_FOLDER'], file.filename)
                file.save(filepath)
                return render_template('index.html', uploaded_image=file.filename, response=None)

        if 'analyze' in request.form:
            image_name = request.form.get('uploaded_image')
            filepath = os.path.join(app.config['UPLOAD_FOLDER'], image_name)

            base64_image = convert_to_png_base64(filepath)
            if not base64_image:
                flash('Error processing image.')
                return redirect(request.url)

            # Send to OpenAI API
            response = client.chat.completions.create(
                model="gpt-4o-mini",
                messages=[
                    {
                        "role": "user",
                        "content": [
                            {
                                "type": "text",
                                "text": '''You are an expert in analyzing infrastructure issues related to manholes. Your task is to analyze an image of a manhole (inside the red rectangle) and generate a detailed report with the following fields:
                                - **Ticket Number:** Assign a unique identifier for the issue.
                                - **GPS Location:** Provide a fake but realistic GPS coordinate.
                                - **Severity:** Categorize the issue as Low, Medium, or High.
                                - **Comment:** Provide a brief explanation of the problem visible in the image.
                                - **Defect:** Identify the specific defect (e.g., cracked cover, open manhole, debris obstruction, etc.).
                                - **Department:** Always set as "Road/Manhole Management."
                                - **Action:** Suggest an appropriate action based on the severity and type of defect.''',
                            },
                            {
                                "type": "image_url",
                                "image_url": {"url": f"data:image/png;base64,{base64_image}"},
                            },
                        ],
                    }
                ],
            )

            message_content = response.choices[0].message.content
            formatted_content = format_response(message_content)
            return render_template('index.html', uploaded_image=image_name, response=formatted_content)

    return render_template('index.html', uploaded_image=None, response=None)

if __name__ == '__main__':
    app.run(host="0.0.0.0", port=7860, debug=False)  # Debug should be False in production