vms/tabs/manage_tab.py

"""
Manage tab for Video Model Studio UI
"""

import gradio as gr
import logging
import shutil
from pathlib import Path
from typing import Dict, Any, List, Optional

from .base_tab import BaseTab
from ..config import (
    HF_API_TOKEN, VIDEOS_TO_SPLIT_PATH, STAGING_PATH, TRAINING_VIDEOS_PATH, 
    TRAINING_PATH, MODEL_PATH, OUTPUT_PATH, LOG_FILE_PATH
)
from ..utils import validate_model_repo

logger = logging.getLogger(__name__)

class ManageTab(BaseTab):
    """Manage tab for storage management and model publication"""
    
    def __init__(self, app_state):
        super().__init__(app_state)
        self.id = "manage_tab"
        self.title = "5️⃣  Manage"
    
    def create(self, parent=None) -> gr.TabItem:
        """Create the Manage tab UI components"""
        with gr.TabItem(self.title, id=self.id) as tab:
            with gr.Column():
                with gr.Row():
                    with gr.Column():
                        gr.Markdown("## Publishing")
                        gr.Markdown("You model can be pushed to Hugging Face (this will use HF_API_TOKEN)")

                        with gr.Row():
                            with gr.Column():
                                self.components["repo_id"] = gr.Textbox(
                                    label="HuggingFace Model Repository",
                                    placeholder="username/model-name",
                                    info="The repository will be created if it doesn't exist"
                                )
                                self.components["make_public"] = gr.Checkbox(
                                    label="Check this to make your model public (ie. visible and downloadable by anyone)",
                                    info="You model is private by default"
                                )
                                self.components["push_model_btn"] = gr.Button(
                                    "Push my model"
                                )

                with gr.Row():
                    with gr.Column():
                        with gr.Row():
                            with gr.Column():
                                gr.Markdown("## Storage management")
                                with gr.Row():
                                    self.components["download_dataset_btn"] = gr.DownloadButton(
                                        "Download dataset (click again if DL doesn't start)",
                                        variant="secondary",
                                        size="lg"
                                    )
                                    self.components["download_model_btn"] = gr.DownloadButton(
                                        "Download model (click again if DL doesn't start)",
                                        variant="secondary",
                                        size="lg"
                                    )

                        with gr.Row():
                            self.components["global_stop_btn"] = gr.Button(
                                "Stop everything and delete my data",
                                variant="stop"
                            )
                            self.components["global_status"] = gr.Textbox(
                                label="Global Status",
                                interactive=False,
                                visible=False
                            )
        
        return tab
    
    def connect_events(self) -> None:
        """Connect event handlers to UI components"""
        # Repository ID validation
        self.components["repo_id"].change(
            fn=self.validate_repo,
            inputs=[self.components["repo_id"]],
            outputs=[self.components["repo_id"]]
        )
        
        # Download buttons
        self.components["download_dataset_btn"].click(
            fn=self.app.trainer.create_training_dataset_zip,
            outputs=[self.components["download_dataset_btn"]]
        )

        self.components["download_model_btn"].click(
            fn=self.app.trainer.get_model_output_safetensors,
            outputs=[self.components["download_model_btn"]]
        )
        
        # Global stop button
        self.components["global_stop_btn"].click(
            fn=self.handle_global_stop,
            outputs=[
                self.components["global_status"],
                self.app.tabs["split_tab"].components["video_list"],
                self.app.tabs["caption_tab"].components["training_dataset"],
                self.app.tabs["train_tab"].components["status_box"],
                self.app.tabs["train_tab"].components["log_box"],
                self.app.tabs["split_tab"].components["detect_status"],
                self.app.tabs["import_tab"].components["import_status"],
                self.app.tabs["caption_tab"].components["preview_status"]
            ]
        )
        
        # Push model button 
        self.components["push_model_btn"].click(
            fn=lambda repo_id: self.upload_to_hub(repo_id),
            inputs=[self.components["repo_id"]],
            outputs=[self.components["global_status"]]
        )
        
    def validate_repo(self, repo_id: str) -> gr.update:
        """Validate repository ID for HuggingFace Hub"""
        validation = validate_model_repo(repo_id)
        if validation["error"]:
            return gr.update(value=repo_id, error=validation["error"])
        return gr.update(value=repo_id, error=None)
        
    def upload_to_hub(self, repo_id: str) -> str:
        """Upload model to HuggingFace Hub"""
        if not repo_id:
            return "Error: Repository ID is required"
        
        # Validate repository name
        validation = validate_model_repo(repo_id)
        if validation["error"]:
            return f"Error: {validation['error']}"
        
        # Check if we have a model to upload
        if not self.app.trainer.get_model_output_safetensors():
            return "Error: No model found to upload"
        
        # Upload model to hub
        success = self.app.trainer.upload_to_hub(OUTPUT_PATH, repo_id)
        
        if success:
            return f"Successfully uploaded model to {repo_id}"
        else:
            return f"Failed to upload model to {repo_id}"
            
    def handle_global_stop(self):
        """Handle the global stop button click"""
        result = self.stop_all_and_clear()
        
        # Format the details for display
        status = result["status"]
        details = "\n".join(f"{k}: {v}" for k, v in result["details"].items())
        full_status = f"{status}\n\nDetails:\n{details}"
        
        # Get fresh lists after cleanup
        videos = self.app.tabs["split_tab"].list_unprocessed_videos()
        clips = self.app.tabs["caption_tab"].list_training_files_to_caption()
        
        return {
            self.components["global_status"]: gr.update(value=full_status, visible=True),
            self.app.tabs["split_tab"].components["video_list"]: videos,
            self.app.tabs["caption_tab"].components["training_dataset"]: clips,
            self.app.tabs["train_tab"].components["status_box"]: "Training stopped and data cleared",
            self.app.tabs["train_tab"].components["log_box"]: "",
            self.app.tabs["split_tab"].components["detect_status"]: "Scene detection stopped",
            self.app.tabs["import_tab"].components["import_status"]: "All data cleared",
            self.app.tabs["caption_tab"].components["preview_status"]: "Captioning stopped"
        }
        
    def stop_all_and_clear(self) -> Dict[str, str]:
        """Stop all running processes and clear data
        
        Returns:
            Dict with status messages for different components
        """
        status_messages = {}
        
        try:
            # Stop training if running
            if self.app.trainer.is_training_running():
                training_result = self.app.trainer.stop_training()
                status_messages["training"] = training_result["status"]
            
            # Stop captioning if running
            if self.app.captioner:
                self.app.captioner.stop_captioning()
                status_messages["captioning"] = "Captioning stopped"
            
            # Stop scene detection if running
            if self.app.splitter.is_processing():
                self.app.splitter.processing = False
                status_messages["splitting"] = "Scene detection stopped"
            
            # Properly close logging before clearing log file
            if self.app.trainer.file_handler:
                self.app.trainer.file_handler.close()
                logger.removeHandler(self.app.trainer.file_handler)
                self.app.trainer.file_handler = None
                
            if LOG_FILE_PATH.exists():
                LOG_FILE_PATH.unlink()
            
            # Clear all data directories
            for path in [VIDEOS_TO_SPLIT_PATH, STAGING_PATH, TRAINING_VIDEOS_PATH, TRAINING_PATH,
                        MODEL_PATH, OUTPUT_PATH]:
                if path.exists():
                    try:
                        shutil.rmtree(path)
                        path.mkdir(parents=True, exist_ok=True)
                    except Exception as e:
                        status_messages[f"clear_{path.name}"] = f"Error clearing {path.name}: {str(e)}"
                    else:
                        status_messages[f"clear_{path.name}"] = f"Cleared {path.name}"
            
            # Reset any persistent state
            self.app.tabs["caption_tab"]._should_stop_captioning = True
            self.app.splitter.processing = False
            
            # Recreate logging setup
            self.app.trainer.setup_logging()
            
            return {
                "status": "All processes stopped and data cleared",
                "details": status_messages
            }
            
        except Exception as e:
            return {
                "status": f"Error during cleanup: {str(e)}",
                "details": status_messages
            }