"""Tool macros for reproducible research workflows.

Tool macros are named sequences of tool calls that represent standard
research patterns. Using macros ensures:
- Reproducibility: Same macro always runs same tools
- Team alignment: Macros map to team workflows
- Efficiency: Optimized tool sequences

Each macro defines:
- name: Identifier for the macro
- purpose: What it investigates
- when: When to use this macro
- steps: Ordered list of tool calls with argument templates
- expected_outputs: What findings to expect
"""

from dataclasses import dataclass, field
from typing import Any, Optional
import uuid

from emdash.agent.toolkit import AgentToolkit
from emdash.agent.research.state import EvidenceItem


@dataclass
class MacroStep:
    """A single step in a tool macro."""
    tool: str
    args_template: dict
    description: str = ""


@dataclass
class ToolMacro:
    """A named sequence of tool calls."""
    name: str
    purpose: str
    when: str
    steps: list[MacroStep]
    expected_outputs: list[str]


# Standard tool macros aligned with team workflows
TOOL_MACROS: dict[str, ToolMacro] = {
    "deep_feature_analysis": ToolMacro(
        name="deep_feature_analysis",
        purpose="Understand feature behavior + impact",
        when="Investigating how a feature works end-to-end",
        steps=[
            MacroStep(
                tool="semantic_search",
                args_template={"query": "{topic}", "limit": 10},
                description="Find entry points related to the feature",
            ),
            MacroStep(
                tool="expand_node",
                args_template={"node_type": "FROM_SEARCH", "identifier": "FROM_SEARCH"},
                description="Expand the most relevant result",
            ),
            MacroStep(
                tool="get_callers",
                args_template={"qualified_name": "FROM_EXPAND"},
                description="Find what calls this code",
            ),
            MacroStep(
                tool="get_callees",
                args_template={"qualified_name": "FROM_EXPAND"},
                description="Find what this code calls",
            ),
            MacroStep(
                tool="get_impact_analysis",
                args_template={"file_path": "FROM_EXPAND"},
                description="Assess change impact",
            ),
            MacroStep(
                tool="get_file_history",
                args_template={"file_path": "FROM_EXPAND"},
                description="Get recent changes and authors",
            ),
        ],
        expected_outputs=[
            "main entry points",
            "dependency edges",
            "changed hotspots",
            "code owners",
        ],
    ),
    
    "team_activity_analysis": ToolMacro(
        name="team_activity_analysis",
        purpose="Identify owners, expertise, velocity risks",
        when="Understanding who owns what and review bottlenecks",
        steps=[
            MacroStep(
                tool="get_recent_prs",
                args_template={"limit": 20},
                description="Get recent PRs",
            ),
            MacroStep(
                tool="get_pr_details",
                args_template={"pr_number": "FROM_PRS"},
                description="Get details of relevant PRs",
            ),
            MacroStep(
                tool="get_author_expertise",
                args_template={"author": "FROM_PRS"},
                description="Get author expertise areas",
            ),
            MacroStep(
                tool="detect_knowledge_silos",
                args_template={},
                description="Identify knowledge concentration risks",
            ),
        ],
        expected_outputs=[
            "who touched what",
            "review bottlenecks",
            "ownership gaps",
            "expertise distribution",
        ],
    ),
    
    "architectural_deep_dive": ToolMacro(
        name="architectural_deep_dive",
        purpose="Map architecture and key modules",
        when="Understanding system structure and coupling",
        steps=[
            MacroStep(
                tool="get_communities",
                args_template={"limit": 10, "include_members": True},
                description="Find code communities/clusters",
            ),
            MacroStep(
                tool="get_top_pagerank",
                args_template={"limit": 20},
                description="Find most important code entities",
            ),
            MacroStep(
                tool="expand_node",
                args_template={"node_type": "FROM_PAGERANK", "identifier": "FROM_PAGERANK"},
                description="Expand top entities",
            ),
            MacroStep(
                tool="get_file_dependencies",
                args_template={"file_path": "FROM_PAGERANK"},
                description="Get dependency graph",
            ),
        ],
        expected_outputs=[
            "module clusters",
            "core files",
            "coupling risks",
            "architectural boundaries",
        ],
    ),
    
    "implementation_trace": ToolMacro(
        name="implementation_trace",
        purpose="Trace how a specific feature is implemented",
        when="Need to understand exact implementation path",
        steps=[
            MacroStep(
                tool="text_search",
                args_template={"query": "{symbol}", "limit": 10},
                description="Find exact symbol/function",
            ),
            MacroStep(
                tool="expand_node",
                args_template={"node_type": "FROM_SEARCH", "identifier": "FROM_SEARCH"},
                description="Get full context",
            ),
            MacroStep(
                tool="get_callees",
                args_template={"qualified_name": "FROM_EXPAND"},
                description="Trace call chain",
            ),
            MacroStep(
                tool="get_class_hierarchy",
                args_template={"class_name": "FROM_EXPAND"},
                description="Get class hierarchy if applicable",
            ),
        ],
        expected_outputs=[
            "implementation location",
            "call chain",
            "class hierarchy",
            "dependencies",
        ],
    ),
    
    "risk_assessment": ToolMacro(
        name="risk_assessment",
        purpose="Assess risks of modifying code",
        when="Before making changes to understand impact",
        steps=[
            MacroStep(
                tool="semantic_search",
                args_template={"query": "{topic}", "limit": 5},
                description="Find relevant code",
            ),
            MacroStep(
                tool="get_impact_analysis",
                args_template={"file_path": "FROM_SEARCH"},
                description="Get impact analysis",
            ),
            MacroStep(
                tool="get_callers",
                args_template={"qualified_name": "FROM_SEARCH"},
                description="Find all callers",
            ),
            MacroStep(
                tool="detect_knowledge_silos",
                args_template={},
                description="Check for ownership risks",
            ),
            MacroStep(
                tool="get_file_history",
                args_template={"file_path": "FROM_SEARCH", "limit": 30},
                description="Get change history",
            ),
        ],
        expected_outputs=[
            "affected components",
            "risk level",
            "ownership risks",
            "change frequency",
        ],
    ),
    
    "pr_context": ToolMacro(
        name="pr_context",
        purpose="Understand context around a PR or change",
        when="Reviewing or investigating a specific PR",
        steps=[
            MacroStep(
                tool="get_pr_details",
                args_template={"pr_number": "{pr_number}"},
                description="Get PR details",
            ),
            MacroStep(
                tool="find_similar_prs",
                args_template={"description": "FROM_PR"},
                description="Find related PRs",
            ),
            MacroStep(
                tool="get_author_expertise",
                args_template={"author": "FROM_PR"},
                description="Get author expertise",
            ),
        ],
        expected_outputs=[
            "PR details",
            "related changes",
            "author context",
        ],
    ),

    # GitHub MCP-powered macros for enhanced research
    # These macros use FROM_REPO_* to auto-detect the repo from the graph
    "github_pr_deep_dive": ToolMacro(
        name="github_pr_deep_dive",
        purpose="Deep analysis of a PR with code context from GitHub",
        when="Investigating a PR with diffs, comments, and reviews",
        steps=[
            MacroStep(
                tool="github_pr_details",
                args_template={
                    "owner": "FROM_REPO_OWNER",  # Auto-detect from graph
                    "repo": "FROM_REPO_NAME",    # Auto-detect from graph
                    "pull_number": "FROM_PRS",   # Get from get_recent_prs
                    "include_diff": True,
                    "include_comments": True,
                    "include_reviews": True,
                },
                description="Get full PR details with diff and comments",
            ),
            MacroStep(
                tool="expand_node",
                args_template={"node_type": "File", "identifier": "FROM_PR_FILES"},
                description="Expand affected files in the graph",
            ),
            MacroStep(
                tool="get_impact_analysis",
                args_template={"file_path": "FROM_PR_FILES"},
                description="Assess impact of PR changes",
            ),
        ],
        expected_outputs=[
            "PR rationale and context",
            "code changes (diff)",
            "reviewer comments",
            "affected components",
        ],
    ),

    "github_code_research": ToolMacro(
        name="github_code_research",
        purpose="Research code patterns in the current repository on GitHub",
        when="Finding implementations, usage examples, or code patterns in this repo",
        steps=[
            MacroStep(
                tool="github_search_code",
                args_template={
                    "query": "{topic}",
                    "owner": "FROM_REPO_OWNER",  # Search within current repo
                    "repo": "FROM_REPO_NAME",
                    "per_page": 15,
                },
                description="Search for code patterns in repository",
            ),
            MacroStep(
                tool="github_get_file_content",
                args_template={
                    "owner": "FROM_REPO_OWNER",
                    "repo": "FROM_REPO_NAME",
                    "path": "FROM_GITHUB_SEARCH",
                },
                description="Get full file content from search results",
            ),
        ],
        expected_outputs=[
            "matching code snippets",
            "implementation patterns",
            "usage examples",
        ],
    ),

    "github_repo_exploration": ToolMacro(
        name="github_repo_exploration",
        purpose="Explore the current repository structure and key files",
        when="Understanding repository layout and current development",
        steps=[
            MacroStep(
                tool="github_view_repo_structure",
                args_template={
                    "owner": "FROM_REPO_OWNER",  # Auto-detect from graph
                    "repo": "FROM_REPO_NAME",
                    "path": "",
                },
                description="View repository root structure",
            ),
            MacroStep(
                tool="github_get_file_content",
                args_template={
                    "owner": "FROM_REPO_OWNER",
                    "repo": "FROM_REPO_NAME", 
                    "path": "README.md",
                },
                description="Get README for context",
            ),
            MacroStep(
                tool="github_list_prs",
                args_template={
                    "owner": "FROM_REPO_OWNER",
                    "repo": "FROM_REPO_NAME",
                    "state": "open",
                    "per_page": 10,
                },
                description="See current development activity",
            ),
        ],
        expected_outputs=[
            "repository structure",
            "project overview",
            "active development",
        ],
    ),
}


