import gradio as gr # type: ignore
from transformers import pipeline

def chat_with_model(user_input, messages, pipe):
    """
    Chat with the model and return only its response.
    
    Args:
        user_input (str): The user's input message
        messages (list): List of conversation messages
        pipe: The transformer pipeline object
    
    Returns:
        tuple: (str, list) - The model's response and updated messages
    """
    messages.append({"role": "user", "content": user_input})
    
    # Convert messages to format expected by model
    formatted_messages = [{'role': m['role'], 'content': str(m['content'])} for m in messages]
    
    # Get response from model
    response = pipe(formatted_messages, max_new_tokens=2048)[0]['generated_text']
    
    # Extract just the assistant's response
    # This assumes the model returns the full conversation including the new response
    # We need to parse out just the new response
    try:
        # Try to find the last assistant response in the generated text
        parts = response.split("assistant', 'content': '")
        if len(parts) > 1:
            last_response = parts[-1].split("'}")[0]
        else:
            # Fallback: just use the whole response if we can't parse it
            last_response = response
    except Exception:
        last_response = response
    
    # Add the response to messages history
    messages.append({"role": "assistant", "content": last_response})
    
    return last_response, messages

# Initialize the pipeline
pipe = pipeline("text-generation", model="HuggingFaceTB/SmolLM2-1.7B-Instruct")  # Replace with your model

# Initialize conversation history
messages = [
    {"role": "system", "content": """You are **ModuAssist**, created by the LearnModu Team to help beginners and experienced developers alike with the **Modu programming language**. You mainly speak english and you're integrated into their blog, which provides resources and tutorials about Modu. 

### Key Information:
- **Modu** was developed by Cyteon and released on **December 11, 2024**.
- The LearnModu blog covers all features of Modu, including installation, syntax, and functionality.

---

### Installation
**1. Through Cargo (Recommended)**  
- Install **Rust**, which includes Cargo.  
- Check if Cargo is installed: cargo --version.  
- Run: cargo +nightly install modu.  
- Verify installation: modu.  
- **VSCode Users:** Download the Modu extension on GitHub.

**2. Through Binaries**  
- Download Modu binaries from GitHub Actions.  
- Add them to your PATH environment variable.  
- Verify installation: modu.  

---

### Syntax Overview  
**Hello World:** 
print("Hello, World!");

**User Input:**  
let string = input("Print something: ");
print(string);

**Variables and Types:**  
- Automatic type assignment for variables.  
let string = "text";
let integer = 34;
let boolean = true;

**If Statements:**
if a == b {
    print("Equal!");
} if a !== b {
    print("Not Equal!");
}

**Custom Functions:**  
fn wave(person) {
    print("Hello, ", person, "!");
}
wave("Alice");

**Importing Libraries:**  
- Import libraries with import.  
import "math" as m;
let value = m.sqrt(16);

---

### Advanced Features  
- **Packages:** Install with modu install <package-name>.  
- **File Imports:**  
  Example with main.modu importing something.modu:
  import "something.modu" as sm;
  sm.doSomething();
  
Unfortunately, Modu does not support loops (workaround is the basic_loops package, that adds function loop(another_function, start, end)) and there are also no arrays or dictionaries.

Your main goal is to assist users in debugging, fixing, and understanding Modu programs."""},  # Your full system prompt here
]

def main():
    global messages  # Add this line to declare messages as global
    while True:
        user_input = input("You: ")
        if user_input.lower() == "exit":
            break
            
        response, messages = chat_with_model(user_input, messages, pipe)
        print(f"Model: {response}")

# Add the Gradio interface
iface = gr.Interface(
    fn=chat_with_model, 
    inputs=[
        gr.inputs.Textbox(lines=2, placeholder="Enter your message here..."),
        "state",
        gr.State(messages),  # Initialize with your messages list
    ],
    outputs=[
        gr.outputs.Textbox(label="Model Response"),
        "state",
    ],
    title="ModuAssist Chatbot",
    description="Chat with ModuAssist for help with the Modu programming language.",
)

if __name__ == "__main__":
    iface.launch()
    main()