#!/usr/bin/env python3
"""
Prompt Augmentor - Enhanced prompt generation with context injection

Augments prompts with:
1. Project context (path, focus, tools, recent activity)
2. Loaded documentation and references
3. Known patterns and best practices
4. Tool availability and capabilities
5. Task history and continuation context
"""

import json
import os
from pathlib import Path
from typing import Dict, List, Any, Optional
from datetime import datetime
import re

class PromptAugmentor:
    """Augments prompts with rich context for better task understanding"""

    def __init__(self, project_config: Dict[str, Any],
                 tools_available: List[str] = None,
                 docs: Dict[str, str] = None):
        """Initialize augmentor with project context

        Args:
            project_config: Project configuration dict
            tools_available: List of available tool names
            docs: Dict of documentation by tool/topic
        """
        self.config = project_config
        self.tools = tools_available or []
        self.docs = docs or {}
        self.project_name = project_config.get("name", "unknown")
        self.project_path = project_config.get("path", "")
        self.project_focus = project_config.get("focus", "")

    def augment(self, prompt: str, task_context: Optional[Dict] = None) -> str:
        """Augment prompt with full context

        Args:
            prompt: Original task prompt
            task_context: Optional task context (history, state, etc)

        Returns:
            Augmented prompt with injected context
        """
        sections = [
            "# Task Context",
            self._system_context(),
            "",
            "# Task Instructions",
            self._task_instructions(),
            "",
            "# Available Resources",
            self._resource_documentation(),
        ]

        if task_context:
            sections.extend([
                "",
                "# Continuation Context",
                self._continuation_context(task_context)
            ])

        sections.extend([
            "",
            "# Original Task",
            prompt,
            "",
            "# Execution Guidelines",
            self._execution_guidelines()
        ])

        return "\n".join(sections)

    def _system_context(self) -> str:
        """Generate system context section"""
        return f"""You are working on the {self.project_name} project.

**Project Focus:** {self.project_focus}
**Working Directory:** {self.project_path}
**Timestamp:** {datetime.now().isoformat()}

Key Responsibilities:
- Execute tasks efficiently in project context
- Use available tools appropriately
- Reference documentation when relevant
- Report clear, structured results"""

    def _task_instructions(self) -> str:
        """Generate task execution instructions"""
        return """Before executing the task:
1. Understand what tools are available and when to use them
2. Check documentation for similar tasks or patterns
3. Consider task continuation context if provided
4. Plan execution steps mentally first
5. Execute with focus on clarity and correctness"""

    def _resource_documentation(self) -> str:
        """Generate documentation for available resources"""
        sections = []

        # Available tools
        if self.tools:
            sections.append("## Available Tools")
            for tool in self.tools:
                doc = self.docs.get(tool, f"See Tool Reference for {tool}")
                sections.append(f"- **{tool}**: {doc}")

        # Project-specific knowledge
        if self.config.get("knowledge"):
            sections.append("\n## Project Knowledge")
            for key, value in self.config.get("knowledge", {}).items():
                sections.append(f"- {key}: {value}")

        # Best practices
        sections.append("\n## Best Practices for This Project")
        best_practices = self._get_best_practices()
        for practice in best_practices:
            sections.append(f"- {practice}")

        return "\n".join(sections) if sections else "No additional documentation available"

    def _get_best_practices(self) -> List[str]:
        """Get project-specific best practices"""
        practices = {
            "admin": [
                "Always check system state before making changes",
                "Use systemctl for service management",
                "Backup state before maintenance",
                "Document all configuration changes"
            ],
            "overbits": [
                "Run type checking before commits",
                "Use npm scripts for builds",
                "Test components in isolation first",
                "Keep bundle size in check"
            ],
            "musica": [
                "Test audio patterns thoroughly",
                "Use Strudel patterns for consistency",
                "Verify performance with heavy patterns",
                "Document musical concepts clearly"
            ],
            "dss": [
                "Verify cryptographic operations independently",
                "Test edge cases for crypto functions",
                "Keep key material secure",
                "Validate all inputs strictly"
            ],
            "librechat": [
                "Test conversation flows end-to-end",
                "Monitor API rate limits",
                "Cache frequently used models",
                "Log conversation history appropriately"
            ],
            "bbot": [
                "Test trading logic with mock data first",
                "Implement circuit breakers for safety",
                "Monitor market conditions closely",
                "Log all trades with rationale"
            ]
        }
        return practices.get(self.project_name, [
            "Write clear, self-documenting code",
            "Test changes thoroughly before committing",
            "Keep commits atomic and focused",
            "Document significant decisions"
        ])

    def _continuation_context(self, task_context: Dict) -> str:
        """Generate continuation context from task history"""
        sections = []

        # Previous results
        if task_context.get("previous_results"):
            sections.append("## Previous Results")
            for key, value in task_context.get("previous_results", {}).items():
                # Truncate long values
                value_str = str(value)[:200]
                if len(str(value)) > 200:
                    value_str += "..."
                sections.append(f"- {key}: {value_str}")

        # Task state
        if task_context.get("state"):
            sections.append("\n## Current State")
            state = task_context.get("state", {})
            for key, value in state.items():
                sections.append(f"- {key}: {value}")

        # Blockers or issues
        if task_context.get("issues"):
            sections.append("\n## Known Issues to Address")
            for issue in task_context.get("issues", []):
                sections.append(f"- {issue}")

        # Next steps hint
        if task_context.get("next_steps"):
            sections.append("\n## Suggested Next Steps")
            for step in task_context.get("next_steps", []):
                sections.append(f"- {step}")

        return "\n".join(sections) if sections else "No previous context available"

    def _execution_guidelines(self) -> str:
        """Generate execution guidelines"""
        return """## How to Execute Effectively

1. **Start Clear**: Restate the task in your understanding
2. **Check Context**: Review continuation context if present
3. **Plan First**: Outline steps before executing
4. **Use Tools**: Choose appropriate tools for each step
5. **Document**: Provide clear output with reasoning
6. **Handle Errors**: Stop and report clearly if blocked
7. **Summarize**: Recap what was done and results

## Output Format

Always provide:
- What was completed
- Key findings or results
- Any errors encountered
- Recommendations for next steps (if applicable)
- Structured data when relevant (JSON/tables)"""

    def create_project_context_file(self, output_path: Path) -> None:
        """Create a JSON file with augmented context for reference

        Useful for keeping state about project knowledge.
        """
        context = {
            "project": self.project_name,
            "path": self.project_path,
            "focus": self.project_focus,
            "tools": self.tools,
            "generated_at": datetime.now().isoformat(),
            "best_practices": self._get_best_practices(),
            "knowledge": self.config.get("knowledge", {})
        }

        output_path.parent.mkdir(parents=True, exist_ok=True)
        output_path.write_text(json.dumps(context, indent=2))


