purpose_agent/easy.py

"""
Easy API — Build self-improving AI agent teams with zero technical expertise.

This module is the ONLY thing a non-technical user needs to touch.
Everything else is auto-configured.

    import purpose_agent as pa

    # One line. That's it.
    team = pa.purpose("Build a research assistant that finds and summarizes papers")
    result = team.run("Find recent papers on climate change solutions")
    print(result)

Three levels of usage:
    Level 1 (Beginner):  pa.purpose("description") → working team
    Level 2 (Intermediate): pa.Team.build(agents=[...]) → custom team
    Level 3 (Advanced):  pa.Agent(), pa.Graph(), pa.Conversation() → full control
"""

from __future__ import annotations

import logging
import os
import re
from typing import Any

from purpose_agent.unified import Agent, Graph, Conversation, parallel, KnowledgeStore, START, END
from purpose_agent.tools import (
    Tool, FunctionTool, CalculatorTool, PythonExecTool,
    ReadFileTool, WriteFileTool, ToolRegistry,
)
from purpose_agent.llm_backend import LLMBackend, MockLLMBackend, ChatMessage
from purpose_agent.types import State
from purpose_agent.orchestrator import TaskResult

logger = logging.getLogger(__name__)


# ═══════════════════════════════════════════════════════════════════════════
# LEVEL 1 — The purpose() function. ONE call does everything.
# ═══════════════════════════════════════════════════════════════════════════

def purpose(
    description: str,
    model: str | LLMBackend | None = None,
    local: bool = True,
    knowledge: list[str] | str | None = None,
    tools: list[Tool] | None = None,
    interactive: bool = False,
) -> "Team":
    """
    Build a complete self-improving agent team from a plain English description.

    This is the entry point for everyone. No technical knowledge required.

    Args:
        description: What you want the team to do, in plain English.
                     e.g. "Research and summarize scientific papers"
                     e.g. "Write Python code and test it"
                     e.g. "Analyze CSV files and create reports"

        model: (optional) Which AI model to use.
               - None → auto-detects (Ollama if installed, else mock for testing)
               - "qwen3:1.7b" → local model via Ollama (free, private)
               - "gpt-4o" → OpenAI (needs OPENAI_API_KEY)
               - "Qwen/Qwen3-32B" → HuggingFace cloud (needs HF_TOKEN)

        local: If True (default), prefer local models. Zero cost, full privacy.

        knowledge: (optional) Give your agents knowledge.
                   - List of strings: ["fact 1", "fact 2", ...]
                   - File path: "./docs" or "./data.txt"

        tools: (optional) Extra tools for the agents to use.

        interactive: If True, agents ask for your approval before acting.

    Returns:
        A Team you can run tasks on. It gets smarter with every task.

    Examples:
        # Simplest possible usage
        team = purpose("Help me with coding tasks")
        result = team.run("Write a function to sort a list")
        print(result)

        # With local SLM (free, private)
        team = purpose("Research assistant", model="qwen3:1.7b")
        result = team.run("What are the latest trends in AI?")

        # With knowledge base
        team = purpose("Answer questions about my docs", knowledge="./my_docs/")
        result = team.run("What is our refund policy?")

        # Interactive mode (approve each action)
        team = purpose("File organizer", interactive=True)
        result = team.run("Organize my downloads folder")
    """
    # Auto-detect the best model
    resolved_model = _auto_detect_model(model, local)

    # Analyze the purpose description to pick the right team template
    template = _analyze_purpose(description)

    # Build tools list
    all_tools = list(tools or [])

    # Add template-specific tools
    for tool_cls in template["tools"]:
        all_tools.append(tool_cls())

    # Add knowledge store if provided
    kb = None
    if knowledge:
        kb = _build_knowledge_store(knowledge)
        all_tools.append(kb.as_tool())

    # Build the team
    team = Team(
        purpose=description,
        agents=template["agents"],
        model=resolved_model,
        tools=all_tools,
        interactive=interactive,
        knowledge=kb,
    )

    logger.info(f"✅ Team created: {template['name']} ({len(template['agents'])} agents)")
    return team


# ═══════════════════════════════════════════════════════════════════════════
# LEVEL 2 — Team class. Customizable but still simple.
# ═══════════════════════════════════════════════════════════════════════════

