import gradio as gr
import os
import json
import torch
import subprocess
from dotenv import load_dotenv
import logging

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    handlers=[
        logging.StreamHandler(),
        logging.FileHandler("app.log")
    ]
)
logger = logging.getLogger(__name__)

# Load environment variables
load_dotenv()

# Load config file
def load_config(config_path="transformers_config.json"):
    try:
        with open(config_path, 'r') as f:
            config = json.load(f)
        return config
    except Exception as e:
        logger.error(f"Error loading config: {str(e)}")
        return {}

# Load configuration
config = load_config()
model_config = config.get("model_config", {})

# Model details from config
MODEL_NAME = model_config.get("model_name_or_path", "unsloth/DeepSeek-R1-Distill-Qwen-14B-bnb-4bit")
SPACE_NAME = os.getenv("HF_SPACE_NAME", "phi4training")

# Function to start the training process
def start_training():
    try:
        # Run the training script directly - IMPORTANT: Don't redirect output so container logs show
        # Using nohup to ensure process continues even if web request ends
        os.system("nohup python run_cloud_training.py > training.log 2>&1 &")
        
        # Log the start of training
        logger.info("Training started - Check Hugging Face logs for details")
        print("Training process initiated! This will appear in Hugging Face logs.")
        
        return """
        ✅ Training process initiated! 
        
        The model is now being fine-tuned in the background.
        
        To monitor progress:
        1. Check the Hugging Face space logs in the "Logs" tab
        2. Training metrics will be available in the Hugging Face UI
        3. The process will continue running in the background
        
        NOTE: This is a research training phase only, no model outputs will be available.
        """
    except Exception as e:
        logger.error(f"Error starting training: {str(e)}")
        return f"❌ Error starting training: {str(e)}"

# Create Gradio interface - training status only, no model outputs
with gr.Blocks(css="footer {visibility: hidden}") as demo:
    gr.Markdown(f"# {SPACE_NAME}: Research Training Dashboard")
    
    with gr.Row():
        with gr.Column():
            status = gr.Markdown(
                f"""
                ## DeepSeek-R1-Distill-Qwen-14B Research Training
                
                **Model**: {MODEL_NAME}
                **Dataset**: phi4-cognitive-dataset
                
                This is a multidisciplinary research training phase. The model is not available for interactive use.
                
                ### Training Configuration:
                - **Epochs**: {config.get("training_config", {}).get("num_train_epochs", 3)}
                - **Batch Size**: {config.get("training_config", {}).get("per_device_train_batch_size", 2)}
                - **Gradient Accumulation Steps**: {config.get("training_config", {}).get("gradient_accumulation_steps", 4)}
                - **Learning Rate**: {config.get("training_config", {}).get("learning_rate", 2e-5)}
                - **Max Sequence Length**: {config.get("training_config", {}).get("max_seq_length", 2048)}
                
                ⚠️ **NOTE**: This space does not provide model outputs during the research training phase.
                All logs are available in the Hugging Face "Logs" tab.
                """
            )
    
    with gr.Row():
        # Add button for starting training
        start_btn = gr.Button("Start Training", variant="primary")
        
    # Output area for training start messages
    training_output = gr.Markdown("")
    
    # Connect start button to function
    start_btn.click(start_training, outputs=training_output)
    
    gr.Markdown("""
    ### Research Training Information
    
    This model is being fine-tuned on research-focused datasets and is not available for interactive querying.
    The training process will run in the background and logs will be available in the Hugging Face UI.
    
    #### Instructions
    1. Click "Start Training" to begin the fine-tuning process
    2. Monitor progress in the Hugging Face "Logs" tab
    3. Training metrics and results will be saved to the output directory
    
    #### About This Project
    The model is being fine-tuned on the phi4-cognitive-dataset with a focus on research capabilities.
    This training phase does not include any interactive features or output generation.
    """)

# Launch the interface
if __name__ == "__main__":
    # Start Gradio with minimal features
    logger.info("Starting research training dashboard")
    print("Research training dashboard started - Logs will be visible here")
    demo.launch(share=False)