import streamlit as st
import os
from PIL import Image
import torch
from manga_translator.translator import MangaTextDetector, process_manga_pages
import tempfile
import cv2
import io

# Initialize session state
if 'processed_results' not in st.session_state:
    st.session_state.processed_results = {}
if 'temp_dir' not in st.session_state:
    st.session_state.temp_dir = tempfile.mkdtemp()

# Set page config for wider layout and title
st.set_page_config(
    page_title="Manga Translator",
    page_icon="🎯",
    layout="wide",
    initial_sidebar_state="auto"
)

# Custom CSS for centering and background
st.markdown("""
    <style>
        /* Background Gradient */
        body {
            background: linear-gradient(135deg, #a18cd1 0%, #fbc2eb 100%) !important;
        }

        /* Main content container */
        .main .block-container {
            max-width: 900px;
            margin: 0 auto;
            padding-top: 2rem;
            padding-bottom: 2rem;
        }

        /* Custom box style */
        .centered-box {
            background: white;
            border-radius: 16px;
            box-shadow: 0 4px 32px rgba(0,0,0,0.08);
            padding: 2rem;
            margin-bottom: 2rem;
        }

        /* Button stretch */
        .stButton>button {
            width: 100%;
        }

        /* Alert style */
        .stAlert {
            background-color: rgb(255, 251, 235);
            border: none;
            padding: 1rem;
            margin: 1rem 0;
        }
    </style>
""", unsafe_allow_html=True)



st.title("🦀 MangaFlow: AI Manga Translator")
st.write("Upload a manga page to translate Japanese text to English. Powered by YOLO, Manga OCR, and DeepL.")

st.warning("**Note:** Translation accuracy may vary due to the complexity of Japanese text and manga-specific expressions. We're continuously working to improve the system!")

with st.expander("📋 Guidelines for Best Results", expanded=False):
    st.markdown("1. **Image Requirements:**")
    st.markdown("""
       - Clear, high-resolution manga page
       - Japanese text should be clearly visible
       - Text bubbles should not be cropped
       - Supported formats: PNG, JPG, JPEG
    """)
    st.markdown("2. **For Best Results:**")
    st.markdown("""
       - Avoid pages with handwritten text
       - Ensure text bubbles are not overlapping
       - Image should be properly oriented
       - Avoid heavily compressed images
    """)
    st.markdown("3. **Privacy & Copyright:**")
    st.markdown("""
       - Only upload content you have rights to use
       - We don't store any uploaded images
       - All processing is done in real-time
    """)

with st.expander("🔍 How It Works", expanded=False):
    st.markdown("""
    1. **Text Detection:** Custom YOLO model detects text bubbles
    2. **OCR Processing:** Extracts Japanese text
    3. **Translation:** Converts to English using DeepL API
    4. **Text Insertion:** Places translated text back into the image
    """)

def process_image(uploaded_file):
    """Process image and store results in session state."""
    if uploaded_file.name not in st.session_state.processed_results:
        try:
            detector = MangaTextDetector('best.pt')
            # Save temporary file
            temp_path = os.path.join(st.session_state.temp_dir, uploaded_file.name)
            with open(temp_path, "wb") as f:
                f.write(uploaded_file.getbuffer())
            # Process image
            image, detections, result_image, processed_regions, translated_image = detector.process_image(temp_path)
            # Store all results in session state
            st.session_state.processed_results[uploaded_file.name] = {
                'image': image,
                'detections': detections,
                'result_image': result_image,
                'processed_regions': processed_regions,
                'translated_image': translated_image
            }
            return True
        except Exception as e:
            st.error(f"❌ Error: {str(e)}")
            return False
    return True

uploaded_files = st.file_uploader(
    "Choose manga pages",
    type=['png', 'jpg', 'jpeg'],
    accept_multiple_files=True,
    help="Drag and drop your manga images here. Supported formats: PNG, JPG, JPEG"
)

if uploaded_files:
    with tempfile.TemporaryDirectory() as temp_dir:
        for uploaded_file in uploaded_files:
            file_path = os.path.join(temp_dir, uploaded_file.name)
            with open(file_path, "wb") as f:
                f.write(uploaded_file.getbuffer())
        detector = MangaTextDetector('best.pt')
        for uploaded_file in uploaded_files:
            st.subheader(f"Processing: {uploaded_file.name}")
            col1, col2 = st.columns(2)
            image_path = os.path.join(temp_dir, uploaded_file.name)
            original_image = Image.open(image_path)
            with col1:
                st.markdown("**Original Image**")
                st.image(original_image, use_column_width=True)
            try:
                image, detections, result_image, processed_regions, translated_image = detector.process_image(image_path)
                with col2:
                    if translated_image is not None:
                        st.markdown("**Translated Image**")
                        st.image(translated_image, use_column_width=True)
                    else:
                        st.error("No text was detected in this image.")
                if processed_regions and processed_regions['text_regions']:
                    with st.expander("View Detected Text and Translations"):
                        for i, region in enumerate(processed_regions['text_regions'], 1):
                            st.markdown(f"**Region {i}**")
                            cols = st.columns(2)
                            with cols[0]:
                                st.markdown("Original Text:")
                                st.code(region['text'])
                            with cols[1]:
                                st.markdown("Translation:")
                                st.code(region['translation'])
                            st.markdown("---")
            except Exception as e:
                st.error(f"Error processing {uploaded_file.name}: {str(e)}")
                continue
            st.markdown("---")
else:
    st.info("👆 Upload a manga page to get started!")

st.markdown("""
    <div style='text-align: center; color: #666666; padding: 1rem;'>
        Made with ❤️ for manga fans
    </div>
""", unsafe_allow_html=True)

st.markdown("---")
st.markdown("""
*Created by Ebhon*

This app translates manga text from Japanese to English using:
- YOLO for text detection
- Manga OCR for Japanese text recognition
- DeepL for translation
""")

st.markdown('</div>', unsafe_allow_html=True)
st.markdown('</div>', unsafe_allow_html=True)