Py.Cafe

"""Main application entry point."""
import streamlit as st
import os
import logging
from openai import AzureOpenAI
from dotenv import load_dotenv
from translation_prompts import TECHNICAL_GLOSSARY, LANGUAGE_CODES
from handlers.format_router import FormatRouter, MIME_TYPES
from handlers.cache_manager import CacheManager
from handlers.translation import validate_azure_config
from handlers.translation_errors import (
    TranslationError,
    AzureDeploymentError,
    ModelNotFoundError,
    APIKeyError,
    EndpointError,
    FormatError,
    ContentTooLongError
)

# Set page config
st.set_page_config(page_title='Kotodama (言霊) 🌐', page_icon='./images/tm-logo.png', layout='wide')

# Load environment variables
load_dotenv()

# Configure logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

try:
    # Validate Azure OpenAI configuration
    validate_azure_config()
    
    # Initialize Azure OpenAI client
    client = AzureOpenAI(
        api_key=os.getenv("AZURE_OPENAI_API_KEY"),
        api_version=os.getenv("AZURE_OPENAI_API_VERSION", "2024-02-15-preview"),
        azure_endpoint=os.getenv("AZURE_OPENAI_ENDPOINT")
    )
except TranslationError as e:
    st.error(str(e))
    st.stop()

# Setup directories
TRANSLATED_FILES_DIR = os.path.join(os.getcwd(), 'translated_files')
CACHE_DIR = os.path.join(os.getcwd(), 'cache')
os.makedirs(TRANSLATED_FILES_DIR, exist_ok=True)
os.makedirs(CACHE_DIR, exist_ok=True)

