vms/tabs/import_tab/hub_tab.py

"""
Hugging Face Hub tab for Video Model Studio UI.
Handles browsing, searching, and importing datasets from the Hugging Face Hub.
"""

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

from ..base_tab import BaseTab

logger = logging.getLogger(__name__)

class HubTab(BaseTab):
    """Hub tab for importing datasets from Hugging Face Hub"""
    
    def __init__(self, app_state):
        super().__init__(app_state)
        self.id = "hub_tab"
        self.title = "Import from Hugging Face"
    
    def create(self, parent=None) -> gr.Tab:
        """Create the Hub tab UI components"""
        with gr.Tab(self.title, id=self.id) as tab:
            with gr.Column():
                with gr.Row():
                    gr.Markdown("## Import from Hub datasets")
                
                with gr.Row():
                    gr.Markdown("Search for datasets with videos or WebDataset archives:")
                
                with gr.Row():
                    self.components["dataset_search"] = gr.Textbox(
                        label="Search Hugging Face Datasets",
                        placeholder="Search for video datasets..."
                    )
                
                with gr.Row():
                    self.components["dataset_search_btn"] = gr.Button("Search Datasets", variant="primary")
                
                # Dataset browser results section
                with gr.Row(visible=False) as dataset_results_row:
                    self.components["dataset_results_row"] = dataset_results_row
                    
                    with gr.Column(scale=3):
                        self.components["dataset_results"] = gr.Dataframe(
                            headers=["id", "title", "downloads"],
                            interactive=False,
                            wrap=True,
                            row_count=10,
                            label="Dataset Results"
                        )
                    
                    with gr.Column(scale=3):
                        # Dataset info and state
                        self.components["dataset_info"] = gr.Markdown("Select a dataset to see details")
                        self.components["dataset_id"] = gr.State(value=None)
                        self.components["file_type"] = gr.State(value=None)
                        
                        # Files section that appears when a dataset is selected
                        with gr.Column(visible=False) as files_section:
                            self.components["files_section"] = files_section
                            
                            gr.Markdown("## Files:")
                            
                            # Video files row (appears if videos are present)
                            with gr.Row(visible=False) as video_files_row:
                                self.components["video_files_row"] = video_files_row
                                
                                with gr.Column(scale=4):
                                    self.components["video_count_text"] = gr.Markdown("Contains 0 video files")
                                
                                with gr.Column(scale=1):
                                    self.components["download_videos_btn"] = gr.Button("Download", variant="primary")
                            
                            # WebDataset files row (appears if tar files are present)
                            with gr.Row(visible=False) as webdataset_files_row:
                                self.components["webdataset_files_row"] = webdataset_files_row
                                
                                with gr.Column(scale=4):
                                    self.components["webdataset_count_text"] = gr.Markdown("Contains 0 WebDataset (.tar) files")
                                
                                with gr.Column(scale=1):
                                    self.components["download_webdataset_btn"] = gr.Button("Download", variant="primary")
                        
                        # Status and loading indicators
                        self.components["dataset_loading"] = gr.Markdown(visible=False)
            
            return tab
    
    def connect_events(self) -> None:
        """Connect event handlers to UI components"""
        # Dataset search event
        self.components["dataset_search_btn"].click(
            fn=self.search_datasets,
            inputs=[self.components["dataset_search"]],
            outputs=[
                self.components["dataset_results"],
                self.components["dataset_results_row"]
            ]
        )
        
        # Dataset selection event - FIX HERE
        self.components["dataset_results"].select(
            fn=self.display_dataset_info,
            outputs=[
                self.components["dataset_info"],
                self.components["dataset_id"],
                self.components["files_section"],
                self.components["video_files_row"],
                self.components["video_count_text"],
                self.components["webdataset_files_row"],
                self.components["webdataset_count_text"]
            ]
        )
        
        # Download videos button
        self.components["download_videos_btn"].click(
            fn=self.set_file_type_and_return,
            outputs=[self.components["file_type"]]
        ).then(
            fn=self.download_file_group,
            inputs=[
                self.components["dataset_id"],
                self.components["enable_automatic_video_split"],
                self.components["file_type"]
            ],
            outputs=[
                self.components["dataset_loading"],
                self.components["import_status"]
            ]
        ).success(
            fn=self.app.tabs["import_tab"].on_import_success,
            inputs=[
                self.components["enable_automatic_video_split"],
                self.components["enable_automatic_content_captioning"],
                self.app.tabs["caption_tab"].components["custom_prompt_prefix"]
            ],
            outputs=[
                self.app.tabs_component,
                self.app.tabs["split_tab"].components["video_list"],
                self.app.tabs["split_tab"].components["detect_status"]
            ]
        )
        
        # Download WebDataset button
        self.components["download_webdataset_btn"].click(
            fn=self.set_file_type_and_return_webdataset,
            outputs=[self.components["file_type"]]
        ).then(
            fn=self.download_file_group,
            inputs=[
                self.components["dataset_id"],
                self.components["enable_automatic_video_split"],
                self.components["file_type"]
            ],
            outputs=[
                self.components["dataset_loading"],
                self.components["import_status"]
            ]
        ).success(
            fn=self.app.tabs["import_tab"].on_import_success,
            inputs=[
                self.components["enable_automatic_video_split"],
                self.components["enable_automatic_content_captioning"],
                self.app.tabs["caption_tab"].components["custom_prompt_prefix"]
            ],
            outputs=[
                self.app.tabs_component,
                self.app.tabs["split_tab"].components["video_list"],
                self.app.tabs["split_tab"].components["detect_status"]
            ]
        )

    def set_file_type_and_return(self):
        """Set file type to video and return it"""
        return "video"
    
    def set_file_type_and_return_webdataset(self):
        """Set file type to webdataset and return it"""
        return "webdataset"
    
    def search_datasets(self, query: str):
        """Search datasets on the Hub matching the query"""
        try:
            logger.info(f"Searching for datasets with query: '{query}'")
            results = self.app.importer.search_datasets(query)
            return results, gr.update(visible=True)
        except Exception as e:
            logger.error(f"Error searching datasets: {str(e)}", exc_info=True)
            return [[f"Error: {str(e)}", "", ""]], gr.update(visible=True)
    
    def display_dataset_info(self, evt: gr.SelectData):
        """Display detailed information about the selected dataset"""
        try:
            if not evt or not evt.value:
                logger.warning("No dataset selected in display_dataset_info")
                return (
                    "No dataset selected",  # dataset_info
                    None,                   # dataset_id
                    gr.update(visible=False), # files_section
                    gr.update(visible=False), # video_files_row
                    "",                     # video_count_text
                    gr.update(visible=False), # webdataset_files_row
                    ""                      # webdataset_count_text
                )
            
            dataset_id = evt.value[0] if isinstance(evt.value, list) else evt.value
            logger.info(f"Getting dataset info for: {dataset_id}")
            
            # Use the importer service to get dataset info
            info_text, file_counts, _ = self.app.importer.get_dataset_info(dataset_id)
            
            # Get counts of each file type
            video_count = file_counts.get("video", 0)
            webdataset_count = file_counts.get("webdataset", 0)
            
            # Return all the required outputs individually
            return (
                info_text,                                # dataset_info
                dataset_id,                              # dataset_id
                gr.update(visible=True),                 # files_section
                gr.update(visible=video_count > 0),      # video_files_row
                f"Contains {video_count} video file{'s' if video_count != 1 else ''}", # video_count_text
                gr.update(visible=webdataset_count > 0), # webdataset_files_row
                f"Contains {webdataset_count} WebDataset (.tar) file{'s' if webdataset_count != 1 else ''}" # webdataset_count_text
            )
        except Exception as e:
            logger.error(f"Error displaying dataset info: {str(e)}", exc_info=True)
            return (
                f"Error loading dataset information: {str(e)}", # dataset_info
                None,                                          # dataset_id
                gr.update(visible=False),                      # files_section
                gr.update(visible=False),                      # video_files_row
                "",                                            # video_count_text
                gr.update(visible=False),                      # webdataset_files_row
                ""                                             # webdataset_count_text
            )
    
    def download_file_group(self, dataset_id: str, enable_splitting: bool, file_type: str) -> Tuple[gr.update, str]:
        """Handle download of a group of files (videos or WebDatasets)"""
        try:
            if not dataset_id:
                return gr.update(visible=False), "No dataset selected"
            
            logger.info(f"Starting download of {file_type} files from dataset: {dataset_id}")
            
            # Show loading indicator
            loading_msg = gr.update(
                value=f"## Downloading {file_type} files from {dataset_id}\n\nThis may take some time...",
                visible=True
            )
            status_msg = f"Downloading {file_type} files from {dataset_id}..."
            
            # Use the async version in a non-blocking way
            asyncio.create_task(self._download_file_group_bg(dataset_id, file_type, enable_splitting))
            
            return loading_msg, status_msg
            
        except Exception as e:
            error_msg = f"Error initiating download: {str(e)}"
            logger.error(error_msg, exc_info=True)
            return gr.update(visible=False), error_msg
    
    async def _download_file_group_bg(self, dataset_id: str, file_type: str, enable_splitting: bool):
        """Background task for group file download"""
        try:
            # This will execute in the background
            await self.app.importer.download_file_group(dataset_id, file_type, enable_splitting)
        except Exception as e:
            logger.error(f"Error in background file group download: {str(e)}", exc_info=True)