import os
from threading import Thread
from typing import Iterator
import json

import gradio as gr
import spaces
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
import subprocess
import copy

import subprocess
import sys

def run_command(command):
    process = subprocess.Popen(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=True)
    output, error = process.communicate()
    if process.returncode != 0:
        print(f"Error executing command: {command}")
        print(f"Error message: {error.decode('utf-8')}")
        sys.exit(1)
    return output.decode('utf-8')

MAX_MAX_NEW_TOKENS = 2048
DEFAULT_MAX_NEW_TOKENS = 1024
MAX_INPUT_TOKEN_LENGTH = int(os.getenv("MAX_INPUT_TOKEN_LENGTH", "8000"))
model_choices = [
    "rubra-ai/Meta-Llama-3-8B-Instruct",
    "rubra-ai/Qwen2-7B-Instruct",
    "rubra-ai/Phi-3-mini-128k-instruct",
    "rubra-ai/Mistral-7B-Instruct-v0.3",
    # "rubra-ai/Mistral-7B-Instruct-v0.2",
    # "rubra-ai/gemma-1.1-2b-it"
]

DESCRIPTION = """\
# Rubra v0.1 - A Collection of Tool (Function) Calling LLMs

This is a demo of the Rubra collection of models. You can use the models for general conversation, 
task completion, and function calling with the provided tools input.

See more at https://docs.rubra.ai/ & https://github.com/rubra-ai/rubra
"""

model_table = """
<p/>

---

## Rubra Benchmarks

| Model                                    | Params (in billions) | Function Calling | MMLU (5-shot) | GPQA (0-shot) | GSM-8K (8-shot, CoT) | MATH (4-shot, CoT) | MT-bench |
|------------------------------------------|----------------------|------------------|---------------|---------------|----------------------|--------------------|----------|
| GPT-4o                                   | -                    | 98.57%           | -             | 53.6          | -                    | -                  | -        |
| Claude-3.5 Sonnet                        | -                    | 98.57%           | 88.7          | 59.4          | -                    | -                  | -        |
| Rubra Llama-3 70B Instruct               | 70.6                 | 97.85%           | 75.90         | 33.93         | 82.26                | 34.24              | 8.36     |
| Rubra Llama-3 8B Instruct                | 8.9                  | 89.28%           | 64.39         | 31.70         | 68.99                | 23.76              | 8.03     |
| Rubra Qwen2-7B-Instruct                  | 8.55                 | 85.71%           | 68.88         | 30.36         | 75.82                | 28.72              | 8.08     |
| Rubra Mistral 7B Instruct v0.3           | 8.12                 | 73.57%           | 59.12         | 29.91         | 43.29                | 11.14              | 7.69     |
| Rubra Phi-3 Mini 128k Instruct           | 4.73                 | 70.00%           | 67.87         | 29.69         | 79.45                | 30.80              | 8.21     |
| Rubra Mistral 7B Instruct v0.2           | 8.11                 | 69.28%           | 58.90         | 29.91         | 34.12                | 8.36               | 7.36     |
| meetkai/functionary-small-v2.5           | 8.03                 | 57.14%           | 63.92         | 32.14         | 66.11                | 20.54              | 7.09     |
| Nexusflow/NexusRaven-V2-13B              | 13.0                 | 53.75% ∔         | 43.23         | 28.79         | 22.67                | 7.12               | 5.36     |
| Mistral Large (closed-source)            | -                    | 48.60%           | -             | -             | 91.21                | 45.0               | -        |
| Rubra Gemma-1.1 2B Instruct              | 2.84                 | 45.00%           | 38.85         | 24.55         | 6.14                 | 2.38               | 5.75     |
| meetkai/functionary-medium-v3.0          | 70.6                 | 46.43%           | 79.85         | 38.39         | 89.54                | 43.02              | 5.49     |
| gorilla-llm/gorilla-openfunctions-v2     | 6.91                 | 41.25% ∔         | 49.14         | 23.66         | 48.29                | 17.54              | 5.13     |
| NousResearch/Hermes-2-Pro-Llama-3-8B     | 8.03                 | 41.25%           | 64.16         | 31.92         | 73.92                | 21.58              | 7.83     |
| Mistral 7B Instruct v0.3                 | 7.25                 | 22.5%            | 62.10         | 30.58         | 53.07                | 12.98              | 7.50     |
| Gemma-1.1 2B Instruct                    | 2.51                 | -                | 37.84         | 22.99         | 6.29                 | 6.14               | 5.82     |
| Llama-3 8B Instruct                      | 8.03                 | -                | 65.69         | 31.47         | 77.41                | 27.58              | 8.07     |
| Llama-3 70B Instruct                     | 70.6                 | -                | 79.90         | 38.17         | 90.67                | 44.24              | 8.88     |
| Mistral 7B Instruct v0.2                 | 7.24                 | -                | 59.27         | 27.68         | 43.21                | 10.30              | 7.50     |
| Phi-3 Mini 128k Instruct                 | 3.82                 | -                | 69.36         | 27.01         | 83.7                 | 32.92              | 8.02     |
| Qwen2-7B-Instruct                        | 7.62                 | -                | 70.78         | 32.14         | 78.54                | 30.10              | 8.29     |

∔ `Nexusflow/NexusRaven-V2-13B` and `gorilla-llm/gorilla-openfunctions-v2` don't accept tool observations, the result of running a tool or function once the LLM calls it, so we appended the observation to the prompt.
"""

