# createsonline/agents/prompts.py
"""
Prompt engineering utilities for CREATESONLINE agents.

Provides PromptTemplate, FewShotPrompt, and ChainOfThought builders.
"""

import re
import json
from typing import Any, Dict, List, Optional, Tuple


class PromptTemplate:
    """
    A template for constructing prompts with variable substitution.
    
    Usage::
    
        template = PromptTemplate(
            "You are a {role}. The user asks: {question}\\n"
            "Respond in {style} style."
        )
        prompt = template.render(role="teacher", question="What is AI?", style="simple")
    """
    
    def __init__(self, template: str, defaults: Dict[str, str] = None,
                 validators: Dict[str, callable] = None):
        self.template = template
        self.defaults = defaults or {}
        self.validators = validators or {}
        self._variables = self._extract_variables()
    
    def _extract_variables(self) -> List[str]:
        """Extract variable names from the template."""
        return list(set(re.findall(r'\{(\w+)\}', self.template)))
    
    @property
    def variables(self) -> List[str]:
        return list(self._variables)
    
    def render(self, **kwargs) -> str:
        """Render the template with provided values."""
        # Apply defaults for missing vars
        values = {**self.defaults, **kwargs}
        
        # Validate
        for var, validator in self.validators.items():
            if var in values:
                if not validator(values[var]):
                    raise ValueError(f"Validation failed for variable: {var}")
        
        # Check all vars are provided
        missing = [v for v in self._variables if v not in values]
        if missing:
            raise ValueError(f"Missing template variables: {missing}")
        
        return self.template.format(**values)
    
    def partial(self, **kwargs) -> 'PromptTemplate':
        """Return a new template with some variables pre-filled."""
        new_defaults = {**self.defaults, **kwargs}
        return PromptTemplate(self.template, new_defaults, self.validators)
    
    def __add__(self, other: 'PromptTemplate') -> 'PromptTemplate':
        """Concatenate two templates."""
        if isinstance(other, str):
            return PromptTemplate(self.template + other, self.defaults)
        return PromptTemplate(
            self.template + other.template,
            {**self.defaults, **other.defaults}
        )
    
    def __repr__(self):
        preview = self.template[:60] + '...' if len(self.template) > 60 else self.template
        return f"PromptTemplate({preview!r})"


class FewShotPrompt:
    """
    Few-shot prompt builder — constructs prompts with examples.
    
    Usage::
    
        prompt = FewShotPrompt(
            instruction="Classify the sentiment of the text.",
            examples=[
                {"input": "I love this!", "output": "positive"},
                {"input": "This is terrible", "output": "negative"},
            ],
            input_key="input",
            output_key="output",
        )
        text = prompt.render(query="The movie was okay")
    """
    
    def __init__(self, instruction: str, examples: List[Dict[str, str]] = None,
                 input_key: str = 'input', output_key: str = 'output',
                 example_separator: str = '\n\n', prefix: str = '',
                 suffix: str = '', input_label: str = 'Input',
                 output_label: str = 'Output'):
        self.instruction = instruction
        self.examples = examples or []
        self.input_key = input_key
        self.output_key = output_key
        self.example_separator = example_separator
        self.prefix = prefix
        self.suffix = suffix
        self.input_label = input_label
        self.output_label = output_label
    
    def add_example(self, **kwargs):
        """Add an example."""
        self.examples.append(kwargs)
    
    def format_example(self, example: dict) -> str:
        """Format a single example."""
        inp = example.get(self.input_key, '')
        out = example.get(self.output_key, '')
        return f"{self.input_label}: {inp}\n{self.output_label}: {out}"
    
    def render(self, query: str, max_examples: int = None) -> str:
        """Render the few-shot prompt with the query."""
        parts = []
        
        if self.prefix:
            parts.append(self.prefix)
        
        parts.append(self.instruction)
        
        examples = self.examples
        if max_examples is not None:
            examples = examples[:max_examples]
        
        if examples:
            parts.append("")  # blank line
            parts.append("Examples:")
            for ex in examples:
                parts.append(self.format_example(ex))
        
        parts.append("")
        parts.append(f"{self.input_label}: {query}")
        parts.append(f"{self.output_label}:")
        
        if self.suffix:
            parts.append(self.suffix)
        
        return self.example_separator.join(parts) if self.example_separator != '\n\n' \
            else '\n'.join(parts)
    
    def select_examples(self, query: str, top_k: int = 3) -> List[dict]:
        """
        Select the most relevant examples for a query using keyword overlap.
        For better selection, use VectorMemory to store examples.
        """
        query_words = set(query.lower().split())
        scored = []
        for ex in self.examples:
            inp = ex.get(self.input_key, '').lower()
            overlap = len(query_words & set(inp.split()))
            scored.append((overlap, ex))
        scored.sort(key=lambda x: x[0], reverse=True)
        return [ex for _, ex in scored[:top_k]]
    
    def render_dynamic(self, query: str, top_k: int = 3) -> str:
        """Render with dynamically selected examples."""
        selected = self.select_examples(query, top_k)
        old_examples = self.examples
        self.examples = selected
        result = self.render(query)
        self.examples = old_examples
        return result
    
    def __repr__(self):
        return f"FewShotPrompt({len(self.examples)} examples)"