class MacroExecutor:
    """Executes tool macros and collects evidence.
    
    The executor handles:
    - Template resolution (FROM_SEARCH, FROM_EXPAND, etc.)
    - Evidence collection with unique IDs
    - Budget tracking
    - Error handling
    """
    
    def __init__(self, toolkit: AgentToolkit, verbose: bool = False):
        """Initialize with an agent toolkit.
        
        Args:
            toolkit: The AgentToolkit instance for tool execution
            verbose: Whether to print debug info
        """
        self.toolkit = toolkit
        self.evidence_counter = 0
        self._last_results: dict[str, Any] = {}
        self.verbose = verbose
    
    def seed_context(self, context: dict) -> None:
        """Seed the executor with prior context.
        
        This allows the executor to use results from prior
        searches when resolving FROM_* dependencies.
        
        Args:
            context: Context dict with keys like 'last_search_results', 'repo_owner', 'repo_name'
        """
        # Convert context to _last_results format
        if "last_search_results" in context:
            self._last_results["semantic_search"] = {
                "results": context["last_search_results"]
            }
            self._last_results["text_search"] = {
                "results": context["last_search_results"]
            }
        
        if "last_search_top" in context:
            # Also seed expand_node if we have a top result
            top = context["last_search_top"]
            self._last_results["expand_node"] = {
                "root_node": top
            }
        
        # Store repository context for FROM_REPO_* dependencies
        # This ensures macros always know which repo they're working with
        self._last_results["_context"] = {
            "repo_owner": context.get("repo_owner"),
            "repo_name": context.get("repo_name"),
            "github_url": context.get("github_url"),
        }
    
    def execute_macro(
        self,
        macro_name: str,
        params: dict,
        budget_remaining: int = 100,
        prior_context: Optional[dict] = None,
    ) -> tuple[list[EvidenceItem], dict]:
        """Execute a tool macro and collect evidence.
        
        Args:
            macro_name: Name of the macro to execute
            params: Parameters to fill in templates (topic, symbol, pr_number, etc.)
            budget_remaining: Remaining tool call budget
            prior_context: Optional prior context to seed dependencies
            
        Returns:
            Tuple of (evidence_items, context_updates)
            
        Raises:
            ValueError: If macro_name is not found
        """
        if macro_name not in TOOL_MACROS:
            raise ValueError(f"Unknown macro: {macro_name}. Available: {list(TOOL_MACROS.keys())}")
        
        macro = TOOL_MACROS[macro_name]
        evidence: list[EvidenceItem] = []
        context_updates: dict[str, Any] = {}
        
        # Reset last results but seed with prior context
        self._last_results = {}
        if prior_context:
            self.seed_context(prior_context)
        
        steps_executed = 0
        steps_skipped = 0
        
        for step in macro.steps:
            if budget_remaining <= 0:
                break
            
            # Resolve arguments
            args = self._resolve_args(step.args_template, params)
            if args is None:
                # Skip step if dependencies not met
                steps_skipped += 1
                continue
            
            # Execute tool
            try:
                result = self.toolkit.execute(step.tool, **args)
                budget_remaining -= 1
                steps_executed += 1
                
                # Create evidence item
                evidence_item = self._create_evidence(
                    tool=step.tool,
                    input_args=args,
                    result=result,
                    description=step.description,
                )
                evidence.append(evidence_item)
                
                # Store result for dependency resolution
                self._store_result(step.tool, result)
                
                # Update context with key findings
                if result.success:
                    context_updates.update(self._extract_context(step.tool, result))
                    
            except Exception as e:
                # Log but continue with other steps
                steps_skipped += 1
        
        return evidence, context_updates
    
    def _resolve_args(self, template: dict, params: dict) -> Optional[dict]:
        """Resolve argument templates.
        
        Handles:
        - Direct parameter substitution ({topic}, {symbol})
        - Dependency resolution (FROM_SEARCH, FROM_EXPAND, etc.)
        
        Args:
            template: The argument template
            params: Parameters for substitution
            
        Returns:
            Resolved arguments or None if dependencies not met
        """
        resolved = {}
        
        for key, value in template.items():
            if isinstance(value, str):
                if value.startswith("{") and value.endswith("}"):
                    # Direct parameter substitution
                    param_name = value[1:-1]
                    if param_name in params:
                        resolved[key] = params[param_name]
                    else:
                        return None  # Missing required parameter
                elif value.startswith("FROM_"):
                    # Dependency resolution
                    resolved_value = self._resolve_dependency(value)
                    if resolved_value is None:
                        return None  # Dependency not available
                    resolved[key] = resolved_value
                else:
                    resolved[key] = value
            else:
                resolved[key] = value
        
        return resolved
    
    def _resolve_dependency(self, dep_key: str) -> Optional[str]:
        """Resolve a FROM_* dependency.
        
        Args:
            dep_key: Dependency key (FROM_SEARCH, FROM_EXPAND, etc.)
            
        Returns:
            Resolved value or None
        """
        if dep_key == "FROM_SEARCH":
            # Get first result from semantic_search or text_search
            for tool in ["semantic_search", "text_search"]:
                if tool in self._last_results:
                    results = self._last_results[tool].get("results", [])
                    if results:
                        return results[0].get("qualified_name") or results[0].get("identifier")
            return None
        
        elif dep_key == "FROM_EXPAND":
            # Get qualified_name or file_path from expand_node
            if "expand_node" in self._last_results:
                data = self._last_results["expand_node"]
                root = data.get("root_node", {})
                return root.get("qualified_name") or root.get("file_path")
            return None
        
        elif dep_key == "FROM_PAGERANK":
            # Get top result from get_top_pagerank
            if "get_top_pagerank" in self._last_results:
                results = self._last_results["get_top_pagerank"].get("results", [])
                if results:
                    return results[0].get("qualified_name") or results[0].get("identifier")
            return None
        
        elif dep_key == "FROM_PRS":
            # Get first PR number from get_recent_prs
            if "get_recent_prs" in self._last_results:
                results = self._last_results["get_recent_prs"].get("results", [])
                if results:
                    return results[0].get("pr_number") or results[0].get("author")
            return None
        
        elif dep_key == "FROM_PR":
            # Get PR description or author from get_pr_details
            if "get_pr_details" in self._last_results:
                data = self._last_results["get_pr_details"]
                return data.get("description") or data.get("title") or data.get("author")
            return None
        
        elif dep_key == "FROM_REPO_OWNER":
            # Get repository owner from context (set by researcher from graph)
            if "_context" in self._last_results:
                return self._last_results["_context"].get("repo_owner")
            return None
        
        elif dep_key == "FROM_REPO_NAME":
            # Get repository name from context (set by researcher from graph)
            if "_context" in self._last_results:
                return self._last_results["_context"].get("repo_name")
            return None
        
        elif dep_key == "FROM_GITHUB_SEARCH":
            # Get file path from github_search_code results
            if "github_search_code" in self._last_results:
                results = self._last_results["github_search_code"].get("results", [])
                if results and len(results) > 0:
                    return results[0].get("path")
            return None
        
        elif dep_key == "FROM_PR_FILES":
            # Get first file from PR details
            if "github_pr_details" in self._last_results:
                data = self._last_results["github_pr_details"]
                files = data.get("files", [])
                if files:
                    return files[0].get("filename") or files[0].get("path")
            return None
        
        return None
    
    def _store_result(self, tool: str, result) -> None:
        """Store tool result for dependency resolution.
        
        Args:
            tool: Tool name
            result: Tool execution result
        """
        if result.success:
            self._last_results[tool] = result.data
    
    def _create_evidence(
        self,
        tool: str,
        input_args: dict,
        result,
        description: str,
    ) -> EvidenceItem:
        """Create an evidence item from a tool result.
        
        Args:
            tool: Tool name
            input_args: Arguments passed to the tool
            result: Tool execution result
            description: Step description
            
        Returns:
            EvidenceItem with unique ID
        """
        self.evidence_counter += 1
        evidence_id = f"E{self.evidence_counter}"
        
        # Extract entities from result
        entities = self._extract_entities(result)
        
        # Create summary
        if result.success:
            summary = self._summarize_result(tool, result.data, description)
        else:
            summary = f"Error: {result.error}"
        
        return EvidenceItem(
            id=evidence_id,
            tool=tool,
            input=input_args,
            output_ref=f"result_{evidence_id}",  # Would store actual output in practice
            summary=summary,
            entities=entities,
        )
    
    def _extract_entities(self, result) -> list[str]:
        """Extract entity identifiers from a tool result.
        
        Args:
            result: Tool execution result
            
        Returns:
            List of entity identifiers (file paths, qualified names, PR IDs)
        """
        entities = []
        
        if not result.success:
            return entities
        
        data = result.data
        
        # Extract from results list
        if "results" in data:
            for item in data["results"][:10]:  # Limit
                if isinstance(item, dict):
                    for key in ["qualified_name", "file_path", "identifier", "pr_number"]:
                        if key in item and item[key]:
                            entities.append(str(item[key]))
                            break
        
        # Extract from root_node
        if "root_node" in data:
            root = data["root_node"]
            for key in ["qualified_name", "file_path"]:
                if key in root and root[key]:
                    entities.append(str(root[key]))
        
        # Extract from files
        if "files" in data:
            for f in data["files"][:10]:
                if isinstance(f, dict) and "path" in f:
                    entities.append(f["path"])
                elif isinstance(f, str):
                    entities.append(f)
        
        return list(set(entities))  # Dedupe
    
    def _summarize_result(self, tool: str, data: dict, description: str) -> str:
        """Create a human-readable summary of a tool result.
        
        Args:
            tool: Tool name
            data: Result data
            description: Step description
            
        Returns:
            1-3 line summary
        """
        count = None
        
        if "results" in data:
            count = len(data["results"])
        elif "callers" in data:
            count = data.get("count", len(data["callers"]))
        elif "callees" in data:
            count = data.get("count", len(data["callees"]))
        elif "communities" in data:
            count = len(data["communities"])
        elif "files" in data:
            count = len(data["files"])
        
        if count is not None:
            return f"{description}: Found {count} results"
        elif "summary" in data:
            s = data["summary"]
            return f"{description}: {s.get('function_count', 0)} functions, {s.get('class_count', 0)} classes"
        else:
            return description
    
    def _extract_context(self, tool: str, result) -> dict:
        """Extract context updates from a tool result.
        
        Args:
            tool: Tool name
            result: Tool execution result
            
        Returns:
            Context updates dict
        """
        context = {}
        data = result.data
        
        if tool in ["semantic_search", "text_search"]:
            if "results" in data and data["results"]:
                context["last_search_results"] = data["results"][:5]
                context["last_search_top"] = data["results"][0]
        
        elif tool == "expand_node":
            if "root_node" in data:
                context["last_expanded_node"] = data["root_node"]
            if "summary" in data:
                context["last_expansion_summary"] = data["summary"]
        
        elif tool == "get_communities":
            if "communities" in data:
                context["communities"] = data["communities"]
        
        elif tool == "get_top_pagerank":
            if "results" in data:
                context["top_entities"] = data["results"][:10]
        
        elif tool == "get_recent_prs":
            if "results" in data:
                context["recent_prs"] = data["results"][:10]
        
        return context