LICENSE = """
<p/>

---
Rubra code is licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.

Rubra models are licensed under the parent model's license. See the parent model card for more information.
"""

if not torch.cuda.is_available():
    DESCRIPTION += "\n<p>Running on CPU 🥶 This demo does not work on CPU.</p>"

if torch.cuda.is_available():
    model_id = "rubra-ai/Meta-Llama-3-8B-Instruct"  # Default model
    model = None
    tokenizer = None

    def load_model(model_name):
        global model, tokenizer
        model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto", load_in_4bit=False)
        tokenizer = AutoTokenizer.from_pretrained(model_name)
        model.generation_config.pad_token_id = tokenizer.pad_token_id

    load_model(model_id)  # Load the default model

def is_valid_json(tools: str) -> bool:
    try:
        json.loads(tools)
        return True
    except ValueError:
        return False

def validate_tools(tools):
    if tools.strip() == "" or is_valid_json(tools):
        return gr.update(visible=False)
    else:
        return gr.update(visible=True)

def json_to_markdown(json_obj):
    """Convert a JSON object to a formatted markdown string."""
    markdown = ""
    for item in json_obj:
        if item.get("type") == "text":
            # For text items, just add the text content
            markdown += item.get("text", "") + "\n\n"
        elif item.get("type") == "function":
            # For function calls, format as JSON
            markdown += "```json\n"
            # markdown += json.dumps(item.get("function", {}), indent=2)
            markdown += json.dumps(item, indent=2)
            markdown += "\n```\n\n"
    return markdown.strip()

def user(user_message, history):
    return "", history + [[user_message, None]]

def bot(history, system_prompt, tools, role, max_new_tokens, temperature):
    user_message = history[-1][0]
    if history[-1][1] is None:
        history[-1][1] = ""  # Ensure it's never None

    ui_history = list(history)  # Clone the history for UI updates
    all_tool_outputs = []  # Store all processed outputs for final aggregation
    output_accumulated = ""  # To accumulate outputs before processing

    for chunk in generate(user_message, history[:-1], system_prompt, tools, role, max_new_tokens, temperature):
        history[-1][1] += chunk
        print(history[-1][1])
        
        if "endtoolcall" in history[-1][1]:
            process_output = postprocess_output(history[-1][1])
            print("process output:\n", process_output)
            if process_output:
                temp_history = copy.deepcopy(history)  # Use deepcopy here
                if isinstance(process_output, list) and len(process_output) > 0 and isinstance(process_output[0], dict):
                    markdown_output = json_to_markdown(process_output)
                    temp_history[-1][1] = markdown_output
                else:
                    temp_history[-1][1] = str(process_output)
                print(temp_history[-1][1])
                print("--------------------------")
                yield temp_history
            else:
                print(history[-1][1])
                print("--------------------------")
                yield history
        else:
            print(history[-1][1])
            print("--------------------------")
            yield history

@spaces.GPU
def generate(
    message: str,
    chat_history: list[tuple[str, str]],
    system_prompt: str,
    tools: str,
    role: str,
    max_new_tokens: int = 1024,
    temperature: float = 0.6,
) -> Iterator[str]:
    global model, tokenizer
    conversation = []
    if system_prompt:
        conversation.append({"role": "system", "content": system_prompt})
    for user, assistant in chat_history:
        conversation.extend([{"role": "user", "content": user}, {"role": "assistant", "content": assistant}])
    conversation.append({"role": role, "content": message})
    
    if tools:
        if not is_valid_json(tools):
            yield "Invalid JSON in tools. Please correct it."
            return
        tools = json.loads(tools)
        formatted_msgs = preprocess_input(msgs=conversation, tools=tools)
    else:
        formatted_msgs = conversation

    input_ids = tokenizer.apply_chat_template(formatted_msgs, return_tensors="pt")
    if input_ids.shape[1] > MAX_INPUT_TOKEN_LENGTH:
        input_ids = input_ids[:, -MAX_INPUT_TOKEN_LENGTH:]
        gr.Warning(f"Trimmed input from conversation as it was longer than {MAX_INPUT_TOKEN_LENGTH} tokens.")
    input_ids = input_ids.to(model.device)

    streamer = TextIteratorStreamer(tokenizer, timeout=10.0, skip_prompt=True, skip_special_tokens=True)
    generate_kwargs = dict(
        input_ids=input_ids,
        streamer=streamer,
        max_new_tokens=max_new_tokens,
        do_sample=True,
        top_p=0.95,
        temperature=temperature,
        num_beams=1,
        repetition_penalty=1.2,
    )
    t = Thread(target=model.generate, kwargs=generate_kwargs)
    t.start()

    for text in streamer:
        # print("Generated text:", text)
        yield text

