main.py

import os
import logging
import shutil
import re
import zipfile
import uvicorn
import asyncio
import threading
import time
from fastapi import FastAPI, Request, BackgroundTasks, HTTPException
from fastapi.responses import HTMLResponse, StreamingResponse, JSONResponse
from fastapi.staticfiles import StaticFiles
from fastapi.templating import Jinja2Templates
from pydantic import BaseModel
from typing import Dict, Set, Tuple, List, Optional
from io import BytesIO
from pathlib import Path

# Configure data directories for Hugging Face Spaces
# HF Spaces can have persistent storage in the /data directory
DATA_DIR = '/data'
# Import the existing telegram bot functionality
# Define these directories here to ensure they're created before importing app.py
PENDING_DIR = os.path.join(DATA_DIR, 'pending')
PENDING_VOICE_DIR = os.path.join(PENDING_DIR, 'voice')
PENDING_TEXT_DIR = os.path.join(PENDING_DIR, 'text')
PROCESSED_DIR = os.path.join(DATA_DIR, 'processed')
PROCESSED_VOICE_DIR = os.path.join(PROCESSED_DIR, 'voice')
PROCESSED_TEXT_DIR = os.path.join(PROCESSED_DIR, 'text')

# Create directories if they don't exist
os.makedirs(PENDING_VOICE_DIR, exist_ok=True)
os.makedirs(PENDING_TEXT_DIR, exist_ok=True)
os.makedirs(PROCESSED_VOICE_DIR, exist_ok=True)
os.makedirs(PROCESSED_TEXT_DIR, exist_ok=True)

# Import from app.py now that we've defined the directories
try:
    from app import extract_sequence_number, get_files_by_sequence_number
except Exception as e:
    # Define fallback functions if import fails
    def extract_sequence_number(filename: str) -> int:
        """Extract sequence number from filename, e.g., 'Sound 100.wav' -> 100."""
        match = re.search(r'(\d+)', filename)
        if match:
            return int(match.group(1))
        return None

    def get_files_by_sequence_number():
        """Create dictionaries mapping sequence numbers to filenames."""
        voice_files = {}
        text_files = {}

        # Map voice files to sequence numbers
        for filename in os.listdir(PENDING_VOICE_DIR):
            if filename.endswith('.wav'):
                seq_num = extract_sequence_number(filename)
                if seq_num is not None:
                    voice_files[seq_num] = filename

        # Map text files to sequence numbers
        for filename in os.listdir(PENDING_TEXT_DIR):
            if filename.endswith('.txt'):
                seq_num = extract_sequence_number(filename)
                if seq_num is not None:
                    text_files[seq_num] = filename

        return voice_files, text_files

    logging.error(f"Error importing from app.py: {str(e)}")

# Enable logging
logging.basicConfig(
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    level=logging.INFO
)
logger = logging.getLogger(__name__)

app = FastAPI(title="Algerian Darija Transcription Service")

# Create templates instance
templates = Jinja2Templates(directory="templates")

# Telegram bot process
bot_process = None
bot_running = False
bot_start_time = None

# Function to start the bot in a separate process instead of a thread


def start_bot_process():
    from multiprocessing import Process
    global bot_process, bot_running

    # Create a new process for the bot
    bot_process = Process(target=run_bot_process)
    bot_process.daemon = True  # Process will terminate when main process ends
    bot_process.start()
    bot_running = True
    return bot_process.pid


def run_bot_process():
    """Function that runs in a separate process to start the bot"""
    import sys
    import os
    import asyncio

    try:
        # Import and run the main function from app.py
        from app import main

        # Create a new event loop for this process
        loop = asyncio.new_event_loop()
        asyncio.set_event_loop(loop)

        # Run the bot
        loop.run_until_complete(main())
        loop.run_forever()
    except Exception as e:
        sys.stderr.write(f"Error in bot process: {str(e)}\n")
        sys.exit(1)


@app.get("/", response_class=HTMLResponse)
async def get_root(request: Request):
    return templates.TemplateResponse("index.html", {"request": request, "bot_running": bot_running})


