from fasthtml.common import *
from fasthtml.xtend import CheckboxX
from translator import languages as VALID_LANGUAGES, translate_to_languages
import tempfile
from starlette.requests import Request
from typing import List
from starlette.responses import FileResponse, PlainTextResponse
import os
import mimetypes
import zipfile
import io

css = Style(':root {--pico-font-size:90%,--pico-font-family: Pacifico, cursive;}')
app = FastHTML(hdrs=(picolink, css))

TEMP_FILES = {}
# Extract form generation to a separate function
def generate_language_form():
    return Form(
        Input(type='file', name='sbv_file', accept='.sbv', required=True),
        H3('Select languages for translation'),
        Fieldset(
            Legend("Languages"),
            Div(
                *[CheckboxX(
                    label=lang,
                    value=lang,
                    id=f'lang_{lang}',
                    name='languages'
                ) for lang in VALID_LANGUAGES],
                class_='grid'
            )
        ),
        Button('Translate', class_="primary"),
        action="/upload", method='post', enctype='multipart/form-data'
    )

@app.get('/')
def home():
    return Container(
        Article(
            H1("SBV File Translator"),
            P("Upload an SBV File and select languages for translation:"),
            generate_language_form(),
        )
    )

# Extract file processing logic to a separate function
def translate_and_process_sbv(content: bytes, filename: str, languages: List[str]) -> dict:
    """
    Translate and process an SBV file for the given languages.

    Args:
        content (bytes): The content of the uploaded SBV file.
        filename (str): The name of the uploaded file.
        languages (List[str]): List of target languages for translation.

    Returns:
        dict: A dictionary containing translation results, file paths, and zip filename.
    """
    try:
        temp_dir = tempfile.mkdtemp()
        input_file = os.path.join(temp_dir, filename)
        
        # Ensure the directory exists
        os.makedirs(os.path.dirname(input_file), exist_ok=True)
        
        with open(input_file, 'wb') as f:
            f.write(content)

        selected_languages = [lang for lang in languages if lang in VALID_LANGUAGES]
        output_dir = os.path.join(temp_dir, "translations")
        os.makedirs(output_dir, exist_ok=True)

        translate_to_languages(input_file, output_dir, selected_languages)

        return create_translation_files(temp_dir, output_dir, filename, selected_languages)
    except Exception as e:
        print(f"Error in translate_and_process_sbv: {str(e)}")
        raise  # Re-raise the exception after logging

# Extract translation file creation logic to a separate function
def create_translation_files(temp_dir: str, output_dir: str, filename: str, languages: List[str]) -> dict:
    """
    Create translation files and organize them for download.

    This function processes the translated SBV files, organizes them in a temporary directory,
    creates a zip file containing all translations, and prepares the data for the response.

    Args:
        temp_dir (str): Path to the temporary directory for storing files.
        output_dir (str): Path to the directory containing the translated SBV files.
        filename (str): Name of the original uploaded file.
        languages (List[str]): List of languages for which translations were created.

    Returns:
        dict: A dictionary containing:
            - 'translations': Dict of language codes to translated content.
            - 'file_paths': Dict of language codes to paths of individual translation files.
            - 'zip_filename': Name of the zip file containing all translations.
    """
    translations = {}
    file_paths = {}
    base_filename = os.path.splitext(filename)[0]

    for lang in languages:
        translated_file = os.path.join(output_dir, f"{base_filename}_{lang}.sbv")
        with open(translated_file, 'r', encoding='utf-8') as f:
            translations[lang] = f.read()
        
        new_filename = f"{base_filename}_{lang}.sbv"
        new_path = os.path.join(temp_dir, new_filename)
        os.rename(translated_file, new_path)
        file_paths[lang] = new_path
        TEMP_FILES[new_filename] = new_path

    zip_filename = f"{base_filename}_{'_'.join(languages)}.zip"
    zip_path = os.path.join(temp_dir, zip_filename)
    
    with zipfile.ZipFile(zip_path, 'w') as zip_file:
        for lang, file_path in file_paths.items():
            zip_file.write(file_path, os.path.basename(file_path))
    
    TEMP_FILES[zip_filename] = zip_path

    return {
        'translations': translations,
        'file_paths': file_paths,
        'zip_filename': zip_filename
    }

@app.post('/upload')
async def upload_file(request: Request, sbv_file: UploadFile):
    content = await sbv_file.read()
    
    # Get form data
    form = await request.form()
    languages = form.getlist('languages')

    # Ensure languages is always a list
    if not languages:
        return PlainTextResponse("Please select at least one language for translation.", status_code=400)

    result = translate_and_process_sbv(content, sbv_file.filename, languages)

    return Container(
        Article(
            H1('Translations'),
            A("Download All Translations", 
              href=f"/download/{result['zip_filename']}", 
              download=True,
              class_="primary",
              role="button"
            ),
            *[Card(
                H3(f'{lang.capitalize()} Translation'),
                Pre(
                    '\n'.join(result['translations'][lang].split('\n')[:9]) + '\n...',
                ),
                Footer(
                    A(f"Download {lang.capitalize()} Translation", 
                      href=f"/download/{os.path.basename(result['file_paths'][lang])}", 
                      download=True, 
                      class_="secondary outline",
                      role="button"
                    )
                )
            ) for lang in result['translations']],
            class_="container-fluid"  # Add this class for full-width content
        )
    )

@app.get("/download/{filename}")
def get(filename: str):
    file_path = TEMP_FILES.get(filename)
    if file_path and os.path.exists(file_path):
        mime_type, _ = mimetypes.guess_type(file_path)
        return FileResponse(
            file_path,
            filename=filename,
            media_type=mime_type or "application/octet-stream",
            headers={
                "Content-Disposition": f"attachment; filename={filename}",
                "X-Content-Type-Options": "nosniff",
            }
        )
    else:
        return PlainTextResponse("File not found", status_code=404)

serve()