#!/usr/bin/env python3
"""
Prompt Integration Framework

Integrates advanced prompt engineering techniques into Luzia task dispatch.

Usage:
    engine = PromptIntegrationEngine(project_config)
    augmented_prompt, metadata = engine.augment_for_task(
        original_task,
        task_type=TaskType.IMPLEMENTATION,
        domain="backend",
        complexity=3
    )
"""

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

from prompt_techniques import (
    TaskType, PromptStrategy, PromptContext, PromptEngineer,
    RoleBasedPrompting, ContextHierarchy, TaskSpecificPatterns
)


class DomainSpecificAugmentor:
    """Domain-specific prompt augmentation for different project types"""

    DOMAIN_CONTEXTS = {
        "backend": {
            "focus": "performance, scalability, and reliability",
            "priorities": ["Error handling", "Concurrency", "Resource efficiency", "Security"],
            "best_practices": [
                "Write defensive code that handles edge cases",
                "Consider performance implications of design choices",
                "Ensure thread-safety for concurrent operations",
                "Log meaningful debugging information",
                "Design for testability from the start"
            ]
        },
        "frontend": {
            "focus": "user experience, accessibility, and performance",
            "priorities": ["User experience", "Accessibility", "Performance", "Cross-browser"],
            "best_practices": [
                "Prioritize user experience and intuitive interactions",
                "Ensure accessibility (WCAG 2.1 AA minimum)",
                "Optimize for performance and smooth interactions",
                "Test on multiple browsers and devices",
                "Keep component logic simple and focused"
            ]
        },
        "devops": {
            "focus": "reliability, automation, and observability",
            "priorities": ["Reliability", "Automation", "Monitoring", "Documentation"],
            "best_practices": [
                "Design for high availability and graceful degradation",
                "Automate all repeatable processes",
                "Implement comprehensive monitoring and alerting",
                "Document operational procedures clearly",
                "Plan for disaster recovery and failover"
            ]
        },
        "crypto": {
            "focus": "correctness, security, and auditability",
            "priorities": ["Correctness", "Security", "Auditability", "Efficiency"],
            "best_practices": [
                "Verify all cryptographic claims independently",
                "Use well-tested libraries, avoid rolling custom crypto",
                "Implement constant-time operations for sensitive comparisons",
                "Document security assumptions explicitly",
                "Include comprehensive test coverage for edge cases"
            ]
        },
        "research": {
            "focus": "rigor, novelty, and reproducibility",
            "priorities": ["Correctness", "Novelty", "Reproducibility", "Clarity"],
            "best_practices": [
                "Clearly state hypotheses and test them systematically",
                "Provide sufficient detail for reproducibility",
                "Distinguish between established facts and speculation",
                "Compare against baselines and prior work",
                "Document all assumptions and limitations"
            ]
        },
        "orchestration": {
            "focus": "coordination, efficiency, and resilience",
            "priorities": ["Correctness", "Efficiency", "Resilience", "Observability"],
            "best_practices": [
                "Design for idempotency and safe retries",
                "Implement clear state transitions and monitoring",
                "Minimize orchestration overhead",
                "Handle failures gracefully with fallbacks",
                "Provide clear visibility into system state"
            ]
        }
    }

    @staticmethod
    def get_domain_context(domain: str) -> Dict[str, Any]:
        """Get domain-specific context"""
        return DomainSpecificAugmentor.DOMAIN_CONTEXTS.get(
            domain,
            DomainSpecificAugmentor.DOMAIN_CONTEXTS["backend"]  # Default
        )


class ComplexityAdaptivePrompting:
    """Adapts prompt complexity based on task complexity"""

    @staticmethod
    def estimate_complexity(task: str, task_type: TaskType) -> int:
        """
        Estimate task complexity from 1-5 scale.

        Heuristics:
        - Word count > 200: complexity += 1
        - Multiple components mentioned: complexity += 1
        - Edge cases mentioned: complexity += 1
        - Performance/security concerns: complexity += 1
        """
        complexity = 1
        task_lower = task.lower()

        # Word count
        if len(task.split()) > 200:
            complexity += 1

        # Multiple concerns
        concerns = ["concurrent", "parallel", "race", "deadlock", "performance",
                    "security", "encrypt", "scale", "distributed", "async"]
        if sum(1 for c in concerns if c in task_lower) >= 2:
            complexity += 1

        # Edge cases mentioned
        if any(word in task_lower for word in ["edge", "corner", "unusual", "error", "exception"]):
            complexity += 1

        # Architectural concerns
        if any(word in task_lower for word in ["architecture", "refactor", "redesign", "migration"]):
            complexity += 1

        return min(complexity, 5)

    @staticmethod
    def get_prompting_strategies(complexity: int) -> List[PromptStrategy]:
        """Get recommended strategies based on complexity"""
        strategy_levels = {
            1: [PromptStrategy.SYSTEM_INSTRUCTION, PromptStrategy.ROLE_BASED],
            2: [PromptStrategy.SYSTEM_INSTRUCTION, PromptStrategy.ROLE_BASED, PromptStrategy.CHAIN_OF_THOUGHT],
            3: [PromptStrategy.SYSTEM_INSTRUCTION, PromptStrategy.ROLE_BASED, PromptStrategy.CHAIN_OF_THOUGHT, PromptStrategy.FEW_SHOT],
            4: [PromptStrategy.SYSTEM_INSTRUCTION, PromptStrategy.ROLE_BASED, PromptStrategy.CHAIN_OF_THOUGHT,
                PromptStrategy.FEW_SHOT, PromptStrategy.TREE_OF_THOUGHT],
            5: [PromptStrategy.SYSTEM_INSTRUCTION, PromptStrategy.ROLE_BASED, PromptStrategy.CHAIN_OF_THOUGHT,
                PromptStrategy.FEW_SHOT, PromptStrategy.TREE_OF_THOUGHT, PromptStrategy.SELF_CONSISTENCY]
        }
        return strategy_levels.get(complexity, strategy_levels[1])


