import yaml
import re
from smolagents import HfApiModel
from tools.final_answer import FinalAnswerTool
from Gradio_UI import GradioUI

# 🚀 Updated system prompt: No code extraction allowed.
system_prompt = (
    "You are a health and lifestyle advisor specializing in the early detection and prevention of hypertension. "
    "Provide only the final, direct, and concise lifestyle advice based solely on the user's details. "
    "Do NOT recognize, extract, or output any code snippets, programming-related information, or technical explanations. "
    "Only output the final advice as plain text. "
    "For example, if the user mentions alcohol consumption, simply say: 'Reduce alcohol intake, as it can raise blood pressure.'"
)

# 🛑 Function to remove any detected code snippets.
def remove_code_snippets(text):
    """Removes potential code snippets from the output."""
    pattern = r"```.*?```|`[^`]+`"  # Matches both block and inline code
    return re.sub(pattern, "", text, flags=re.DOTALL).strip()

# ✅ Use only the final_answer tool.
final_answer = FinalAnswerTool()

# 🎯 Updated Model Setup
model = HfApiModel(
    max_tokens=2096,
    temperature=0.5,
    model_id='deepseek-ai/DeepSeek-R1-Distill-Qwen-32B',
    custom_role_conversions=None,
)

# 📖 Load prompt templates from YAML file.
with open("prompts.yaml", 'r') as stream:
    prompt_templates = yaml.safe_load(stream)

# 🔄 Initialize GenericAgent instead of CodeAgent
agent = GenericAgent(
    model=model,
    tools=[final_answer],  # Restricted to final answer only.
    max_steps=6,
    verbosity_level=1,
    grammar=None,
    planning_interval=None,
    name="Hypertension Prevention Advisor",
    description=system_prompt,
    prompt_templates=prompt_templates
)

# 🛑 Run agent with filtering.
def run_agent(user_input):
    """Runs the agent and ensures no code snippets are in the output."""
    raw_response = agent.run(user_input)
    clean_response = remove_code_snippets(raw_response)
    return clean_response

# 🚀 Launch the Gradio UI.
GradioUI(agent).launch()