class PromptTemplateBuilder:
    """Builds structured prompts for common task types"""

    @staticmethod
    def build_analysis_prompt(topic: str, context: str, focus_areas: List[str]) -> str:
        """Build analysis task prompt"""
        template = f"""# Analysis Task: {topic}

## Context
{context}

## Focus Areas
{chr(10).join(f"- {area}" for area in focus_areas)}

## Required Output
1. Summary of findings
2. Key patterns identified
3. Risk assessment (if applicable)
4. Recommendations
5. Next steps

Please provide clear, structured analysis."""
        return template

    @staticmethod
    def build_debug_prompt(issue: str, symptoms: str, relevant_files: List[str]) -> str:
        """Build debugging task prompt"""
        template = f"""# Debugging Task: {issue}

## Symptoms
{symptoms}

## Relevant Files
{chr(10).join(f"- {f}" for f in relevant_files)}

## Required Output
1. Root cause hypothesis
2. Evidence supporting hypothesis
3. Proposed fix
4. Testing approach
5. Prevention measures

Please debug systematically."""
        return template

    @staticmethod
    def build_implementation_prompt(feature: str, requirements: List[str],
                                   constraints: List[str]) -> str:
        """Build implementation task prompt"""
        template = f"""# Implementation Task: {feature}

## Requirements
{chr(10).join(f"- {req}" for req in requirements)}

## Constraints
{chr(10).join(f"- {c}" for c in constraints)}

## Required Output
1. Implementation plan
2. Code changes
3. Testing strategy
4. Documentation updates
5. Deployment considerations

Please implement systematically."""
        return template