import streamlit as st
import requests
import os
import sys
from PIL import Image
import io
import time
from pathlib import Path

# Set API URL
API_URL = "http://localhost:8001"  # Local FastAPI server URL

st.set_page_config(
    page_title="NAFNet Image Deblurring",
    page_icon="🔍",
    layout="wide",
)

st.title("NAFNet Image Deblurring Application")
st.markdown("""
Transform your blurry photos into clear, sharp images using the state-of-the-art NAFNet AI model.
Upload an image to get started!
""")

# File uploader
uploaded_file = st.file_uploader(
    "Choose a blurry image...", type=["jpg", "jpeg", "png", "bmp"])

# Sidebar controls
with st.sidebar:
    st.header("About NAFNet")
    
    st.markdown("""
    **NAFNet** (Nonlinear Activation Free Network) is a state-of-the-art image restoration model designed for tasks like deblurring.
    
    Key features:
    - High-quality image deblurring
    - Fast processing time
    - Preservation of image details
    """)
    
    st.markdown("---")
    
    # Check API status
    if st.button("Check API Status"):
        try:
            response = requests.get(f"{API_URL}/status/", timeout=5)
            if response.status_code == 200 and response.json().get("status") == "ok":
                st.success("✅ API is running and ready")
                
                # Display additional info if available
                memory_info = response.json().get("memory", {})
                if memory_info:
                    st.info(f"CUDA Memory: {memory_info.get('cuda_memory_allocated', 'N/A')}")
            else:
                st.error("❌ API is not responding properly")
        except:
            st.error("❌ Cannot connect to API")

# Process when upload is ready
if uploaded_file is not None:
    # Display the original image
    col1, col2 = st.columns(2)

    with col1:
        st.subheader("Original Image")
        image = Image.open(uploaded_file)
        st.image(image, use_container_width=True)

    # Process image button
    process_button = st.button("Deblur Image")

    if process_button:
        with st.spinner("Deblurring your image... Please wait."):
            try:
                # Prepare simplified file structure
                files = {
                    "file": ("image.jpg", uploaded_file.getvalue(), "image/jpeg")
                }
                
                # Send request to API
                response = requests.post(f"{API_URL}/deblur/", files=files, timeout=60)

                if response.status_code == 200:
                    with col2:
                        st.subheader("Deblurred Result")
                        deblurred_img = Image.open(io.BytesIO(response.content))
                        st.image(deblurred_img, use_column_width=True)

                    # Option to download the deblurred image
                    st.download_button(
                        label="Download Deblurred Image",
                        data=response.content,
                        file_name=f"deblurred_{uploaded_file.name}",
                        mime="image/png"
                    )
                else:
                    try:
                        error_details = response.json().get('detail', 'Unknown error')
                    except:
                        error_details = response.text
                    
                    st.error(f"Error: {error_details}")
            except Exception as e:
                st.error(f"An error occurred: {str(e)}")

# Footer
st.markdown("---")
st.markdown("Powered by NAFNet - Image Restoration Project")

def main():
    pass  # Streamlit already runs the script from top to bottom

if __name__ == "__main__":
    main()