import gradio as gr
from openai import OpenAI
import os
import time
from typing import List, Tuple, Generator

# Constants
MODEL_ID = "Qwen/Qwen2.5-Coder-32B-Instruct"
HF_API_TOKEN = os.getenv("HF_TOKEN")
DEFAULT_TEMPERATURE = 0.5
MAX_TOKENS_LIMIT = 8192
STREAM_DELAY = 0.015

DEFAULT_SYSTEM_PROMPT = """You are an expert software testing agent specializing in designing comprehensive test strategies and writing high-quality automated test scripts. Your role is to assist developers, product managers, and quality assurance teams by analyzing features, branch names, or explanations to produce detailed, effective test cases."""

FORMATTING_TAGS = [
    "[Understand]", "[Plan]", "[Conclude]", 
    "[Reason]", "[Verify]", "[Capabilities]",
    "[Response Guidelines]"
]

def initialize_client() -> OpenAI:
    """Initialize and return the OpenAI client with Hugging Face configuration."""
    return OpenAI(
        base_url="https://api-inference.huggingface.co/v1/",
        api_key=HF_API_TOKEN
    )

def format_response(text: str) -> str:
    """Apply HTML formatting to special tags in the response text."""
    for tag in FORMATTING_TAGS:
        text = text.replace(
            tag, 
            f'<strong class="special-tag">{tag}</strong>'
        )
    return text

def construct_messages(
    user_input: str,
    chat_history: List[Tuple[str, str]],
    system_prompt: str
) -> List[dict]:
    """Construct the message history for the API request."""
    messages = [{"role": "system", "content": system_prompt}]
    
    for user_msg, bot_msg in chat_history:
        messages.extend([
            {"role": "user", "content": user_msg},
            {"role": "assistant", "content": bot_msg}
        ])
    
    messages.append({"role": "user", "content": user_input})
    return messages

def handle_api_error(e: Exception) -> str:
    """Generate user-friendly error messages for different error types."""
    error_type = type(e).__name__
    if "Authentication" in str(e):
        return "🔒 Authentication Error: Check your API token"
    elif "Timeout" in str(e):
        return "⏳ Request Timeout: Try again later"
    return f"⚠️ Error ({error_type}): {str(e)}"

def generate_response(
    message: str,
    chat_history: List[Tuple[str, str]],
    system_prompt: str,
    temperature: float,
    max_tokens: int
) -> Generator[List[Tuple[str, str]], None, None]:
    """Generate streaming response using Hugging Face inference API."""
    client = initialize_client()
    new_history = chat_history + [(message, "")]
    partial_message = ""
    
    try:
        messages = construct_messages(message, chat_history, system_prompt)
        
        stream = client.chat.completions.create(
            model=MODEL_ID,
            messages=messages,
            temperature=temperature,
            max_tokens=min(max_tokens, MAX_TOKENS_LIMIT),
            stream=True
        )

        for chunk in stream:
            if chunk.choices[0].delta.content:
                partial_message += chunk.choices[0].delta.content
                formatted_response = format_response(partial_message + "▌")
                new_history[-1] = (message, formatted_response)
                yield new_history
                time.sleep(STREAM_DELAY)

        new_history[-1] = (message, format_response(partial_message))
        
    except Exception as e:
        error_msg = handle_api_error(e)
        new_history[-1] = (message, error_msg)
    
    yield new_history

def create_interface() -> gr.Blocks:
    """Create and configure the Gradio interface."""
    css = """
    .gr-chatbot { min-height: 500px; border-radius: 15px; }
    .special-tag { color: #2ecc71; font-weight: 600; }
    footer { visibility: hidden; }
    """
    
    with gr.Blocks(css=css, theme=gr.themes.Soft()) as interface:
        gr.Markdown("""
        # 🧠 AI Test Engineering Assistant
        ## Specialized in Automated Testing Strategies
        """)
        
        chatbot = gr.Chatbot(label="Testing Discussion", elem_classes="gr-chatbot")
        user_input = gr.Textbox(label="Feature Description", placeholder="Describe feature or paste branch name...")
        
        with gr.Accordion("Engine Parameters", open=False):
            system_prompt = gr.TextArea(value=DEFAULT_SYSTEM_PROMPT, label="System Instructions")
            temperature = gr.Slider(0, 1, value=DEFAULT_TEMPERATURE, label="Creativity Level")
            max_tokens = gr.Slider(128, MAX_TOKENS_LIMIT, value=2048, label="Response Length")
        
        with gr.Row():
            clear_btn = gr.Button("🧹 Clear History")
            submit_btn = gr.Button("🚀 Generate Tests")
        
        user_input.submit(
            generate_response,
            [user_input, chatbot, system_prompt, temperature, max_tokens],
            [chatbot]
        )
        
        submit_btn.click(
            generate_response,
            [user_input, chatbot, system_prompt, temperature, max_tokens],
            [chatbot]
        )
        
        clear_btn.click(lambda: [], None, chatbot)

    return interface

if __name__ == "__main__":
    if not HF_API_TOKEN:
        raise ValueError("HF_API_TOKEN environment variable not set - add it in Spaces settings!")
    
    interface = create_interface()
    interface.launch(server_name="0.0.0.0", server_port=7860)