import streamlit as st
import base64
import openai
import numpy as np
import cv2
from tensorflow.keras.models import load_model
from keras.preprocessing.image import img_to_array
from keras.applications.inception_v3 import preprocess_input
import requests
import json_repair
import json
from PIL import Image
import io
from datetime import datetime
import pandas as pd
import re
from json_repair import repair_json
from dotenv import load_dotenv
import os

# Load environment variables
load_dotenv()

# Class labels for defect prediction
class_labels = [
    "algae",    
    "bubbles",
    "Cracks",
    "Fungus",
    "peeling",
    
]

# Set OpenAI API key
openai.api_key = os.getenv("OPENAI_API_KEY")

# Page configuration
st.set_page_config(
    page_title="Painting Defect Detection System",
    page_icon="🔍",
    layout="wide",
    initial_sidebar_state="expanded"
)

# Custom CSS for better styling
st.markdown("""
<style>
    .main {
        padding-top: 2rem;
    }
    .stAlert {
        margin-top: 1rem;
    }
    .defect-card {
        border: 1px solid #ddd;
        border-radius: 8px;
        padding: 1rem;
        margin: 1rem 0;
        background-color: #f9f9f9;
    }
    .metric-card {
        background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
        padding: 1rem;
        border-radius: 8px;
        color: white;
        text-align: center;
        margin: 0.5rem 0;
    }
</style>
""", unsafe_allow_html=True)

# Cache the model loading to improve performance
@st.cache_resource
def load_trained_model():
    try:
        return load_model('painting.keras')
    except Exception as e:
        st.error(f"Error loading model: {str(e)}")
        return None

def compress_image(image_bytes, max_size_kb=500):
    img = Image.open(io.BytesIO(image_bytes))
    quality = 95
    output_bytes = io.BytesIO()
    while True:
        output_bytes.seek(0)
        output_bytes.truncate()
        img.save(output_bytes, format='JPEG', quality=quality)
        if len(output_bytes.getvalue()) <= max_size_kb * 1024 or quality <= 5:
            break
        quality -= 5
    return output_bytes.getvalue()

def get_direct_drive_url(share_url):
    if "drive.google.com/file/d/" in share_url:
        file_id = share_url.split("/file/d/")[1].split("/")[0]
        return f"https://drive.google.com/uc?export=download&id={file_id}"
    return share_url

def process_image(url):
    try:
        file_id = re.search(r'/d/(.*?)/', url)
        if file_id:  # Google Drive share link
            url = f'https://drive.google.com/uc?export=download&id={file_id.group(1)}'
        
        resp = requests.get(url, timeout=15)
        resp.raise_for_status()
        
        if 'image' not in resp.headers.get('Content-Type', ''):
            raise ValueError('URL does not point to an image')

        img = cv2.imdecode(np.frombuffer(resp.content, np.uint8), cv2.IMREAD_COLOR)
        input_img_resized = cv2.resize(img, (256, 256))
        x = img_to_array(input_img_resized)
        x = np.expand_dims(x, axis=0)
        x = preprocess_input(x)
        return x, resp.content
    except Exception as e:
        st.error(f"Error processing image from URL: {e}")
        return None, None

def process_uploaded_image(uploaded_file):
    try:
        image_bytes = uploaded_file.read()
        img = Image.open(io.BytesIO(image_bytes))
        img_array = np.array(img)
        
        # Convert RGB to BGR for OpenCV
        if len(img_array.shape) == 3:
            img_array = cv2.cvtColor(img_array, cv2.COLOR_RGB2BGR)
        
        input_img_resized = cv2.resize(img_array, (256, 256))
        x = img_to_array(input_img_resized)
        x = np.expand_dims(x, axis=0)
        x = preprocess_input(x)
        
        return x, image_bytes
    except Exception as e:
        st.error(f"Error processing uploaded image: {e}")
        return None, None

