import gradio as gr
import os
from huggingface_hub import InferenceClient
from utils import validate_token, format_chat_history
from config import DEFAULT_MODEL, DEFAULT_PROVIDER, MAX_TOKENS, TEMPERATURE

# Initialize the client with secure token handling
def get_client():
    """Initialize and return the Hugging Face Inference Client"""
    token = os.getenv("HF_TOKEN")
    if not token:
        raise ValueError("HF_TOKEN environment variable is not set")
    
    return InferenceClient(token=token)

def chat_response(message, history, model, provider, temperature, max_tokens):
    """
    Generate chat response using Hugging Face Inference API
    """
    try:
        client = get_client()
        
        # Format the history for the API
        messages = format_chat_history(history, message)
        
        # Make the API call - try different approaches for provider
        try:
            # First try with provider parameter (newer versions)
            response = client.chat_completion(
                model=model,
                messages=messages,
                max_tokens=max_tokens,
                temperature=temperature,
                stream=False
            )
        except TypeError as e:
            if "provider" in str(e):
                # If provider parameter not supported, try without it
                # Some providers might be specified in the model name or URL
                response = client.chat_completion(
                    model=f"{provider}/{model}" if provider else model,
                    messages=messages,
                    max_tokens=max_tokens,
                    temperature=temperature,
                    stream=False
                )
            else:
                raise e
        
        # Extract the assistant's response
        assistant_message = response.choices[0].message.content
        
        return assistant_message
        
    except Exception as e:
        return f"Error: {str(e)}"

def user_message(user_message, history):
    """Add user message to history and return empty string"""
    return "", history + [{"role": "user", "content": user_message}]

def bot_message(history, model, provider, temperature, max_tokens):
    """Generate bot response and add to history"""
    try:
        if not history:
            return history
        
        # Get the last user message
        last_message = history[-1]["content"] if history else ""
        
        # Get previous messages (excluding the current user message)
        previous_history = history[:-1] if len(history) > 1 else []
        
        # Generate response
        response = chat_response(last_message, previous_history, model, provider, temperature, max_tokens)
        
        # Add assistant response to history
        history.append({"role": "assistant", "content": response})
        
        return history
        
    except Exception as e:
        history.append({"role": "assistant", "content": f"Error: {str(e)}"})
        return history

def create_interface():
    """Create and configure the Gradio interface"""
    
    with gr.Blocks(
        title="Hugging Face Chat Interface",
        theme=gr.themes.Soft(),
        css="""
        .gradio-container {
            max-width: 800px;
            margin: auto;
        }
        """
    ) as demo:
        
        gr.Markdown("""
        # 🤖 Hugging Face Chat Interface
        Built with [anycoder](https://huggingface.co/spaces/akhaliq/anycoder)
        
        Chat with Llama-3.1-8B-Instruct using Hugging Face Inference API.
        """)
        
        with gr.Row():
            with gr.Column(scale=1):
                model_input = gr.Textbox(
                    value=DEFAULT_MODEL,
                    label="Model",
                    placeholder="Enter model name"
                )
                provider_input = gr.Textbox(
                    value=DEFAULT_PROVIDER,
                    label="Provider",
                    placeholder="Enter provider name (may not be supported by all API versions)"
                )
                
            with gr.Column(scale=1):
                temperature_slider = gr.Slider(
                    minimum=0.0,
                    maximum=2.0,
                    value=TEMPERATURE,
                    step=0.1,
                    label="Temperature"
                )
                max_tokens_slider = gr.Slider(
                    minimum=1,
                    maximum=4096,
                    value=MAX_TOKENS,
                    step=1,
                    label="Max Tokens"
                )
        
        chatbot = gr.Chatbot(
            type="messages",
            label="Chat",
            height=400,
            show_copy_button=True
        )
        
        msg = gr.Textbox(
            label="Your message",
            placeholder="Type your message here...",
            lines=2
        )
        
        with gr.Row():
            submit_btn = gr.Button("Send", variant="primary")
            clear_btn = gr.Button("Clear")
        
        # State to maintain chat history
        chat_history = gr.State([])
        
        # Handle message submission - proper flow for type="messages"
        msg.submit(
            user_message,
            inputs=[msg, chat_history],
            outputs=[msg, chat_history]
        ).then(
            bot_message,
            inputs=[chat_history, model_input, provider_input, temperature_slider, max_tokens_slider],
            outputs=[chat_history]
        ).then(
            lambda history: history,
            inputs=[chat_history],
            outputs=[chatbot]
        )
        
        submit_btn.click(
            user_message,
            inputs=[msg, chat_history],
            outputs=[msg, chat_history]
        ).then(
            bot_message,
            inputs=[chat_history, model_input, provider_input, temperature_slider, max_tokens_slider],
            outputs=[chat_history]
        ).then(
            lambda history: history,
            inputs=[chat_history],
            outputs=[chatbot]
        )
        
        clear_btn.click(
            lambda: [],
            outputs=[chat_history]
        ).then(
            lambda: [],
            outputs=[chatbot]
        )
        
        # Example inputs
        gr.Examples(
            examples=[
                ["Hello! How are you?"],
                ["Can you explain quantum computing in simple terms?"],
                ["What are the benefits of renewable energy?"],
                ["Write a short poem about nature."]
            ],
            inputs=[msg]
        )
    
    return demo

if __name__ == "__main__":
    # Validate environment setup
    if not validate_token():
        print("Warning: HF_TOKEN environment variable is not set.")
        print("Please set your Hugging Face token as an environment variable.")
        print("Example: export HF_TOKEN=your_token_here")
    
    demo = create_interface()
    demo.launch(
        share=False,
        show_error=True,
        show_api=True
    )