import random
import re
import torch
import transformers

from rich.console import Console
from rich.panel import Panel
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline, GenerationConfig

console = Console()

transformers.logging.set_verbosity_error()

# Set seeds for reproducibility
random.seed(42)
torch.manual_seed(42)
torch.cuda.manual_seed_all(42)

#########
# Model #
#########

model_name = "Qwen/Qwen2-1.5B-Instruct" # (1.5B)

print(f"Loading model {model_name} …")
tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=True)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    device_map="cpu",
    torch_dtype="auto"
)

gen = pipeline(
    "text-generation",
    model=model,
    tokenizer=tokenizer,
    device_map="cpu",
    return_full_text=False,
)

generation_config = GenerationConfig(max_new_tokens=150, do_sample=False)

def generate(prompt: str) -> str:
    out = gen(prompt, generation_config=generation_config)
    return out[0]["generated_text"]

#########
# Tools #
#########

def add(a: float, b: float) -> float:
    """Add two numbers and return the sum."""
    # Fill in

def subtract(a: float, b: float) -> float:
    """Add two numbers and return the difference."""
    # Fill in

def multiply(a: float, b: float) -> float:
    """Multiply two numbers and return the product."""
    # Fill in

def divide(a: float, b: float) -> float:
    """Divide two numbers and return the product."""
    # Fill in

def execute_action(action: str) -> float | str:
    """Parse and execute an action, return observation."""
    action = action.strip()

    # Parse add[a, b]
    # Fill in

    # Parse subtract[a, b]
    # Fill in

    # Parse multiply[a, b]
    # Fill in

    # Parse divide[a, b]
    # Fill in

    # Parse finish[answer]
    match = re.search(r'finish\[(.+)\]', action, re.IGNORECASE)
    if match:
        return f"FINISH: {match.group(1)}"

    return f"Unknown action: {action}"

##########
# Prompt #
##########
system_prompt = f"""
# Fill in
"""

def apply_system_prompt(input: str) -> str:
    return system_prompt + f"\nQuestion: {input}\n"


def react(question: str, max_steps: int = 7, verbose = True):
    """Run ReAct loop following the paper's format."""
    prompt = apply_system_prompt(question)

    console.print(Panel(question, title="Question", border_style="blue"))

    for i in range(1, max_steps + 1):

        # Generate thought and action together, stop at observation
        thought_action = generate(prompt + f"Thought {i}:")

        if verbose:
            console.print(Panel(prompt + f"Thought {i}:", title=f"Raw Prompt: Step {i}", border_style="black"))
            console.print(Panel(thought_action, title=f"Raw Response: Step {i}", border_style="black"))

        # Parse thought and action from first two lines
        lines = thought_action.strip().split("\n")
        thought = # Fill in 
        action = # Fill in

        # Removes the prefix Action x: from response
        action = re.sub(r'^Action\s*\d*:\s*', '', action)

        console.print(Panel(f"Thought {i}: {thought}", title=f"Step {i}", border_style="green"))
        console.print(Panel(f"Action {i}: {action}", border_style="cyan"))

        # Execute the action
        observation = # Fill in 

        console.print(Panel(f"Observation {i}: {observation}", border_style="yellow"))

        # Check if finished
        if isinstance(observation, str) and observation.startswith("FINISH:"):
            answer = observation[7:].strip()
            console.print(Panel(f"Final Answer: {answer}", title="Done", border_style="magenta"))
            return answer

        # Add step to prompt
        step_str = f"""Thought {i}: {thought}
Action {i}: {action}
Observation {i}: {observation}
"""
        prompt += step_str

    console.print("[red]Max steps reached[/red]")
    return ""

if __name__ == "__main__":
    question = "What is 24880 + 220?"
    # question = "What is (24880 / 10) + 220?"
    # question = "What is ((24880 / 10) + 220) * 1234?"
    # question = "What is ((24880 / 10) + 220 - 110) * 1234?"
    # question = "What is ((24880 / 10) + 220 - 110) * 1234?"
    # question = "Bob and Charlie go out to dinner and the bill comes out to $50. Both order the same entree. How much should each pay?"

    react(question, max_steps=7, verbose=True)