def predict_defect(processed_image, model):
    try:
        preds = model.predict(processed_image)[0]
        class_index = int(np.argmax(preds))
        confidence = float(preds[class_index])
        class_name = class_labels[class_index]

        return {
            "type": class_name,
            "severity": confidence
        }, class_index
    except Exception as e:
        return {"error": str(e)}, []

def compress_image_to_base64(image_bytes, max_size_kb=500):
    img = Image.open(io.BytesIO(image_bytes)).convert("RGB")
    output_bytes = io.BytesIO()
    quality = 95
    while True:
        output_bytes.seek(0)
        output_bytes.truncate()
        img.save(output_bytes, format='JPEG', quality=quality)
        if len(output_bytes.getvalue()) <= max_size_kb * 1024 or quality <= 5:
            break
        quality -= 5
    return base64.b64encode(output_bytes.getvalue()).decode('utf-8')

def generate_openai_description(classification, image_bytes):
    try:
        if isinstance(classification, dict) and 'type' in classification:
            defect_type = classification['type']
        elif isinstance(classification, dict) and 'error' in classification:
            defect_type = "Unknown defect"
        else:
            defect_type = str(classification)

        # Compress and encode image to base64
        compressed_base64 = compress_image_to_base64(image_bytes)
        ai_prompt = (
                f"You are an expert painting inspector analyzing surface defects."
                f"DETECTED DEFECTS: {defect_type}"
                f"TASK: Examine the painted surface image and provide a technical description of the visible defects."
                f"REQUIREMENTS:"
                f"- Write 200 characters describing what you observe"
                f"- Focus on visible paint failures, surface irregularities, or coating issues"
                f"- Use professional painting and coating terminology"
                f"- Be specific about location, extent, and characteristics of defects"
                f"- If multiple defects are present, describe each one"
                f"- Maintain an objective, technical tone"
                f"Dont use \\n"
                f"IMPORTANT: Always provide a description based on what you can see in the image. "
                f"Even if the detected defect types seem unclear, describe the actual visible conditions "
                f"in the painted surface. Only respond with 'Unable to generate description due to "
                f"image quality issues' if the image is genuinely too blurry, dark, or corrupted to analyze."
                f"Begin your description now:")

       

        # Call OpenAI with image and text
        response = openai.ChatCompletion.create(
            model="gpt-4o",
            messages=[
                {
                    "role": "user",
                    "content": [
                        {"type": "text", "text": ai_prompt},
                        {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{compressed_base64}"}}
                    ]
                }
            ]
        )

        result_text = response.choices[0].message.content.strip()
        return result_text

    except Exception as e:
        return f"Error generating description: {str(e)}"


    

def get_sps(defect_type, flooring_type,defect_category="flooring"):
    ai_content = """ You are an expert in identifying and resolving painting and wall surface defects.
Your task is to accept a defect name and surface type as input and return applicable solutions using the structured dataset below.
IMPORTANT: You MUST respond with valid JSON only. Do not include any explanatory text before or after the JSON.
Instructions:
Return TWO separate categories of solutions:

MANDATORY SPS: Solutions specifically matching the input defect
OPTIONAL SPS: Solutions from 'All Defects' records (universal solutions)

For MANDATORY SPS:

Match the input defect against the 'Defects' column using:

Exact match
Partial match (e.g. 'Cracks' matches all 'Cracks' records)


EXCLUDE any 'All Defects' records from this section
Only return records where the 'Category' matches the specified surface type (NonWaterproofing or Waterproofing)

For OPTIONAL SPS:

ONLY include records where 'Defects' is 'All Defects', 'ALL Defects', or 'All Defect'
Filter by surface type:

If surface_type is "NonWaterproofing": Return records with Category="NonWaterproofing"
If surface_type is "Waterproofing": Return records with Category="Waterproofing"


Return only those matching the specified surface type

For each valid record in both categories:

Use the id from the dataset. If not available, generate a unique numeric ID.
Use the title from the dataset. If not available, use appropriate fallback title based on surface type
Use the description from the dataset exactly (verbatim). If not found, generate fallback content based on best practices.
CRITICAL FOR PRODUCT EXTRACTION:

Carefully analyze the description content for ANY product names, brand names, or material references
Extract ALL product names mentioned in the description, including:

Brand names (e.g., "Asian Paints SmartCare Crack Seal")
Generic product names (e.g., "Economy Primer", "Ace Suprema")
Material names (e.g., "white cement", "fine sand")
Tool names (e.g., "Putty Knife", "Paint Scraper")
Chemical names (e.g., "oxygen bleach", "chlorinated bleach")


List them in products array, assigning a unique ID per product
NEVER leave products array empty - if no specific products are mentioned, include generic alternatives like "Primer", "Paint", "Cleaning Solution", etc.


Include all records where Category matches the filtering condition
IMPORTANT: Remove all HTML tags from the description content. Convert HTML lists to plain text format.

Product Extraction Examples:

"Use Asian Paints SmartCare Crack Seal" → Extract: "Asian Paints SmartCare Crack Seal"
"Economy Primer, Ace Suprema" → Extract: "Economy Primer", "Ace Suprema"
"Use boiling water, pressure washers, oxygen bleach" → Extract: "Pressure Washer", "Oxygen Bleach"
"Sand with medium-grit sandpaper" → Extract: "Medium-grit Sandpaper"



If no matching dataset entry is found for MANDATORY SPS, return one fallback solution following industry-standard practices for the specified surface type (without stating it is a fallback).
Each solution must follow standard practices such as:

Surface preparation
Cleaning and washing
Priming
Crack filling
Sanding
Paint application
Sealing
Moisture treatment
Finishing
    Each solution must follow standard practices such as:
    - Surface preparation
    - Cleaning and washing
    - Priming
    - Crack filling
    - Sanding
    - Paint application
    - Sealing
    - Moisture treatment
    - Finishing

    Return ONLY this JSON format (no other text):

    {
      "mandatory_sps": [
        {
          "id": 181,
          "title": "Surface Preparation for Cracks - NonWaterproofing",
          "content": "New masonry surfaces must be allowed to cure completely. It is recommended to allow 28 days as the curing time for new masonry surfaces. Opening the cracks with the help of a grinder/cutter and cleaning the surface with pressure washers for further treatment. FILLING FOR CRACKS For filling cracks up to 3mm, use Asian Paints SmartCare Crack Seal. For filling cracks more than 3mm, use Asian Paints SmartCare Textured Crack Filler. FILLING FOR HOLES & DENTS In case of dents and holes, use TruCare Wall Putty Suprema or a mix of white cement and fine sand in the ratio 1:3.",
          "products": [
            {
              "id": 1,
              "product_name": "Asian Paints SmartCare Crack Seal"
            },
            {
              "id": 2,
              "product_name": "Asian Paints SmartCare Textured Crack Filler"
            },
            {
              "id": 3,
              "product_name": "TruCare Wall Putty Suprema"
            }
          ]
        }
      ],
      "optional_sps": [
        {
          "id": 211,
          "title": "4 Year Paint Warranty System - NonWaterproofing",
          "content": "Economy Primer, Ace Suprema",
          "products": [
            {
              "id": 4,
              "product_name": "Economy Primer"
            },
            {
              "id": 5,
              "product_name": "Ace Suprema"
            }
          ]
        }
      ]
    }

    Dataset:

    id: 181
    Defects: Cracks
    Category: NonWaterproofing
    title: Surface Preparation for Cracks - NonWaterproofing
    description: New masonry surfaces must be allowed to cure completely. It is recommended to allow 28 days as the curing time for new masonry surfaces. Opening the cracks with the help of a grinder/cutter and cleaning the surface with pressure washers for further treatment.

    FILLING FOR CRACKS
    For filling cracks up to 3mm, use Asian Paints SmartCare Crack Seal.
    For filling cracks more than 3mm, use Asian Paints SmartCare Textured Crack Filler.

    FILLING FOR HOLES & DENTS
    In case of dents and holes, use TruCare Wall Putty Suprema or a mix of white cement and fine sand in the ratio 1:3.

    id: 182
    Defects: Cracks
    Category: Waterproofing
    title: Surface Preparation for Cracks - Waterproofing
    description: New masonry surfaces must be allowed to cure completely. It is recommended to allow 28 days as the curing time for new masonry surfaces. Opening the cracks with the help of a grinder/cutter and cleaning the surface with pressure washers for further treatment.

    FILLING FOR CRACKS
    For filling cracks up to 3mm, use Asian Paints SmartCare Crack Seal.
    For filling cracks more than 3mm, use Asian Paints SmartCare Textured Crack Filler.

    FILLING FOR HOLES & DENTS
    In case of dents and holes, use TruCare Wall Putty Suprema or a mix of white cement and fine sand in the ratio 1:3.

    id: 183
    Defects: Patchiness
    Category: NonWaterproofing
    title: Surface Preparation for Patchiness - NonWaterproofing
    description: Sand the surface with medium-grit sandpaper to remove major surface irregularities bumps and rough patches. Switch to fine-grit sandpaper for final smoothing. Work in small circular motions overlapping each area. Remove all sanding dust by cleaning the surface with pressure washers for further treatment

    id: 184
    Defects: Patchiness
    Category: Waterproofing
    title: Surface Preparation for Patchiness - Waterproofing
    description: Sand the surface with medium-grit sandpaper to remove major surface irregularities bumps and rough patches. Switch to fine-grit sandpaper for final smoothing. Work in small circular motions overlapping each area. Remove all sanding dust by cleaning the surface with pressure washers for further treatment

    id: 185
    Defects: Peeling
    Category: NonWaterproofing
    title: Surface Preparation for Peeling - NonWaterproofing
    description: Use Putty Knife or Paint Scraper for removing loose and peeling paint without damaging the underlying surface.Sand the surface with medium-grit sandpaper to remove major surface irregularities bumps and rough patches. Switch to fine-grit sandpaper for final smoothing. Work in small circular motions overlapping each area. Remove all sanding dust by cleaning the surface for further treatment
    Surface should be free from any loose paint, dust or grease.Growths of fungus, algae or moss should be removed by wire brushing and water.

    id: 186
    Defects: Peeling
    Category: Waterproofing
    title: Surface Preparation for Peeling - Waterproofing
    description: Use Putty Knife or Paint Scraper for removing loose and peeling paint without damaging the underlying surface.Sand the surface with medium-grit sandpaper to remove major surface irregularities bumps and rough patches. Switch to fine-grit sandpaper for final smoothing. Work in small circular motions overlapping each area. Remove all sanding dust by cleaning the surface for further treatment
    Surface should be free from any loose paint, dust or grease.Growths of fungus, algae or moss should be removed by wire brushing and water.

    id: 187
    Defects: Algae
    Category: NonWaterproofing
    title: Surface Preparation for Algae - NonWaterproofing
    description: Use boiling water, pressure washers, oxygen bleach, chlorinated bleach, and commercial algae removers to get rid of green algae on concrete. Pressure wash and scrape the wall with a soft brush

    id: 188
    Defects: Algae
    Category: Waterproofing
    title: Surface Preparation for Algae - Waterproofing
    description: Use boiling water, pressure washers, oxygen bleach, chlorinated bleach, and commercial algae removers to get rid of green algae on concrete. Pressure wash and scrape the wall with a soft brush

    id: 189
    Defects: Bubble
    Category: NonWaterproofing
    title: Surface Preparation for Bubble - NonWaterproofing
    description: If moisture is the issue, fix any leaks or dampness before proceeding. Moisture must be fully resolved to prevent new bubbling. Remove the buubled paint by using the scraper to peel off the loose and raised areas.Sand the surface with medium-grit sandpaper to remove major surface irregularities bumps and rough patches. Switch to fine-grit sandpaper for final smoothing. Work in small circular motions overlapping each area. Remove all sanding dust by cleaning the surface with pressure washers for further treatment

    id: 190
    Defects: Bubble
    Category: Waterproofing
    title: Surface Preparation for Bubble - Waterproofing
    description: If moisture is the issue, fix any leaks or dampness before proceeding. Moisture must be fully resolved to prevent new bubbling. Remove the buubled paint by using the scraper to peel off the loose and raised areas.Sand the surface with medium-grit sandpaper to remove major surface irregularities bumps and rough patches. Switch to fine-grit sandpaper for final smoothing. Work in small circular motions overlapping each area. Remove all sanding dust by cleaning the surface with pressure washers for further treatment

    id: 191
    Defects: Blisters
    Category: NonWaterproofing
    title: Surface Preparation for Blisters - NonWaterproofing
    description: Remove blisters by scraping, sanding or pressure-washing down to underlying coats of paint or primer.Repair loose caulking and improve ventilation of the building to prevent a recurring problem

    id: 192
    Defects: Blisters
    Category: Waterproofing
    title: Surface Preparation for Blisters - Waterproofing
    description: Remove blisters by scraping, sanding or pressure-washing down to underlying coats of paint or primer.Repair loose caulking and improve ventilation of the building to prevent a recurring problem

    id: 193
    Defects: Efflorescense
    Category: NonWaterproofing
    title: Surface Preparation for Efflorescense - NonWaterproofing
    description: Stiff brushes are used to sweep away efflorescence from smoother surfaces.Apply a forceful water rinse to efflorescence with a pressure washer to clean it up

    id: 194
    Defects: Efflorescense
    Category: Waterproofing
    title: Surface Preparation for Efflorescense - Waterproofing
    description: Stiff brushes are used to sweep away efflorescence from smoother surfaces.Apply a forceful water rinse to efflorescence with a pressure washer to clean it up

    id: 195
    Defects: Fungus
    Category: NonWaterproofing
    title: Surface Preparation for Fungus - NonWaterproofing
    description: Use boiling water, pressure washers, oxygen bleach, chlorinated bleach, and commercial fungus removers to get rid of fungus on concrete. Pressure wash and scrape the wall with a soft brush

    id: 196
    Defects: Fungus
    Category: Waterproofing
    title: Surface Preparation for Fungus - Waterproofing
    description: Use boiling water, pressure washers, oxygen bleach, chlorinated bleach, and commercial fungus removers to get rid of fungus on concrete. Pressure wash and scrape the wall with a soft brush

    id: 197
    Defects: Poor Hiding
    Category: NonWaterproofing
    title: Surface Preparation for Poor Hiding - NonWaterproofing
    description: Sand the surface with medium-grit sandpaper to remove major surface irregularities bumps and rough patches. Switch to fine-grit sandpaper for final smoothing. Work in small circular motions overlapping each area. Remove all sanding dust by cleaning the surface with pressure washers for further treatment

    id: 198
    Defects: Poor Hiding
    Category: Waterproofing
    title: Surface Preparation for Poor Hiding - Waterproofing
    description: Sand the surface with medium-grit sandpaper to remove major surface irregularities bumps and rough patches. Switch to fine-grit sandpaper for final smoothing. Work in small circular motions overlapping each area. Remove all sanding dust by cleaning the surface with pressure washers for further treatment

    id: 199
    Defects: Shade variation
    Category: NonWaterproofing
    title: Surface Preparation for Shade variation - NonWaterproofing
    description: Remove Loose Material: Carefully strip away loose or blistered paint to expose a firm surface.Spot Priming: Apply an appropriate primer to the affected area, ensuring adhesion.

    id: 200
    Defects: Shade variation
    Category: Waterproofing
    title: Surface Preparation for Shade variation - Waterproofing
    description: Remove Loose Material: Carefully strip away loose or blistered paint to expose a firm surface.Spot Priming: Apply an appropriate primer to the affected area, ensuring adhesion.

    id: 211
    Defects: All Defects
    Category: NonWaterproofing
    title: 4 Year Paint Warranty System - NonWaterproofing
    description: Economy Primer, Ace Suprema

    id: 212
    Defects: All Defects
    Category: NonWaterproofing
    title: 7 year Paint Warranty System - NonWaterproofing
    description: Economy Primer, Apex Suprema

    id: 213
    Defects: All Defects
    Category: NonWaterproofing
    title: 8 year Paint Warranty System - NonWaterproofing
    description: Exterior Primer, Ultima Strech Suprema

    id: 214
    Defects: All Defects
    Category: NonWaterproofing
    title: 9 year Paint Warranty System - NonWaterproofing
    description: Exterior Primer, Ultima Suprema

    id: 215
    Defects: All Defects
    Category: NonWaterproofing
    title: 12 year Paint Warranty System - NonWaterproofing
    description: Exterior Primer, Protek Topcoat

    id: 216
    Defects: All Defects
    Category: NonWaterproofing
    title: 15 year Paint Warranty System - NonWaterproofing
    description: Exterior Primer, Duralife Topcoat

    id: 217
    Defects: All Defects
    Category: Waterproofing
    title: 4 year Complete Paint Warranty System - Waterproofing
    description: Damp Sheath suprema, Ace Suprema

    id: 218
    Defects: All Defects
    Category: Waterproofing
    title: 6 year Complete Paint Warranty System - Waterproofing
    description: Damp Proof Suprema, Apex Suprema

    id: 219
    Defects: All Defects
    Category: Waterproofing
    title: 8 year Complete Paint Warranty System - Waterproofing
    description: Damp Proof Suprema, Ultima Strech Suprema

    id: 220
    Defects: All Defects
    Category: Waterproofing
    title: 10 year Complete Paint Warranty System - Waterproofing
    description: Damp Prime Xtreme, Ultima Stretch Suprema

    id: 221
    Defects: All Defects
    Category: Waterproofing
    title: 12 year Complete Paint Warranty System - Waterproofing
    description: Protek Basecoat, Protek Topcoat

    id: 222
    Defects: All Defects
    Category: Waterproofing
    title: 15 year Complete Paint Warranty System - Waterproofing
    description: Duralife Basecoat, Duralife Topcoat

    """
    
   
   



    def remove_html_tags(text):
        """Remove HTML tags and convert lists to plain text format"""
        if not text:
            return text

        text = re.sub(r'<li>(.*?)</li>', r'• \1', text, flags=re.IGNORECASE | re.DOTALL)

        text = re.sub(r'<[^>]+>', '', text)

        text = re.sub(r'\s+', ' ', text).strip()
        
        text = text.replace('\n', ' ').replace('\r', '')

        return text



    try:
      
        


        response = openai.ChatCompletion.create(
            model="gpt-4o",
            messages=[
                {
                    "role": "system",
                    "content": ai_content
                },
                {
                    "role": "user",
                    "content": f"Generate recommended solutions for defect '{defect_type}' with painting  type '{flooring_type}'. Return valid JSON only with both mandatory SPS and optional SPS that match the '{type}' . Remove HTML tags from content."
                }
            ],
            max_tokens=2000,  # Increased token limit
            temperature=0.1,  # Lower temperature for more consistent responses
        )


        ai_description = response.choices[0].message.content.strip()
        
        ai_description = ai_description.replace('\n', ' ').replace('\r', '')
        
        clean_response = ai_description
        if clean_response.startswith('```json'):
            clean_response = clean_response[7:]
        elif clean_response.startswith('```'):
            clean_response = clean_response[3:]
        if clean_response.endswith('```'):
            clean_response = clean_response[:-3]
        clean_response = clean_response.strip()

        try:
            import json
            response_data = json.loads(clean_response)

            mandatory_sps = []
            if 'mandatory_sps' in response_data and isinstance(response_data['mandatory_sps'], list):
                for item in response_data['mandatory_sps']:
                    if isinstance(item, dict) and 'content' in item:
                        item['content'] = remove_html_tags(item['content'])
                    mandatory_sps.append(item)

            optional_sps = []
            if 'optional_sps' in response_data and isinstance(response_data['optional_sps'], list):
                for item in response_data['optional_sps']:
                    if isinstance(item, dict) and 'content' in item:
                        item['content'] = remove_html_tags(item['content'])
                    optional_sps.append(item)

            if not mandatory_sps and not optional_sps:
                print("AI returned empty response")
                # return create_fallback_response(defect_type, flooring_type)

            return mandatory_sps, optional_sps

        except json.JSONDecodeError as json_error:
        
            
            try:
                json_start = clean_response.find('{')
                json_end = clean_response.rfind('}')
                
                if json_start != -1 and json_end != -1 and json_end > json_start:
                    potential_json = clean_response[json_start:json_end+1]
                    
                    potential_json = potential_json.replace("'", '"')  # Replace single quotes
                    potential_json = re.sub(r',\s*}', '}', potential_json)  # Remove trailing commas
                    potential_json = re.sub(r',\s*]', ']', potential_json)  # Remove trailing commas in arrays
                    
                    response_data = json.loads(potential_json)
                    
                    mandatory_sps = []
                    if 'mandatory_sps' in response_data and isinstance(response_data['mandatory_sps'], list):
                        for item in response_data['mandatory_sps']:
                            if isinstance(item, dict) and 'content' in item:
                                item['content'] = remove_html_tags(item['content'])
                            mandatory_sps.append(item)

                    optional_sps = []
                    if 'optional_sps' in response_data and isinstance(response_data['optional_sps'], list):
                        for item in response_data['optional_sps']:
                            if isinstance(item, dict) and 'content' in item:
                                item['content'] = remove_html_tags(item['content'])
                            optional_sps.append(item)

                    return mandatory_sps, optional_sps
                    
            except Exception as recovery_error:
                print(f"JSON recovery failed: {recovery_error}")
            
            print("Using fallback response due to JSON parsing failure")
           
    except Exception as e:
        print(f"Error in OpenAI API call: {str(e)}")
       

def main():
    st.title("🔍 Painting Defect Detection System")
    st.markdown("---")

    # Load the model
    model = load_trained_model()
    if model is None:
        st.error("Failed to load the defect detection model. Please check if the model file exists.")
        return

    # Sidebar for input options
    st.sidebar.header("📋 Input Configuration")
    
    # Report details
    st.sidebar.subheader("Report Details")
    report_submission_id = st.sidebar.text_input("Report Submission Id", value=1)
    report_section_page_id = st.sidebar.text_input("Page ID", value=7)
    report_section_id = st.sidebar.text_input("Section ID", value=10)
    defect_type = st.sidebar.text_input("Defect Type", value="painting")
    type = st.sidebar.text_input("Type", value="NonWaterproofing")

    # Input method selection
    input_method = "Image URLs"

    # Main content area
    col1, col2 = st.columns([2, 1])

    with col1:
        st.header("📸 Image Input")
        st.subheader("Enter Image URLs")

        
        image_data_list = []
        classifications = []
        
      
        
                 
        num_urls = st.number_input("Number of images", min_value=1, max_value=10, value=1)
        
        urls = []
        for i in range(num_urls):
            url = st.text_input(f"Image URL {i+1}", key=f"url_{i}")
            if url:
                urls.append(url)
        
        if st.button("Process URLs", type="primary"):
            for i, url in enumerate(urls):
                if url.strip():
                    st.subheader(f"Image {i+1}")
                    
                    with st.spinner(f"Processing image {i+1}..."):
                        processed_img, image_bytes = process_image(url)
                        
                        if processed_img is not None and image_bytes is not None:
                            # Display the image
                            image = Image.open(io.BytesIO(image_bytes))
                            st.image(image, caption=f"Image {i+1}", width=300)
                            
                            classification, dpoc = predict_defect(processed_img, model)
                            classifications.append(classification['type'] if 'type' in classification else 'Unknown')
                            
                            # Generate OpenAI description
                            openai_desc = generate_openai_description(classification, image_bytes)
                            
                            image_data_list.append({
                                "url": url,
                                "defect_classification": classification,
                                "individual_dopc": dpoc,
                                "tags": [classification['type']] if 'type' in classification else ["Unknown"],
                                "openai_desc": openai_desc,
                            })
                            
                            # Display results
                            if 'type' in classification:
                                st.success(f"**Detected Defect:** {classification['type']}")
                                st.info(f"**Confidence:** {classification['severity']:.2%}")
                            else:
                                st.error(f"**Error:** {classification.get('error', 'Unknown error')}")
                    
                    st.markdown("---")

    with col2:
        st.header("📊 Analysis Summary")
        
        if classifications:
            # Display metrics
            st.markdown('<div class="metric-card">', unsafe_allow_html=True)
            st.metric("Total Images", len(classifications))
            st.markdown('</div>', unsafe_allow_html=True)
            
            st.markdown('<div class="metric-card">', unsafe_allow_html=True)
            st.metric("Unique Defects", len(set(classifications)))
            st.markdown('</div>', unsafe_allow_html=True)
            
            # Defect distribution
            if len(classifications) > 1:
                st.subheader("Defect Distribution")
                defect_counts = pd.Series(classifications).value_counts()
                st.bar_chart(defect_counts)
        
        # Display class labels
        st.subheader("🏷️ Supported Defect Types")
        for label in class_labels:
            st.write(f"• {label}")

    # Generate comprehensive report
    if image_data_list:
        st.markdown("---")
        st.header("📋 Detailed Analysis Report")
        
        # Generate SPS recommendations
        if classifications:
            with st.spinner("Generating recommended solutions..."):
                user_message = ", ".join(set(classifications))
                mandatory_sps, optional_sps = get_sps(user_message,type)
                
              

        # Display detailed results for each image
        for i, image_data in enumerate(image_data_list):
            with st.expander(f"📷 Image {i+1} - Detailed Analysis", expanded=True):
                col1, col2 = st.columns([1, 2])
                
                with col1:
                    if 'file_name' in image_data:
                        st.write(f"**File:** {image_data['file_name']}")
                    if 'url' in image_data:
                        st.write(f"**URL:** {image_data['url'][:50]}...")
                    
                    classification = image_data['defect_classification']
                    if 'type' in classification:
                        st.write(f"**Defect Type:** {classification['type']}")
                        st.write(f"**Severity:** {classification['severity']:.2%}")
                
                with col2:
                    st.write("**OpenAI Analysis:**")
                    st.text_area("", image_data['openai_desc'], height=150, key=f"desc_{i}")

        # Display SPS recommendations
        
        # Export results
        # Export results
        ai_response = {}
        boq_list = []
        ids = [sps["id"] for sps in optional_sps]
        sps_boq_map = {
            211:58,
            212:59,
            213:60,
            214:61,
            215:62,
            216:63,
            217:64,
            218:65,
            219:66,
            220:67,
            221:57,
            222:49
        }
        boq_list =  list([sps_boq_map[sps_id] for sps_id in ids if sps_id in sps_boq_map])

   
        ai_response = {
                "report_id": report_submission_id,
                "page_id": report_section_page_id,
                "section_id": report_section_id,
                "defect_type": defect_type,
                "images": image_data_list,
                "mandatory_sps": mandatory_sps,
                "optional_sps": optional_sps,
                "boq_list": boq_list
            }
        # if st.button("📤 Export Results as JSON", type="secondary"):
        #     ai_response = {
        #         "report_id": report_submission_id,
        #         "page_id": report_section_page_id,
        #         "section_id": report_section_id,
        #         "defect_type": defect_type,
        #         "images": image_data_list,
        #         "mandatory_sps": mandatory_sps,
        #         "optional_sps": optional_sps
        #     }
        st.subheader("📦 JSON Response Preview")
        with st.expander("View JSON Response", expanded=False):
            st.json(json.loads(json.dumps(ai_response, indent=2, ensure_ascii=False)))


        # 📥 Download button
        st.download_button(
            label="Download JSON Report",
            data=json.dumps(ai_response, indent=2),
            file_name=f"defect_report_{report_submission_id}_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json",
            mime="application/json"
        )

       

if __name__ == "__main__":
    main()