def get_macro(name: str) -> Optional[ToolMacro]:
    """Get a tool macro by name.
    
    Args:
        name: Macro name
        
    Returns:
        ToolMacro or None if not found
    """
    return TOOL_MACROS.get(name)


def list_macros() -> list[dict]:
    """List all available macros.
    
    Returns:
        List of macro info dicts
    """
    return [
        {
            "name": macro.name,
            "purpose": macro.purpose,
            "when": macro.when,
            "steps": len(macro.steps),
            "expected_outputs": macro.expected_outputs,
        }
        for macro in TOOL_MACROS.values()
    ]


def suggest_macros(goal: str, include_github: bool = True) -> list[str]:
    """Suggest macros based on research goal.
    
    Args:
        goal: The research goal
        include_github: Whether to include GitHub MCP macros
        
    Returns:
        List of suggested macro names
    """
    goal_lower = goal.lower()
    suggestions = []
    
    # Feature/behavior understanding
    if any(w in goal_lower for w in ["feature", "how does", "understand", "works"]):
        suggestions.append("deep_feature_analysis")
    
    # Team/ownership questions
    if any(w in goal_lower for w in ["who", "owner", "team", "expertise", "author"]):
        suggestions.append("team_activity_analysis")
    
    # Architecture questions
    if any(w in goal_lower for w in ["architecture", "structure", "module", "design"]):
        suggestions.append("architectural_deep_dive")
    
    # Implementation details
    if any(w in goal_lower for w in ["implement", "code", "function", "class"]):
        suggestions.append("implementation_trace")
    
    # Risk assessment
    if any(w in goal_lower for w in ["risk", "impact", "change", "modify", "refactor"]):
        suggestions.append("risk_assessment")
    
    # PR context - prefer GitHub MCP version for richer data
    if any(w in goal_lower for w in ["pr", "pull request", "review", "diff", "comment"]):
        if include_github:
            suggestions.append("github_pr_deep_dive")
        suggestions.append("pr_context")
    
    # GitHub-specific research
    if include_github:
        # Code search / patterns
        if any(w in goal_lower for w in ["search", "find", "pattern", "example", "usage"]):
            suggestions.append("github_code_research")
        
        # Repository exploration
        if any(w in goal_lower for w in ["repo", "repository", "explore", "discover"]):
            suggestions.append("github_repo_exploration")
    
    # Default to feature analysis if no match
    if not suggestions:
        suggestions.append("deep_feature_analysis")
    
    return suggestions

