import os
import re
import smolagents
from smolagents import CodeAgent, DuckDuckGoSearchTool, tool

@tool
def calculate_math(expression: str) -> str:
    """Useful for evaluating math expressions or numerical problems safely.
    Args:
        expression: A valid mathematical string to evaluate, like '2 + 2' or '312 * 43'.
    """
    try:
        allowed_chars = "0123456789+-*/(). "
        cleaned = "".join(c for c in expression if c in allowed_chars)
        if cleaned:
            return str(eval(cleaned))
        return "Error: Invalid characters in math expression."
    except Exception as e:
        return f"Error executing calculation: {str(e)}"

class CourseEvaluationAgent:
    def __init__(self):
        print("[Agent System] Initializing Adaptive Bootstrapper...")
        hf_token = os.getenv("HF_TOKEN")
        model_name = "Qwen/Qwen2.5-Coder-32B-Instruct"
        
        # --- DYNAMIC CLASS DETECTION ---
        # We check the library metadata directly to find the right class
        self.model = None
        
        for class_name in ["InferenceClientModel", "HfApiModel", "HfApiEngine", "LiteLLMModel"]:
            if hasattr(smolagents, class_name):
                try:
                    TargetClass = getattr(smolagents, class_name)
                    print(f"[Agent System] Trying class: {class_name}")
                    
                    # Try instantiating with keyword parameters
                    if class_name == "HfApiEngine":
                        self.model = TargetClass(model_id=model_name, token=hf_token)
                    else:
                        self.model = TargetClass(model_id=model_name, token=hf_token)
                        
                    print(f"[Agent System] Successfully bound engine via {class_name}!")
                    break
                except Exception as inner_e:
                    print(f"[Agent System] {class_name} failed setup: {inner_e}")
                    continue

        # Ultimate fallback: If classes fail, see if ApiModel works directly
        if self.model is None and hasattr(smolagents, "ApiModel"):
            try:
                self.model = getattr(smolagents, "ApiModel")(model_id=model_name, token=hf_token)
                print("[Agent System] Bound engine via base ApiModel fallback.")
            except Exception:
                pass

        # If everything failed, we let CodeAgent try raw string fallback
        if self.model is None:
            print("[Agent System] Warning: Forcing direct string fallback mode.")
            self.model = model_name

        # --- BUILD CODE AGENT ---
        # Safely pass our resolved backend configuration to the agent
        self.agent = CodeAgent(
            tools=[DuckDuckGoSearchTool(), calculate_math], 
            model=self.model,
            max_steps=5,
            verbosity_level=1
        )
        
        # Inject styling rules cleanly
        if "system_prompt" in self.agent.prompt_templates:
            self.agent.prompt_templates["system_prompt"] += (
                "\n\nCRITICAL SYSTEM RULE:\n"
                "Your final response after using tools MUST contain ONLY the direct, raw answer text "
                "(e.g., '42', 'Paris', 'John Doe'). Absolutely DO NOT include conversational padding like "
                "'The answer is' or summary sentences. Output just the clean answer string."
            )

    def run_task(self, question: str) -> str:
        try:
            output = self.agent.run(question)
            final_ans = str(output).strip()
            
            # Post-processing to strip out accidental conversational wrappers
            clean_patterns = [
                r"^[Tt]he answer is\s*[:\-]?\s*",
                r"^[Ff]inal [Aa]nswer\s*[:\-]?\s*",
                r"^[\"']", 
                r"[\"']$"
            ]
            for pattern in clean_patterns:
                final_ans = re.sub(pattern, "", final_ans)
                
            return final_ans.strip()
            
        except Exception as e:
            print(f"[CRITICAL AGENT EXCEPTION]: {str(e)}")
            return f"Runtime error: {str(e)[:50]}"