class Team:
    """
    A self-improving team of AI agents.

    Created automatically by purpose(), or build your own:

        team = Team.build(
            purpose="code review assistant",
            agents=["researcher", "coder", "reviewer"],
            model="qwen3:1.7b",
        )
        result = team.run("Review this pull request: ...")
    """

    def __init__(
        self,
        purpose: str,
        agents: list[dict[str, str]],
        model: str | LLMBackend | None = None,
        tools: list[Tool] | None = None,
        interactive: bool = False,
        knowledge: KnowledgeStore | None = None,
    ):
        self.purpose = purpose
        self.interactive = interactive
        self.knowledge = knowledge
        self._history: list[dict] = []

        # Create Agent instances
        self._agents: list[Agent] = []
        first_agent = None
        for spec in agents:
            agent = Agent(
                name=spec["name"],
                instructions=spec.get("role", ""),
                model=model,
                tools=tools,
                handoff_from=first_agent,  # Shared learning!
            )
            self._agents.append(agent)
            if first_agent is None:
                first_agent = agent

        # Create conversation for multi-agent collaboration
        if len(self._agents) > 1:
            self._conversation = Conversation(self._agents)
        else:
            self._conversation = None

    def run(self, task: str, verbose: bool = True) -> str:
        """
        Run a task. Returns a human-readable result string.

        The team gets smarter with every task you give it.

        Args:
            task: What you want done, in plain English.
            verbose: If True, print progress as it happens.

        Returns:
            The result as a simple string.
        """
        if verbose:
            print(f"\n🚀 Working on: {task}")
            print(f"   Team: {', '.join(a.name for a in self._agents)}")
            print(f"   Purpose: {self.purpose}")
            print()

        # Single agent → direct run
        if len(self._agents) == 1:
            result = self._agents[0].run(task)
            output = self._format_result(result)
        else:
            # Multi-agent → conversation
            conv_result = self._conversation.run(
                topic=task,
                rounds=2,
                initial_context=f"Team purpose: {self.purpose}",
            )
            output = conv_result.summary or str(conv_result.data.get("final_summary", ""))

        self._history.append({"task": task, "result": output[:500]})

        if verbose:
            print(f"\n✅ Done!")
            print(f"   Tasks completed: {len(self._history)}")
            print(f"   (The team learns from each task — it gets better over time)")

        return output

    def ask(self, question: str) -> str:
        """Shorthand for run() — more natural for Q&A use cases."""
        return self.run(question, verbose=False)

    def teach(self, lesson: str) -> None:
        """
        Teach the team something. This goes directly into their memory.

        Example:
            team.teach("Always cite your sources")
            team.teach("When writing code, add docstrings to every function")
        """
        for agent in self._agents:
            from purpose_agent.types import Heuristic, MemoryTier
            h = Heuristic(
                pattern="Always",
                strategy=lesson,
                steps=[],
                tier=MemoryTier.STRATEGIC,
                q_value=1.0,
            )
            agent.orch.optimizer.heuristic_library.append(h)
            agent.orch.sync_memory()
        print(f"📝 Taught all {len(self._agents)} agents: \"{lesson}\"")

    def status(self) -> str:
        """Show what the team has learned."""
        lines = [f"🧠 Team Status: {self.purpose}", ""]

        # Agents
        for agent in self._agents:
            n_heuristics = len(agent.orch.optimizer.heuristic_library)
            n_experiences = agent.orch.experience_replay.size
            lines.append(f"  🤖 {agent.name}: {n_heuristics} lessons learned, {n_experiences} experiences")

        # History
        lines.append(f"\n  📋 Tasks completed: {len(self._history)}")
        for i, h in enumerate(self._history[-5:], 1):
            lines.append(f"     {i}. {h['task'][:60]}")

        return "\n".join(lines)

    @staticmethod
    def _format_result(result: TaskResult) -> str:
        """Convert a TaskResult into a readable string."""
        data = result.final_state.data
        # Try to get the most useful output
        for key in ["_last_result", "_result", "result", "output", "answer"]:
            if key in data and data[key]:
                return str(data[key])
        if result.final_state.summary:
            return result.final_state.summary
        return str(data)

    @classmethod
    def build(
        cls,
        purpose: str,
        agents: list[str] | list[dict],
        model: str | LLMBackend | None = None,
        tools: list[Tool] | None = None,
    ) -> "Team":
        """
        Build a custom team with named agents.

        Args:
            purpose: What the team does.
            agents: List of agent names or {"name": ..., "role": ...} dicts.
            model: AI model to use.
            tools: Tools available to all agents.

        Example:
            team = Team.build(
                purpose="Content creation",
                agents=["writer", "editor", "fact_checker"],
                model="qwen3:1.7b",
            )
        """
        agent_specs = []
        for a in agents:
            if isinstance(a, str):
                agent_specs.append({"name": a, "role": f"You are the {a}."})
            else:
                agent_specs.append(a)
        return cls(purpose=purpose, agents=agent_specs, model=model, tools=tools)


