import os
import io
import base64
import json
import re
import numpy as np
from PIL import Image
import matplotlib.pyplot as plt
import matplotlib.patches as patches
import streamlit as st
from openai import OpenAI
# Set your API key and instantiate the client (make sure your OpenAI client is imported/defined)
from dotenv import load_dotenv
load_dotenv()
client = OpenAI()  # Assumes you have an OpenAI client available

# The prompt used to instruct the model
prompt = """
You are an expert road quality analyst. You will be shown an image of a road segment. Your task is to thoroughly inspect the condition of the road surface in the image and provide a detailed evaluation.

Your output must be in the form of a JSON object with the following structure:

{
  "description": "<A detailed analysis of the visible road surface condition. Describe texture, cracks, potholes, construction quality, and overall appearance.>",
  "label": "<One of: 'Excellent', 'Adequate', 'Basic', 'Poor', or 'Not Constructed'>",
  "score": <A numerical value from the following scale: 5, 3, 2, 1, or 0>
}

### Use the following scoring rubric to assign both the 'label' and the 'score':

- **5 – Excellent**: The road surface is in pristine condition, with no visible damage. It may be a newly constructed or expressway-type road with smooth asphalt and no visible cracks, potholes, or any other surface imperfections. It is **ideal for fast, high-volume traffic**.
  
- **3 – Adequate**: The road is constructed and mostly functional. There are a few minor defects such as small cracks or potholes, but the road is still in **generally good condition**. It is suitable for use but could use maintenance to address minor issues.
  
- **2 – Basic**: The road shows clear signs of wear and construction defects, such as visible potholes, cracks, or surface degradation. The **road is still usable but needs repair**, and it may be difficult or uncomfortable to drive on for extended periods.

- **1 – Poor**: The road has **major issues** such as large potholes, severe cracks, or other types of significant damage that make it **hard to use**. While the road may still be passable, it poses risks to vehicles and drivers due to the level of deterioration.

- **0 – Not Constructed**: The road is **incomplete or not constructed**. This could include images of roads with severe mud, large holes, missing surface material, or roads that are **barely passable or not usable** for normal traffic. It is **not fit for use** and potentially dangerous.

### Specific Classification Criteria:
- **0 (Not Constructed)**: Roads that are missing large sections, filled with severe mud, or have very large holes making them essentially unusable or extremely hazardous.
- **1 (Poor)**: Roads with major potholes, severe cracks, or significant surface degradation where it is barely usable. The road is **dangerous for normal traffic**.
- **2 (Basic)**: Roads with **multiple defects**, such as potholes, cracks, or surface degradation. These roads are **difficult to drive on** but are still passable.
- **3 (Adequate)**: Roads that are **mostly intact** with only a few minor defects, such as small potholes or cracks, but still in **acceptable condition** for regular use.
- **5 (Excellent)**: Pristine roads, like expressways or newly constructed highways, that are **in perfect condition** with no visible damage, cracks, or potholes.

Make sure your response contains **only the JSON output**, with no extra text or commentary.
"""

# prompt = '''You are an expert road quality analyst. You will be provided with an image of a road segment. Your task is to perform a detailed visual analysis of the road surface and output your evaluation as a JSON object strictly in the format below, with no extra commentary:

# {
#   "description": "<A detailed analysis of the visible road surface condition. Describe texture, cracks, potholes, construction quality, and overall appearance.>",
#   "label": "<One of: 'Excellent', 'Adequate', 'Basic', 'Poor', or 'Not Constructed'>",
#   "score": <A numerical value from the following scale: 5, 3, 2, 1, or 0>
# }

# Scoring Rubric:
# - **5 – Excellent**: The road surface is pristine, with no damage. It may be a newly constructed highway or expressway with smooth asphalt and no visible cracks, potholes, or imperfections. Ideal for fast, high-volume traffic.
  
# - **3 – Adequate**: The road is properly constructed and mostly functional. There might be a few minor defects like small cracks or potholes, but overall it is in generally good condition. Some minor maintenance could be beneficial.
  
# - **2 – Basic**: The road shows signs of wear and some construction defects, such as moderate cracks, potholes, or surface degradation. It remains passable but clearly needs repair. **For example, dirt roads that look flat and generally good should be classified as 'Basic' with a score of 2.**
  
# - **1 – Poor**: The road has major issues including large potholes, severe cracks, or significant surface degradation that make it hazardous. Although passable in some cases, it poses a risk to vehicles and drivers.
  