class ChainOfThought:
    """
    Chain-of-thought prompt builder for step-by-step reasoning.
    
    Usage::
    
        cot = ChainOfThought(
            instruction="Solve the problem step by step.",
            steps=[
                "Identify the key information",
                "Set up the equation",
                "Solve the equation",
                "Verify the answer",
            ],
        )
        text = cot.render(query="If John has 5 apples and gives 2 away, how many remain?")
    """
    
    def __init__(self, instruction: str = '', steps: List[str] = None,
                 step_prefix: str = 'Step {n}: ', reasoning_prompt: str = None,
                 answer_prefix: str = 'Therefore, the answer is: '):
        self.instruction = instruction or "Think step by step to solve this problem."
        self.steps = steps or []
        self.step_prefix = step_prefix
        self.reasoning_prompt = reasoning_prompt or "Let's think step by step."
        self.answer_prefix = answer_prefix
    
    def render(self, query: str, include_steps: bool = True) -> str:
        """Render the chain-of-thought prompt."""
        parts = [self.instruction, '', f"Problem: {query}", '']
        
        if include_steps and self.steps:
            parts.append("Follow these steps:")
            for i, step in enumerate(self.steps, 1):
                parts.append(f"{self.step_prefix.format(n=i)}{step}")
            parts.append('')
        
        parts.append(self.reasoning_prompt)
        parts.append('')
        
        return '\n'.join(parts)
    
    def render_with_scratchpad(self, query: str) -> str:
        """Render with a scratchpad format for intermediate reasoning."""
        parts = [
            self.instruction,
            '',
            f"Question: {query}",
            '',
            "Scratchpad (show your work):",
            "---",
        ]
        
        if self.steps:
            for i, step in enumerate(self.steps, 1):
                parts.append(f"{self.step_prefix.format(n=i)}{step}")
                parts.append(f"  [Your reasoning here]")
                parts.append('')
        
        parts.append("---")
        parts.append(self.answer_prefix)
        
        return '\n'.join(parts)
    
    def __repr__(self):
        return f"ChainOfThought({len(self.steps)} steps)"


class ReActPrompt:
    """
    ReAct (Reasoning + Acting) prompt format for tool-using agents.
    
    Generates prompts in the format:
        Thought: I need to search for information about X
        Action: search_web
        Action Input: {"query": "X"}
        Observation: [result from tool]
        Thought: Now I know the answer
        Final Answer: ...
    """
    
    def __init__(self, system: str = '', tools_description: str = '',
                 max_iterations: int = 10):
        self.system = system or "You are a helpful assistant that reasons step by step."
        self.tools_description = tools_description
        self.max_iterations = max_iterations
    
    def build_system_prompt(self, tools_desc: str = None) -> str:
        """Build the complete system prompt with tool descriptions."""
        desc = tools_desc or self.tools_description
        return f"""{self.system}

You have access to the following tools:
{desc}

Use the following format:

Question: the input question you must answer
Thought: reason about what to do next
Action: the tool to use (tool name only)
Action Input: the input for the tool (as JSON)
Observation: the result of the action
... (repeat Thought/Action/Action Input/Observation as needed)
Thought: I now know the final answer
Final Answer: the final answer to the question

Begin!
"""
    
    def format_iteration(self, thought: str, action: str = None,
                         action_input: dict = None, observation: str = None,
                         final_answer: str = None) -> str:
        """Format a single ReAct iteration."""
        parts = [f"Thought: {thought}"]
        
        if final_answer is not None:
            parts.append(f"Final Answer: {final_answer}")
        elif action:
            parts.append(f"Action: {action}")
            if action_input is not None:
                parts.append(f"Action Input: {json.dumps(action_input)}")
            if observation is not None:
                parts.append(f"Observation: {observation}")
        
        return '\n'.join(parts)
    
    @staticmethod
    def parse_response(text: str) -> dict:
        """
        Parse a ReAct-format response from the LLM.
        
        Returns dict with keys: thought, action, action_input, final_answer.
        """
        result = {
            'thought': None,
            'action': None,
            'action_input': None,
            'final_answer': None,
        }
        
        lines = text.strip().split('\n')
        for line in lines:
            line = line.strip()
            if line.startswith('Thought:'):
                result['thought'] = line[len('Thought:'):].strip()
            elif line.startswith('Action:'):
                result['action'] = line[len('Action:'):].strip()
            elif line.startswith('Action Input:'):
                raw = line[len('Action Input:'):].strip()
                try:
                    result['action_input'] = json.loads(raw)
                except json.JSONDecodeError:
                    result['action_input'] = {'input': raw}
            elif line.startswith('Observation:'):
                pass  # Observations come from tools, not the LLM
            elif line.startswith('Final Answer:'):
                result['final_answer'] = line[len('Final Answer:'):].strip()
        
        return result
    
    def __repr__(self):
        return f"ReActPrompt(max_iter={self.max_iterations})"