def main():
    st.image("./images/trend-logo.png", width=400)
    st.title("Kotodama (言霊) 🌐")
    st.subheader("AI-powered translation tool for technical content")
    st.divider()
    st.write(f"Upload your file for translation ({', '.join(MIME_TYPES.keys())})")
    
    # Initialize session state
    if 'translation_stopped' not in st.session_state:
        st.session_state.translation_stopped = False
    if 'translation_in_progress' not in st.session_state:
        st.session_state.translation_in_progress = False
    if 'cache_manager' not in st.session_state:
        st.session_state.cache_manager = CacheManager(CACHE_DIR)
    
    # Cache status and clear button in sidebar
    with st.sidebar:
        st.subheader("Cache Status")
        stats = st.session_state.cache_manager.get_stats()
        st.write(f"Hits: {stats['hits']}")
        st.write(f"Misses: {stats['misses']}")
        st.write(f"Total entries: {stats['total_entries']}")
        
        if st.button("Clear Cache"):
            st.session_state.cache_manager.clear()
            st.success("Cache cleared! 🧹")
            st.rerun()
    
    # File uploader with supported extensions
    uploaded_file = st.file_uploader("Choose a file", type=list(MIME_TYPES.keys()))
    
    # Language selector
    available_languages = sorted(list(LANGUAGE_CODES.keys()))
    default_idx = available_languages.index("Japanese")
    language = st.selectbox(
        "Select target language",
        available_languages,
        index=default_idx
    )
    
    # Button container
    button_container = st.container()

    # Translate button
    translate_disabled = st.session_state.translation_in_progress
    translate_button = button_container.button(
        "Translate",
        type="primary",
        disabled=translate_disabled,
        key="translate_button"
    )

    if translate_button:
        if uploaded_file is None:
            st.error("Please upload a file first!")
            st.stop()
        
        file_extension = uploaded_file.name.split('.')[-1].lower()
        st.session_state.translation_stopped = False
        st.session_state.translation_in_progress = True
        
        # Create status container for progress
        with st.status("Translating...", expanded=True) as status:
            try:
                # Initialize format router with OpenAI client
                router = FormatRouter(
                    st.session_state.cache_manager,
                    TRANSLATED_FILES_DIR,
                    client
                )
                
                # Track initial cache stats
                initial_stats = st.session_state.cache_manager.get_stats()
                
                # Create progress elements and stop button in columns
                col1, col2 = status.columns([3, 1])
                with col1:
                    progress_text = st.text("Initializing...")
                    progress_bar = st.progress(0.0)
                with col2:
                    if st.button("Stop", type="secondary", key="stop_button"):
                        st.session_state.translation_stopped = True
                        st.warning("Translation stopped by user")
                        st.session_state.translation_in_progress = False
                        st.rerun()
                
                # Create a progress callback that updates both bar and text
                def update_progress(value: float) -> None:
                    if st.session_state.translation_stopped:
                        status.update(label="Translation stopped", state="error")
                        st.stop()
                    # Ensure value is between 0 and 1
                    value = max(0.0, min(1.0, float(value)))
                    progress_bar.progress(value)
                    percent = int(value * 100)
                    progress_text.text(f"Translating... {percent}%")
                
                # Translate file
                translated_content, mime_type = router.translate_file(
                    uploaded_file,
                    file_extension,
                    language,
                    progress_callback=update_progress 
                )
                
                # Update cache stats
                final_stats = st.session_state.cache_manager.get_stats()
                new_hits = final_stats['hits'] - initial_stats['hits']
                new_misses = final_stats['misses'] - initial_stats['misses']
                
                # Update status with results
                status.update(label="Translation completed! 🎉", state="complete")
                
                if new_hits > 0:
                    st.info(f"✨ Used {new_hits} cached translations!")
                if new_misses > 0:
                    st.info(f"🆕 Generated {new_misses} new translations!")
                
                # Save translated file
                translated_file_path = os.path.join(TRANSLATED_FILES_DIR, f'translated.{file_extension}')
                
                # Handle binary vs text content
                if isinstance(translated_content, bytes):
                    write_mode = 'wb'
                    file_content = translated_content
                else:
                    write_mode = 'w'
                    file_content = str(translated_content)
                    if not isinstance(file_content, bytes):
                        file_content = file_content.encode('utf-8').decode('utf-8')
                    
                with open(translated_file_path, write_mode, encoding='utf-8' if write_mode == 'w' else None) as f:
                    f.write(file_content)
                
                # Center the download button
                col1, col2, col3 = st.columns([1, 2, 1])
                with col2:
                    # Ensure content is properly encoded for download
                    if isinstance(file_content, str):
                        download_content = file_content.encode('utf-8')
                    else:
                        download_content = file_content
                        
                    st.download_button(
                        label=f"Download translated {file_extension.upper()}",
                        data=download_content,
                        file_name=f"translated.{file_extension}",
                        mime=mime_type,
                        use_container_width=True
                    )
                
                st.session_state.translation_in_progress = False
            
            except AzureDeploymentError as e:
                st.session_state.translation_in_progress = False
                status.update(label="Azure OpenAI Service Error", state="error")
                st.error(str(e))
                st.info("Please check your Azure OpenAI deployment and model configuration.")
                
            except ModelNotFoundError as e:
                st.session_state.translation_in_progress = False
                status.update(label="Model Configuration Error", state="error")
                st.error(str(e))
                st.info("Make sure the model is deployed in your Azure OpenAI resource.")
                
            except APIKeyError as e:
                st.session_state.translation_in_progress = False
                status.update(label="Authentication Error", state="error")
                st.error(str(e))
                st.info("Check your .env file and Azure OpenAI API key.")
                
            except EndpointError as e:
                st.session_state.translation_in_progress = False
                status.update(label="Configuration Error", state="error")
                st.error(str(e))
                st.info("Check your .env file and Azure OpenAI endpoint URL.")
                
            except FormatError as e:
                st.session_state.translation_in_progress = False
                status.update(label="File Format Error", state="error")
                st.error(str(e))
                st.info("Please check your file format and try again.")
                
            except ContentTooLongError as e:
                st.session_state.translation_in_progress = False
                status.update(label="Content Length Error", state="error")
                st.error(str(e))
                st.info("Try splitting your file into smaller parts.")
                
            except Exception as e:
                st.session_state.translation_in_progress = False
                status.update(label="An unexpected error occurred", state="error")
                st.error(str(e))
                logging.error(f"Translation error: {str(e)}")
                st.info("Please try again or contact support if the issue persists.")
    st.divider()
    st.write("Made with ❤️ by [Fernando Tucci](mailto:fernando_tucci@trendmicro.com)")
if __name__ == '__main__':
    main()