@app.post("/start-bot")
async def start_bot_endpoint(background_tasks: BackgroundTasks):
    global bot_running, bot_start_time

    if bot_running:
        return {"status": "error", "message": "Bot is already running"}

    # Check if TELEGRAM_TOKEN is set
    if not os.environ.get("TELEGRAM_TOKEN"):
        return {"status": "error", "message": "TELEGRAM_TOKEN environment variable is not set. Please configure it in Hugging Face Space settings."}

    try:
        # Start bot in a separate process
        pid = start_bot_process()
        bot_start_time = time.time()
        return {"status": "success", "message": f"Bot started successfully with PID {pid}"}
    except Exception as e:
        logger.error(f"Error starting bot: {str(e)}")
        return {"status": "error", "message": f"Failed to start bot: {str(e)}"}


@app.post("/stop-bot")
async def stop_bot_endpoint():
    global bot_process, bot_running, bot_start_time

    if not bot_running or bot_process is None:
        return {"status": "error", "message": "Bot is not running"}

    try:
        # Terminate the process
        bot_process.terminate()
        bot_process.join(timeout=5)  # Wait for process to terminate

        # If process didn't terminate, force kill it
        if bot_process.is_alive():
            bot_process.kill()
            bot_process.join()

        bot_running = False
        bot_start_time = None
        return {"status": "success", "message": "Bot stopped successfully"}
    except Exception as e:
        logger.error(f"Error stopping bot: {str(e)}")
        return {"status": "error", "message": f"Failed to stop bot: {str(e)}"}


@app.get("/status")
async def get_status():
    # Check if directories exist and create if needed
    os.makedirs(PENDING_VOICE_DIR, exist_ok=True)
    os.makedirs(PENDING_TEXT_DIR, exist_ok=True)
    os.makedirs(PROCESSED_VOICE_DIR, exist_ok=True)
    os.makedirs(PROCESSED_TEXT_DIR, exist_ok=True)

    # Count pending files
    pending_voice_count = len(
        [f for f in os.listdir(PENDING_VOICE_DIR) if f.endswith('.wav')])
    pending_text_count = len(
        [f for f in os.listdir(PENDING_TEXT_DIR) if f.endswith('.txt')])

    # Get sequence numbers
    voice_files, text_files = get_files_by_sequence_number()
    matching_pairs = len(
        set(voice_files.keys()).intersection(set(text_files.keys())))

    # Count processed files - Only count text files now as voice files are no longer saved to processed
    processed_voice_count = 0  # Set to 0 since we no longer save voice files
    processed_text_count = 0

    # Count processed files by user
    user_stats = {}

    # Process text files only (voice files are no longer saved to processed directory)
    for root, dirs, files in os.walk(PROCESSED_TEXT_DIR):
        txt_files = [f for f in files if f.endswith('.txt')]
        processed_text_count += len(txt_files)

        # Get user ID from directory path
        if os.path.basename(root).startswith("user_"):
            user_id = os.path.basename(root)
            if user_id not in user_stats:
                user_stats[user_id] = {"voice": 0, "text": 0}
            user_stats[user_id]["text"] += len(txt_files)

    # Calculate uptime if bot is running
    uptime = None
    if bot_running and bot_start_time:
        uptime = int(time.time() - bot_start_time)

    return {
        "bot_running": bot_running,
        "uptime_seconds": uptime,
        "pending_voice_count": pending_voice_count,
        "pending_text_count": pending_text_count,
        "matching_pairs": matching_pairs,
        "processed_voice_count": processed_voice_count,
        "processed_text_count": processed_text_count,
        "user_stats": user_stats,
        "is_huggingface": True
    }


@app.get("/list-users")
async def list_users():
    users = set()

    # Get users from text directories only (since voice files are no longer saved)
    for item in os.listdir(PROCESSED_TEXT_DIR):
        if item.startswith("user_"):
            users.add(item)

    return {"users": sorted(list(users))}


@app.get("/download/{user_id}")
async def download_processed_data(user_id: str):
    if not user_id.startswith("user_"):
        raise HTTPException(status_code=400, detail="Invalid user ID format")

    # Check if user exists - only check text directory now
    text_dir = os.path.join(PROCESSED_TEXT_DIR, user_id)

    if not os.path.exists(text_dir):
        raise HTTPException(
            status_code=404, detail=f"No data found for user {user_id}")

    # Create a zip file in memory
    zip_buffer = BytesIO()
    with zipfile.ZipFile(zip_buffer, "w", zipfile.ZIP_DEFLATED) as zip_file:
        # Add text files only
        if os.path.exists(text_dir):
            for file_name in os.listdir(text_dir):
                file_path = os.path.join(text_dir, file_name)
                if os.path.isfile(file_path) and file_name.endswith('.txt'):
                    # Read the file and add it to the zip
                    with open(file_path, 'rb') as f:
                        zip_file.writestr(f"text/{file_name}", f.read())

    # Reset buffer position
    zip_buffer.seek(0)

    # Return the zip file as a response
    return StreamingResponse(
        zip_buffer,
        media_type="application/zip",
        headers={
            "Content-Disposition": f"attachment; filename={user_id}_processed_data.zip"
        }
    )


