import gradio as gr
from openai import OpenAI
import os
import glob
from pypdf import PdfReader

# Initialize client using Hugging Face routing
client = OpenAI(
    base_url="https://router.huggingface.co/v1",
    api_key=os.environ.get("HF_TOKEN"),
)

# --- Document Context Logic ---
DOCS_DIR = "docs" # Folder for your JMeter cheat sheets, project specs, etc.

def load_pdf_context():
    """Reads all PDFs in the DOCS_DIR once at startup and extracts text."""
    context = ""
    if not os.path.exists(DOCS_DIR):
        os.makedirs(DOCS_DIR)
        print(f"Created '{DOCS_DIR}' directory. Upload your PDFs here.")
        return context
        
    pdf_files = glob.glob(os.path.join(DOCS_DIR, "*.pdf"))
    if not pdf_files:
        print("No PDFs found in the docs folder.")
        return context
        
    print(f"Loading {len(pdf_files)} PDF(s) into context...")
    for file_path in pdf_files:
        try:
            reader = PdfReader(file_path)
            context += f"\n--- Start Document: {os.path.basename(file_path)} ---\n"
            for page in reader.pages:
                text = page.extract_text()
                if text:
                    context += text + "\n"
            context += f"--- End Document: {os.path.basename(file_path)} ---\n\n"
        except Exception as e:
            print(f"Error loading {file_path}: {e}")
            
    return context

DOCUMENT_CONTEXT = load_pdf_context()
# -----------------------------------

def respond(message, history):
    # System Message tailored for an experienced Dev learning JMeter
    system_content = (
        "You are a highly experienced Performance Engineer and Apache JMeter expert. "
        "The user you are helping is an experienced software developer who is new to JMeter and stress testing. "
        "Skip basic programming concepts and focus strictly on JMeter architecture (Test Plans, Thread Groups, Samplers, Listeners, Timers), "
        "performance testing theory (Load vs. Stress vs. Spike testing), and best practices (e.g., using JSR223 Samplers with Groovy instead of BeanShell, running in CLI/non-GUI mode). "
        "Be concise, technical, and highly practical. Provide code snippets or XML structures where helpful.\n\n"
        "Here is the project-specific context and documentation you should use to answer questions:\n"
        f"{DOCUMENT_CONTEXT}"
    )
    
    messages = [{"role": "system", "content": system_content}]
    
    # Add conversation history
    for val in history:
        if val['role'] == 'user':
            messages.append({"role": "user", "content": val['content']})
        else:
            messages.append({"role": "assistant", "content": val['content']})

    # Process the current message
    user_content = []
    if message["text"]:
        user_content.append({"type": "text", "text": message["text"]})
    
    # Image/File handling placeholder
    for file in message["files"]:
        pass 

    messages.append({"role": "user", "content": user_content})

    response = ""
    
    try:
        stream = client.chat.completions.create(
            # You can keep this model or switch to another available on HF-Pro
            model="Qwen/Qwen2.5-Coder-32B-Instruct:nscale", 
            messages=messages,
            stream=True
        )

        for chunk in stream:
            if hasattr(chunk, 'choices') and len(chunk.choices) > 0:
                token = chunk.choices[0].delta.content
                if token:
                    response += token
                    yield response
    except Exception as e:
        yield f"Error encountered: {str(e)} - Please verify the model endpoint or your API token."

# Interface Setup
demo = gr.ChatInterface(
    respond,
    multimodal=True, 
    title="Apache JMeter Performance Expert ⚙️",
    description="Your technical guide for building robust JMeter test plans, analyzing performance metrics, and executing effective stress tests. Backed by your project documentation.",
    examples=[
        {"text": "What is the difference between a Load Test and a Stress Test?"},
        {"text": "How do I extract a dynamic token from a JSON response to use in my next request?"},
        {"text": "Why should I run my tests in non-GUI mode, and what is the command for it?"}
    ]
)

if __name__ == "__main__":
    demo.launch()