import os
import yaml
from smolagents import GradioUI, CodeAgent, OpenAIServerModel
from tools.web_search import DuckDuckGoSearchTool as WebSearch
from tools.final_answer import FinalAnswerTool as FinalAnswer

# Get current directory path
CURRENT_DIR = os.path.dirname(os.path.abspath(__file__))

# Initialize the model
model = OpenAIServerModel(
    model_id='qwen/qwen2.5-coder-32b-instruct',
    api_key=os.getenv("OPENAI_API_KEY", "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"),
    api_base=os.getenv("OPENAI_API_BASE", "https://integrate.api.nvidia.com/v1"),
)

# Initialize tools
web_search = WebSearch()
final_answer = FinalAnswer()

# Load YAML prompt templates
with open(os.path.join(CURRENT_DIR, "prompts.yaml"), "r") as f:
    raw_templates = yaml.safe_load(f)

# ✅ FIXED: Function name corrected from unwrap_text_object to unwrap_text_objects
def unwrap_text_objects(obj):
    if isinstance(obj, dict):
        if obj.get("type") == "text" and "text" in obj:
            return obj["text"]
        return {k: unwrap_text_objects(v) for k, v in obj.items()}
    elif isinstance(obj, list):
        return "\n".join(unwrap_text_objects(i) for i in obj)
    return obj

# Apply the unwrapping to raw templates
prompt_templates = unwrap_text_objects(raw_templates)

print("✅ Loaded YAML keys:", list(raw_templates.keys()))
print("🧾 System prompt preview:\n", raw_templates["system_prompt"][:200])


# ✅ Debug print to verify structure
print("🟢 System Prompt:", prompt_templates.get("system_prompt"))




# Create the agent
agent = CodeAgent(
    model=model,
    tools=[web_search, final_answer],
    managed_agents=[],
    max_steps=20,
    verbosity_level=1,
    name="CodeAgent",
    description="An autonomous coding assistant",
    executor_type='local',
    executor_kwargs={},
    prompt_templates=prompt_templates
)

# Launch the Gradio UI
if __name__ == "__main__":
    GradioUI(agent).launch()