@app.get("/download-all")
async def download_all_processed_data():
    # Create a zip file in memory
    zip_buffer = BytesIO()
    with zipfile.ZipFile(zip_buffer, "w", zipfile.ZIP_DEFLATED) as zip_file:
        # Add all processed text files only
        for root, dirs, files in os.walk(PROCESSED_TEXT_DIR):
            for file_name in files:
                if file_name.endswith('.txt'):
                    file_path = os.path.join(root, file_name)
                    # Get the relative path from the PROCESSED_TEXT_DIR
                    rel_path = os.path.relpath(file_path, PROCESSED_TEXT_DIR)
                    # Read the file and add it to the zip
                    with open(file_path, 'rb') as f:
                        zip_file.writestr(f"text/{rel_path}", f.read())

    # Reset buffer position
    zip_buffer.seek(0)

    # Return the zip file as a response
    return StreamingResponse(
        zip_buffer,
        media_type="application/zip",
        headers={
            "Content-Disposition": "attachment; filename=all_processed_data.zip"
        }
    )


@app.on_event("startup")
async def setup_app():
    # Ensure templates directory exists
    os.makedirs("templates", exist_ok=True)

    # Only create index.html if it doesn't exist
    if not os.path.exists("templates/index.html"):
        # Create a simple default template if needed
        with open("templates/index.html", "w") as f:
            f.write("""
<!DOCTYPE html>
<html>
<head>
    <title>Algerian Darija Transcription Service</title>
    <meta charset="UTF-8">
    <meta name="viewport" content="width=device-width, initial-scale=1.0">
    <style>
        body {
            font-family: Arial, sans-serif;
            max-width: 800px;
            margin: 0 auto;
            padding: 20px;
            line-height: 1.6;
        }
        .container {
            border: 1px solid #ddd;
            border-radius: 8px;
            padding: 20px;
            margin-bottom: 20px;
        }
        button {
            background-color: #4CAF50;
            color: white;
            padding: 8px 16px;
            border: none;
            border-radius: 4px;
            cursor: pointer;
            margin-right: 10px;
        }
        button.stop {
            background-color: #f44336;
        }
        button:disabled {
            background-color: #cccccc;
            cursor: not-allowed;
        }
        .status {
            margin-top: 20px;
        }
        table {
            width: 100%;
            border-collapse: collapse;
            margin-top: 10px;
        }
        th, td {
            border: 1px solid #ddd;
            padding: 8px;
            text-align: left;
        }
        th {
            background-color: #f2f2f2;
        }
        .users-container {
            margin-top: 20px;
        }
        .user-item {
            display: flex;
            justify-content: space-between;
            padding: 8px;
            border-bottom: 1px solid #eee;
        }
        .note {
            background-color: #f0f8ff;
            border-left: 4px solid #1e90ff;
            padding: 10px;
            margin: 10px 0;
            font-style: italic;
        }
    </style>
</head>
<body>
    <h1>Algerian Darija Transcription Service</h1>
    
    <div class="container">
        <h2>Bot Control</h2>
        <p>Start or stop the Telegram bot service.</p>
        <button id="startBot" onclick="startBot()">Start Bot</button>
        <button id="stopBot" class="stop" onclick="stopBot()" disabled>Stop Bot</button>
        <div id="botMessage" class="status"></div>
    </div>
    
    <div class="container">
        <h2>Status</h2>
        <div class="note">
            Note: Only corrected text files are saved to the processed directory. Voice files are removed after processing to save storage space.
        </div>
        <div id="statusInfo" class="status">Loading status...</div>
    </div>
    
    <div class="container">
        <h2>Users Data</h2>
        <div id="users" class="users-container">Loading users...</div>
    </div>
    
    <div class="container">
        <h2>Data Download</h2>
        <p>Download all processed corrected text files or data for specific users.</p>
        <button onclick="downloadAll()">Download All Text Data</button>
    </div>

    <script>
        // Check status when page loads
        document.addEventListener('DOMContentLoaded', function() {
            updateStatus();
            updateUsers();
            setInterval(updateStatus, 10000); // Update status every 10 seconds
        });
        
        function updateStatus() {
            fetch('/status')
                .then(response => response.json())
                .then(data => {
                    console.log('Status data:', data);
                    
                    // Update bot status
                    const botRunning = data.bot_running;
                    document.getElementById('startBot').disabled = botRunning;
                    document.getElementById('stopBot').disabled = !botRunning;
                    
                    // Format uptime
                    let uptimeDisplay = 'Not running';
                    if (data.uptime_seconds) {
                        const hours = Math.floor(data.uptime_seconds / 3600);
                        const minutes = Math.floor((data.uptime_seconds % 3600) / 60);
                        const seconds = data.uptime_seconds % 60;
                        uptimeDisplay = `${hours}h ${minutes}m ${seconds}s`;
                    }
                    
                    // Update status table
                    const statusHtml = `
                        <table>
                            <tr>
                                <th>Metric</th>
                                <th>Value</th>
                            </tr>
                            <tr>
                                <td>Bot Status</td>
                                <td>${botRunning ? '✅ Running' : '❌ Stopped'}</td>
                            </tr>
                            <tr>
                                <td>Uptime</td>
                                <td>${uptimeDisplay}</td>
                            </tr>
                            <tr>
                                <td>Pending Voice Files</td>
                                <td>${data.pending_voice_count}</td>
                            </tr>
                            <tr>
                                <td>Pending Text Files</td>
                                <td>${data.pending_text_count}</td>
                            </tr>
                            <tr>
                                <td>Matching Pending Pairs</td>
                                <td>${data.matching_pairs}</td>
                            </tr>
                            <tr>
                                <td>Processed Text Files</td>
                                <td>${data.processed_text_count}</td>
                            </tr>
                        </table>
                    `;
                    
                    document.getElementById('statusInfo').innerHTML = statusHtml;
                })
                .catch(error => {
                    console.error('Error fetching status:', error);
                    document.getElementById('statusInfo').innerHTML = 'Error fetching status';
                });
        }
        
        function updateUsers() {
            fetch('/list-users')
                .then(response => response.json())
                .then(data => {
                    console.log('Users data:', data);
                    const users = data.users;
                    
                    if (users.length === 0) {
                        document.getElementById('users').innerHTML = 'No users found.';
                        return;
                    }
                    
                    let usersHtml = '';
                    users.forEach(user => {
                        usersHtml += `
                            <div class="user-item">
                                <span>${user}</span>
                                <button onclick="downloadUser('${user}')">Download Text Data</button>
                            </div>
                        `;
                    });
                    
                    document.getElementById('users').innerHTML = usersHtml;
                })
                .catch(error => {
                    console.error('Error fetching users:', error);
                    document.getElementById('users').innerHTML = 'Error fetching users';
                });
        }
        
        function startBot() {
            document.getElementById('botMessage').innerHTML = 'Starting bot...';
            
            fetch('/start-bot', {
                method: 'POST',
            })
                .then(response => response.json())
                .then(data => {
                    document.getElementById('botMessage').innerHTML = data.message;
                    updateStatus();
                })
                .catch(error => {
                    console.error('Error starting bot:', error);
                    document.getElementById('botMessage').innerHTML = 'Error starting bot: ' + error;
                });
        }
        
        function stopBot() {
            document.getElementById('botMessage').innerHTML = 'Stopping bot...';
            
            fetch('/stop-bot', {
                method: 'POST',
            })
                .then(response => response.json())
                .then(data => {
                    document.getElementById('botMessage').innerHTML = data.message;
                    updateStatus();
                })
                .catch(error => {
                    console.error('Error stopping bot:', error);
                    document.getElementById('botMessage').innerHTML = 'Error stopping bot: ' + error;
                });
        }
        
        function downloadUser(userId) {
            window.location.href = `/download/${userId}`;
        }
        
        function downloadAll() {
            window.location.href = '/download-all';
        }
    </script>
</body>
</html>
            """)

# Entrypoint for running the server
if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=7860)