"""
DocuMint Train - Gradio UI for LoRA Training
"""

import os
import threading
import gradio as gr
from train import train, get_status, authenticate

# ============ GLOBAL STATE ============

training_thread = None


# ============ HANDLERS ============

def start_training(dataset_name: str, epochs: int, batch_size: int, learning_rate: float):
    """Start training in background thread."""
    global training_thread
    
    if training_thread and training_thread.is_alive():
        return "⚠️ Training already in progress!"
    
    def run():
        train(
            dataset_name=dataset_name if dataset_name.strip() else None,
            epochs=int(epochs),
            batch_size=int(batch_size),
            learning_rate=float(learning_rate)
        )
    
    training_thread = threading.Thread(target=run, daemon=True)
    training_thread.start()
    
    return "🚀 Training started! Check status below."


def refresh_status():
    """Get formatted training status."""
    status = get_status()
    
    output = "## 📊 Training Status\n\n"
    output += f"**Status:** {'🔄 Training' if status['is_training'] else '⏸️ Idle'}\n"
    output += f"**Message:** {status['message']}\n\n"
    
    if status['is_training'] or status['progress'] > 0:
        output += f"**Progress:** {status['progress']:.1f}%\n"
        output += f"**Step:** {status['current_step']} / {status['total_steps']}\n"
        output += f"**Loss:** {status['loss']:.4f}\n"
        
        # Progress bar
        bar_len = 30
        filled = int(bar_len * status['progress'] / 100)
        bar = "█" * filled + "░" * (bar_len - filled)
        output += f"\n`[{bar}]`"
    
    return output


def check_auth():
    """Check HuggingFace authentication."""
    if os.environ.get("HF_TOKEN"):
        return "✅ HF_TOKEN is set"
    return "❌ HF_TOKEN not found - set it in Space secrets!"


# ============ GRADIO UI ============

with gr.Blocks(
    title="DocuMint Train - LoRA Training",
    theme=gr.themes.Soft(primary_hue="orange")
) as demo:
    
    gr.Markdown("""
    # 🏆 DocuMint Train
    ### LoRA Fine-tuning for Qwen2-0.5B
    
    Train custom LoRA adapters for document processing tasks.
    """)
    
    with gr.Row():
        auth_status = gr.Textbox(label="Authentication", value=check_auth(), interactive=False)
    
    with gr.Tabs():
        # Training Tab
        with gr.Tab("🎯 Train"):
            with gr.Row():
                with gr.Column():
                    dataset_input = gr.Textbox(
                        label="Dataset",
                        placeholder="Leave empty for himu1780/DocuMint-Data",
                        info="HuggingFace dataset name or empty for default"
                    )
                    
                    epochs_input = gr.Slider(
                        minimum=1, maximum=10, value=3, step=1,
                        label="Epochs"
                    )
                    
                    batch_input = gr.Slider(
                        minimum=1, maximum=4, value=1, step=1,
                        label="Batch Size",
                        info="Keep low for CPU training"
                    )
                    
                    lr_input = gr.Number(
                        value=2e-4,
                        label="Learning Rate"
                    )
                    
                    train_btn = gr.Button("🚀 Start Training", variant="primary", size="lg")
                
                with gr.Column():
                    train_output = gr.Textbox(label="Output", interactive=False, lines=3)
                    
                    status_display = gr.Markdown()
                    refresh_btn = gr.Button("🔄 Refresh Status")
            
            train_btn.click(
                fn=start_training,
                inputs=[dataset_input, epochs_input, batch_input, lr_input],
                outputs=train_output
            )
            
            refresh_btn.click(fn=refresh_status, outputs=status_display)
            demo.load(fn=refresh_status, outputs=status_display)
        
        # Config Tab
        with gr.Tab("⚙️ Configuration"):
            gr.Markdown("""
            ## Model Configuration
            
            | Setting | Value |
            |---------|-------|
            | Base Model | `Qwen/Qwen2-0.5B-Instruct` |
            | Output Repo | `himu1780/DocuMint-Models` |
            | Data Repo | `himu1780/DocuMint-Data` |
            
            ## LoRA Configuration
            
            | Setting | Value |
            |---------|-------|
            | Rank (r) | 8 |
            | Alpha | 16 |
            | Dropout | 0.05 |
            | Target Modules | q_proj, k_proj, v_proj, o_proj |
            
            ## Training Settings
            
            | Setting | Value |
            |---------|-------|
            | Max Length | 512 tokens |
            | Gradient Accumulation | 4 |
            | Warmup Steps | 100 |
            | Scheduler | Cosine |
            """)
        
        # Help Tab
        with gr.Tab("❓ Help"):
            gr.Markdown("""
            ## How to Use
            
            ### 1. Set HF_TOKEN
            Add your HuggingFace token as a Space secret named `HF_TOKEN`.
            
            ### 2. Prepare Dataset
            Upload your dataset to `himu1780/DocuMint-Data` with one of these formats:
            
            **Instruction Format (Alpaca-style):**
            ```json
            {"instruction": "Summarize this document", "output": "Summary here..."}
            ```
            
            **Q&A Format:**
            ```json
            {"question": "What is in this document?", "answer": "The document contains..."}
            ```
            
            **Plain Text:**
            ```json
            {"text": "Document text here..."}
            ```
            
            ### 3. Start Training
            - Leave dataset empty to use DocuMint-Data
            - Or specify any HuggingFace dataset name
            - Click "Start Training"
            - Monitor progress with "Refresh Status"
            
            ### 4. Use Trained Model
            After training, LoRA adapters will be saved to `himu1780/DocuMint-Models`.
            The main DocuMint app will automatically load these adapters!
            """)
    
    gr.Markdown("""
    ---
    <center>
    
    **DocuMint Train** | [DocuMint](https://huggingface.co/spaces/himu1780/DocuMint) | [Models](https://huggingface.co/himu1780/DocuMint-Models)
    
    </center>
    """)


# ============ LAUNCH ============

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)