# - **0 – Not Constructed**: The road is incomplete or appears unconstructed. This includes images showing severe gaps, extensive mud, or missing sections, making the road extremely unsafe or unusable.

# Important Instructions:
# - **Consistency:** The "label" field must be exactly one of the specified text values ('Excellent', 'Adequate', 'Basic', 'Poor', or 'Not Constructed'). Do not output a numerical value as the label.
# - **Accurate Mapping:** The "score" field must correspond exactly to the assigned label based on the rubric. For instance, if you determine the road condition is "Basic", the score must be 2.
# - **Visual-Only Analysis:** Your evaluation must be based solely on what is visible in the provided image. Do not infer conditions beyond what the image shows.
# - **Output Format:** Only output a valid JSON object adhering to the above format—no additional text, explanation, or markdown formatting.

# Follow these instructions precisely to ensure accuracy and to minimize hallucinations.
# '''

# Function to create a base64 image URL from an uploaded file
def make_links_from_file(uploaded_file):
    # Ensure we're reading from the start
    uploaded_file.seek(0)
    with Image.open(uploaded_file) as img:
        img = img.convert("RGB")
        if img.width > 512 or img.height > 512:
            img.thumbnail((512, 512))
        buffer = io.BytesIO()
        img.save(buffer, format="PNG")
        encoded_image = base64.b64encode(buffer.getvalue()).decode("utf-8")
    return f"data:image/png;base64,{encoded_image}"

# Function to get the response from the model
def get_response(prompt, img_url, model="gpt-4o-mini"):
    if model == "gpt-4o":
        print("Using gpt-4o model")
        messages = [
        {
            "role": "user",
            "content": [
                { "type": "text", "text": prompt },
                {
                    "type": "image_url", 
                    "image_url": {"url": img_url}
                },
            ],
        }
        ]

    else:
        messages = [
            {
                "role": "user",
                "content": [
                    {"type": "text", "text": prompt},
                    {"type": "image_url", "image_url": {"url": img_url}}
                ]
            }
        ]
    response = client.chat.completions.create(
        model=model,
        messages=messages
    )
    return response

# Function to clean the output (removes markdown formatting and parses JSON)
def clean_response(output):
    raw_output = output.choices[0].message.content.strip()
    cleaned_output = re.sub(r"^```(?:json)?\s*|\s*```$", "", raw_output.strip())
    data = json.loads(cleaned_output)
    return data

# Function to display the image with an overlay and description
def show_outs(img, data):
    fig, ax = plt.subplots(figsize=(10, 6))
    ax.imshow(img)
    ax.axis('off')
    label_text = f"Label: {data['label']}  |  Score: {data['score']}"
    ax.add_patch(patches.Rectangle((0, 0), 1, 0.1, transform=ax.transAxes,
                                   color='black', alpha=0.5))
    ax.text(0.01, 0.03, label_text, transform=ax.transAxes,
            fontsize=14, color='white', fontweight='bold', va='center')
    plt.tight_layout()
    st.pyplot(fig)
    plt.close(fig)
    st.write("\n📝 Road Surface Description:\n")
    st.write(data.get("description", "No description provided."))

# Modified inference pipeline that works with an uploaded file
def inference_pipeline_uploaded(uploaded_file, prompt, model):
    img_url = make_links_from_file(uploaded_file)
    st.write("Fetching Response...")
    output = get_response(prompt, img_url, model)
    data = clean_response(output)
    # Reset file pointer and reopen image for display
    uploaded_file.seek(0)
    img = Image.open(uploaded_file).convert("RGB")
    if img.width > 512 or img.height > 512:
        img.thumbnail((512, 512))
    img_array = np.array(img)
    show_outs(img_array, data)

# Main app layout
st.title("Road Quality Analysis")
st.write("Upload an image of a road segment for analysis.")

# Toggle between models using a radio button
model_choice = st.radio("Select Model", options=["GPT 40", "GPT 40 mini"])
model = "gpt-4o-2024-08-06" if model_choice == "GPT 40" else "gpt-4o-mini"

# Image file uploader
uploaded_file = st.file_uploader("Choose an image file", type=["jpg", "jpeg", "png"])

if uploaded_file is not None:
    st.image(uploaded_file, caption="Uploaded Road Segment", use_column_width=True)
    if st.button("Run Analysis"):
        inference_pipeline_uploaded(uploaded_file, prompt, model)