bot_message = """Hello! How can I assist you today? If you have any questions or need information on a specific topic, feel free to ask. I can also utilize `tools` that you input to help you better. For example:
```
[
    {
        "type": "function",
        "function": {
            "name": "get_stock_information",
            "description": "Get the current stock market information for a given company",
            "parameters": {
                "type": "object",
                "properties": {
                    "ticker_symbol": {
                        "type": "string",
                        "description": "The stock ticker symbol of the company, e.g., 'AAPL' for Apple Inc."
                    },
                    "exchange": {
                        "type": "string",
                        "description": "The stock exchange where the company is listed, e.g., 'NASDAQ'. If not provided, default to the primary exchange for the ticker symbol."
                    },
                    "data_type": {
                        "type": "string",
                        "enum": ["price", "volume", "market_cap"],
                        "description": "The type of stock data to retrieve: 'price' for current price, 'volume' for trading volume, 'market_cap' for market capitalization."
                    }
                },
                "required": ["ticker_symbol", "data_type"]
            }
        }
    }
]
```

You can also define `functions` (deprecated in favor of `tools` in OpenAI):
```
[
    {
        "name": "get_current_date",
        "description": "Gets the current date at the given location. Results are in ISO 8601 date format; e.g. 2024-04-25",
        "parameters": {
            "type": "object",
            "properties": {
                "location": {
                    "type": "string",
                    "description": "The city and state to get the current date at, e.g. San Francisco, CA"
                }
            },
            "required":["location"]
        }
    }
]
```
"""

def create_chat_interface():
    with gr.Blocks(css="style.css") as demo:
        gr.Markdown(DESCRIPTION)
        
        with gr.Row(equal_height=True, elem_id="main-row"):
            with gr.Column(scale=3, min_width=500):
                # Initialize the chatbot with the welcome message
                chatbot = gr.Chatbot(
                    value=[("Hi", bot_message)],
                    show_copy_button=True,
                    elem_id="chatbot",
                    show_label=False,
                    render_markdown=True,
                    height="100%",
                    layout='bubble',
                    avatar_images=("human.png", "bot.png")
                )
                
                error_box = gr.Markdown(visible=False, elem_id="error-box")

            with gr.Column(scale=2, min_width=300):
                model_dropdown = gr.Dropdown(
                    choices=model_choices,
                    label="Select Model",
                    value="rubra-ai/Meta-Llama-3-8B-Instruct"
                )
                model_dropdown.change(load_model, inputs=[model_dropdown])
                
                with gr.Accordion("Settings", open=False):
                    max_new_tokens = gr.Slider(
                        label="Max new tokens",
                        minimum=1,
                        maximum=MAX_MAX_NEW_TOKENS,
                        step=1,
                        value=DEFAULT_MAX_NEW_TOKENS,
                    )
                    temperature = gr.Slider(
                        label="Temperature",
                        minimum=0.0,
                        maximum=1.2,
                        step=0.01,
                        value=0.01,
                    )

                with gr.Row():
                    role = gr.Dropdown(choices=["user", "observation"], value="user", label="Role", scale=4)
                system_prompt = gr.Textbox(label="System Prompt", lines=1, info="Optional")
                tools = gr.Textbox(label="Tools", lines=1, placeholder="Enter tools in JSON format", info="Optional")
                
                
                
                with gr.Row():
                    user_input = gr.Textbox(
                        label="User Input", 
                        placeholder="Type your message here...",
                        show_label=True,
                        scale=8
                    )
                submit_btn = gr.Button("Submit", variant="primary", elem_id="submit-button")
                clear_btn = gr.Button("Clear Conversation", elem_id="clear-button")
                

        tools.change(validate_tools, tools, error_box)

        submit_btn.click(
            user,
            [user_input, chatbot],
            [user_input, chatbot],
            queue=False
        ).then(
            bot,
            [chatbot, system_prompt, tools, role, max_new_tokens, temperature],
            chatbot
        )

        clear_btn.click(lambda: ([], None), outputs=[chatbot, error_box])

        gr.Markdown(model_table)
        gr.Markdown(LICENSE)

    return demo

if __name__ == "__main__":
    # Initialize npm project if package.json doesn't exist
    if not os.path.exists('package.json'):
        print("Initializing npm project...")
        run_command("npm init -y")

    # Install jsonrepair locally
    print("Installing jsonrepair...")
    run_command("npm install jsonrepair")

    # Verify installation
    print("Verifying jsonrepair installation:")
    run_command("npm list jsonrepair")

    # Add node_modules/.bin to PATH
    os.environ['PATH'] = f"{os.path.join(os.getcwd(), 'node_modules', '.bin')}:{os.environ['PATH']}"
    from preprocess import preprocess_input
    from postprocess import postprocess_output
    demo = create_chat_interface()
    demo.queue(max_size=20).launch()