class PromptIntegrationEngine:
    """Main engine integrating all prompt augmentation techniques"""

    def __init__(self, project_config: Dict[str, Any] = None):
        """
        Initialize prompt integration engine.

        Args:
            project_config: Project configuration with name, path, focus, etc.
        """
        self.project_config = project_config or {}
        self.engineer = PromptEngineer()
        self.domain_augmentor = DomainSpecificAugmentor()
        self.complexity_adapter = ComplexityAdaptivePrompting()
        self.context_hierarchy = ContextHierarchy()

    def augment_for_task(self,
                        task: str,
                        task_type: TaskType,
                        domain: str = "backend",
                        complexity: Optional[int] = None,
                        context: Optional[Dict[str, Any]] = None,
                        strategies: Optional[List[PromptStrategy]] = None) -> tuple:
        """
        Augment a task prompt with advanced techniques.

        Args:
            task: Original task description
            task_type: Type of task (from TaskType enum)
            domain: Domain context (backend, frontend, crypto, etc.)
            complexity: Estimated complexity (1-5), auto-detected if None
            context: Additional context (project state, history, etc.)
            strategies: Specific strategies to use, auto-selected if None

        Returns:
            (augmented_prompt, metadata)
        """
        # Auto-detect complexity if not provided
        if complexity is None:
            complexity = self.complexity_adapter.estimate_complexity(task, task_type)

        # Auto-select strategies based on complexity
        if strategies is None:
            strategies = self.complexity_adapter.get_prompting_strategies(complexity)

        # Get domain context
        domain_context = self.domain_augmentor.get_domain_context(domain)

        # Build context object
        prompt_context = PromptContext(
            task_type=task_type,
            primary={
                "project": self.project_config.get("name", "unknown"),
                "domain": domain,
                "focus": domain_context["focus"]
            },
            secondary={
                "priorities": ", ".join(domain_context["priorities"]),
                "best_practices": "; ".join(domain_context["best_practices"][:3])
            },
            tertiary={
                "timestamp": datetime.now().isoformat(),
                "context_provided": "yes" if context else "no"
            },
            examples=[],
            complexity_level=complexity
        )

        # Engineer the prompt
        engineered_prompt, metadata = self.engineer.engineer_prompt(
            task,
            task_type,
            strategies=strategies,
            context=prompt_context
        )

        # Add domain-specific augmentation
        final_prompt = self._add_domain_augmentation(
            engineered_prompt,
            domain,
            domain_context
        )

        # Add project context if available
        if self.project_config:
            final_prompt = self._add_project_context(final_prompt, self.project_config)

        # Add task context if provided
        if context:
            final_prompt = self._add_task_context(final_prompt, context)

        # Update metadata
        metadata.update({
            "domain": domain,
            "complexity": complexity,
            "strategies": [s.value for s in strategies],
            "project": self.project_config.get("name", "unknown"),
            "final_token_estimate": len(final_prompt.split())
        })

        return final_prompt, metadata

    def _add_domain_augmentation(self, prompt: str, domain: str,
                                domain_context: Dict[str, Any]) -> str:
        """Add domain-specific augmentation"""
        domain_section = f"""
## Domain-Specific Requirements ({domain})

**Focus Areas:** {domain_context['focus']}

**Priorities:**
{chr(10).join(f'- {p}' for p in domain_context['priorities'])}

**Best Practices for This Domain:**
{chr(10).join(f'- {p}' for p in domain_context['best_practices'])}
"""
        return prompt + domain_section

    def _add_project_context(self, prompt: str, project_config: Dict[str, Any]) -> str:
        """Add project-specific context"""
        project_section = f"""
## Project Context

**Project:** {project_config.get('name', 'unknown')}
**Path:** {project_config.get('path', 'unknown')}
**Focus:** {project_config.get('focus', 'unknown')}
"""
        return prompt + project_section

    def _add_task_context(self, prompt: str, context: Dict[str, Any]) -> str:
        """Add task execution context"""
        context_section = "\n## Task Context\n\n"

        if context.get("previous_results"):
            context_section += "**Previous Results:**\n"
            for key, value in context["previous_results"].items():
                context_section += f"- {key}: {str(value)[:100]}...\n"

        if context.get("state"):
            context_section += "\n**Current State:**\n"
            for key, value in context["state"].items():
                context_section += f"- {key}: {value}\n"

        if context.get("blockers"):
            context_section += "\n**Known Blockers:**\n"
            for blocker in context["blockers"]:
                context_section += f"- {blocker}\n"

        return prompt + context_section

    def save_augmentation_report(self, output_path: Path,
                                original_task: str,
                                augmented_prompt: str,
                                metadata: Dict[str, Any]) -> None:
        """Save augmentation report for analysis"""
        report = {
            "timestamp": datetime.now().isoformat(),
            "original_task": original_task,
            "augmented_prompt": augmented_prompt,
            "metadata": metadata,
            "stats": {
                "original_length": len(original_task),
                "augmented_length": len(augmented_prompt),
                "augmentation_ratio": len(augmented_prompt) / max(len(original_task), 1)
            }
        }

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