# ═══════════════════════════════════════════════════════════════════════════
# Auto-Detection & Templates
# ═══════════════════════════════════════════════════════════════════════════

# Pre-built team templates matched by keywords in the purpose description
TEAM_TEMPLATES = {
    "research": {
        "name": "Research Team",
        "keywords": ["research", "find", "search", "discover", "learn", "papers", "study", "investigate", "summarize", "analyze information"],
        "agents": [
            {"name": "researcher", "role": "Find and gather relevant information. Be thorough and cite sources."},
            {"name": "analyst", "role": "Analyze the gathered information. Identify patterns, draw conclusions, and summarize findings clearly."},
        ],
        "tools": [CalculatorTool],
    },
    "coding": {
        "name": "Coding Team",
        "keywords": ["code", "program", "develop", "build", "software", "python", "javascript", "debug", "fix bug", "function", "api", "script"],
        "agents": [
            {"name": "architect", "role": "Design the solution. Break the problem into clear steps before coding."},
            {"name": "coder", "role": "Write clean, well-documented code. Include error handling and comments."},
            {"name": "tester", "role": "Review the code for bugs, edge cases, and improvements. Suggest fixes."},
        ],
        "tools": [PythonExecTool, CalculatorTool],
    },
    "writing": {
        "name": "Writing Team",
        "keywords": ["write", "blog", "article", "essay", "content", "copy", "draft", "edit", "proofread", "report", "documentation"],
        "agents": [
            {"name": "writer", "role": "Write clear, engaging content. Focus on the reader's needs."},
            {"name": "editor", "role": "Review and improve the writing. Fix grammar, clarity, and flow. Be constructive."},
        ],
        "tools": [],
    },
    "data": {
        "name": "Data Team",
        "keywords": ["data", "csv", "excel", "spreadsheet", "database", "sql", "chart", "graph", "statistics", "analytics", "dashboard"],
        "agents": [
            {"name": "analyst", "role": "Analyze data, find patterns, and compute statistics."},
            {"name": "reporter", "role": "Present findings in clear, non-technical language with key takeaways."},
        ],
        "tools": [PythonExecTool, CalculatorTool, ReadFileTool],
    },
    "assistant": {
        "name": "General Assistant",
        "keywords": ["help", "assist", "answer", "question", "explain", "general", "task", "do"],
        "agents": [
            {"name": "assistant", "role": "Help the user with their request. Be helpful, clear, and thorough."},
        ],
        "tools": [CalculatorTool],
    },
}


def _analyze_purpose(description: str) -> dict:
    """Match a purpose description to the best team template."""
    desc_lower = description.lower()
    best_template = None
    best_score = 0

    for template_key, template in TEAM_TEMPLATES.items():
        score = 0
        for keyword in template["keywords"]:
            if keyword in desc_lower:
                score += 1
                # Bonus for exact match
                if f" {keyword} " in f" {desc_lower} ":
                    score += 0.5
        if score > best_score:
            best_score = score
            best_template = template

    # Default to general assistant
    if best_template is None or best_score < 0.5:
        best_template = TEAM_TEMPLATES["assistant"]

    return best_template