# Example usage and best practices documentation
INTEGRATION_EXAMPLES = """
# Prompt Integration Examples

## Basic Usage

```python
from prompt_integration import PromptIntegrationEngine, TaskType

# Initialize engine
engine = PromptIntegrationEngine({
    "name": "luzia",
    "path": "/opt/server-agents/orchestrator",
    "focus": "orchestration and task dispatch"
})

# Augment a task
task = "Implement a caching layer for database queries"
augmented, metadata = engine.augment_for_task(
    task,
    task_type=TaskType.IMPLEMENTATION,
    domain="backend",
)

print(augmented)
print(f"Complexity: {metadata['complexity']}")
print(f"Strategies: {metadata['strategies']}")
```

## Domain-Specific Example

```python
# Security-focused analysis task
task = "Analyze the security implications of storing tokens in localStorage"
augmented, metadata = engine.augment_for_task(
    task,
    task_type=TaskType.ANALYSIS,
    domain="security",  # Triggers security-specific best practices
    complexity=4
)
```

## With Task Context

```python
# Continuation from previous work
context = {
    "previous_results": {
        "bottleneck_identified": "Database query N+1 problem",
        "estimated_improvement": "60% faster"
    },
    "state": {
        "implementation_status": "In progress",
        "current_focus": "Cache invalidation strategy"
    },
    "blockers": [
        "Need to decide on cache backend (Redis vs Memcached)"
    ]
}

augmented, metadata = engine.augment_for_task(
    "Continue implementation: select and integrate cache backend",
    task_type=TaskType.IMPLEMENTATION,
    context=context
)
```

## Auto-Complexity Detection

```python
# Complexity is auto-detected from task description
tasks = [
    "Fix typo in README",                    # -> Complexity 1
    "Add logging to error handler",          # -> Complexity 2
    "Implement distributed caching layer",   # -> Complexity 4
    "Refactor auth system for multi-tenancy" # -> Complexity 5
]

for task in tasks:
    augmented, metadata = engine.augment_for_task(
        task,
        task_type=TaskType.IMPLEMENTATION
    )
    print(f"Task: {task[:50]}...")
    print(f"Complexity: {metadata['complexity']}, Strategies: {len(metadata['strategies'])}")
```

## Key Features

1. **Chain-of-Thought**: Breaks complex tasks into reasoning steps
2. **Few-Shot Learning**: Provides examples of similar completed tasks
3. **Role-Based Prompting**: Sets appropriate expertise level for task type
4. **Context Hierarchies**: Prioritizes context by importance
5. **Complexity Adaptation**: Adjusts strategies based on task difficulty
6. **Domain Awareness**: Applies domain-specific best practices
7. **Self-Consistency**: For high-complexity tasks, can request multiple approaches

## Best Practices

- **Match Complexity to Strategies**: Higher complexity tasks benefit from more augmentation
- **Use Domain Context**: Domain-specific context dramatically improves quality
- **Provide Context**: Previous results and state help with continuations
- **Track Metadata**: Monitor augmentation ratios to prevent prompt bloat
- **Review Results**: Save augmentation reports to analyze and improve patterns
"""

__all__ = [
    'DomainSpecificAugmentor',
    'ComplexityAdaptivePrompting',
    'PromptIntegrationEngine',
]