def _auto_detect_model(model: str | LLMBackend | None, prefer_local: bool) -> str | LLMBackend:
    """Auto-detect the best available model."""
    if model is not None:
        return model

    # Check for Ollama
    if prefer_local:
        try:
            import urllib.request
            urllib.request.urlopen("http://localhost:11434/api/tags", timeout=2)
            logger.info("🟢 Ollama detected — using local models (free, private)")
            return "qwen3:1.7b"
        except Exception:
            pass

    # Check for API keys
    if os.environ.get("OPENAI_API_KEY"):
        logger.info("🔑 OpenAI API key found — using gpt-4o-mini")
        return "gpt-4o-mini"

    # Fallback: mock for testing
    logger.info(
        "💡 No local model detected. Using mock backend for testing.\n"
        "   To use a real model:\n"
        "   • Install Ollama: https://ollama.ai (free, local, private)\n"
        "   • Or set OPENAI_API_KEY for OpenAI\n"
        "   • Or set HF_TOKEN for HuggingFace"
    )
    return MockLLMBackend()


def _build_knowledge_store(knowledge: list[str] | str) -> KnowledgeStore:
    """Build a KnowledgeStore from various input types."""
    if isinstance(knowledge, list):
        return KnowledgeStore.from_texts(knowledge)
    elif os.path.isdir(knowledge):
        return KnowledgeStore.from_directory(knowledge, glob="*.*")
    elif os.path.isfile(knowledge):
        kb = KnowledgeStore()
        kb.add_file(knowledge)
        return kb
    else:
        # Treat as a single text
        return KnowledgeStore.from_texts([knowledge])


# ═══════════════════════════════════════════════════════════════════════════
# CLI Quickstart Wizard
# ═══════════════════════════════════════════════════════════════════════════

def quickstart():
    """
    Interactive wizard for creating an agent team. Run from command line:

        python -m purpose_agent

    Walks the user through setup step by step.
    """
    print()
    print("╔══════════════════════════════════════════════════════════╗")
    print("║        🧠 Purpose Agent — Quickstart Wizard            ║")
    print("║  Build a self-improving AI team in under 60 seconds.   ║")
    print("╚══════════════════════════════════════════════════════════╝")
    print()

    # Step 1: What's your purpose?
    print("Step 1: What do you want your AI team to do?")
    print("  Examples: 'research assistant', 'code helper', 'content writer'")
    print()
    user_purpose = input("  Your purpose: ").strip()
    if not user_purpose:
        user_purpose = "general assistant"
    print()

    # Step 2: Model selection
    print("Step 2: Which AI model? (press Enter for auto-detect)")
    print("  • Enter     → auto-detect (recommended)")
    print("  • 'local'   → use Ollama (free, private)")
    print("  • 'cloud'   → use HuggingFace cloud")
    print("  • 'openai'  → use OpenAI")
    print()
    model_choice = input("  Model: ").strip().lower()

    if model_choice == "local":
        model = "qwen3:1.7b"
    elif model_choice == "cloud":
        model = "Qwen/Qwen3-32B"
    elif model_choice == "openai":
        model = "gpt-4o-mini"
    elif model_choice:
        model = model_choice
    else:
        model = None  # auto-detect
    print()

    # Step 3: Knowledge?
    print("Step 3: Do you have documents for your team to learn from?")
    print("  • Enter         → no documents")
    print("  • folder path   → load all files from folder")
    print("  • file path     → load a specific file")
    print()
    knowledge_input = input("  Documents: ").strip()
    knowledge = knowledge_input if knowledge_input else None
    print()

    # Build!
    print("━" * 50)
    print("Building your team...")
    print()

    team = purpose(user_purpose, model=model, knowledge=knowledge)

    print()
    print("✅ Your team is ready!")
    print()
    print("Try it now — type a task (or 'quit' to exit):")
    print()

    while True:
        try:
            task = input("📝 Task: ").strip()
        except (EOFError, KeyboardInterrupt):
            print("\n👋 Goodbye!")
            break

        if not task or task.lower() in ("quit", "exit", "q"):
            print("\n👋 Goodbye!")
            break

        if task.lower() == "status":
            print(team.status())
            continue

        if task.lower().startswith("teach:"):
            lesson = task[6:].strip()
            team.teach(lesson)
            continue

        result = team.run(task)
        print(f"\n💡 Result:\n{result}\n")