"""
Prompt templates for different generation tasks.

Each prompt is designed for specific extraction/generation tasks
in the cheat sheet pipeline.
"""

from dataclasses import dataclass
from typing import Any, Dict, List, Optional
from string import Template


@dataclass
class PromptTemplate:
    """Template for generating prompts."""
    template: str
    input_variables: List[str]
    description: str = ""
    
    def format(self, **kwargs) -> str:
        """Format the template with provided variables."""
        # Validate all required variables are provided
        missing = set(self.input_variables) - set(kwargs.keys())
        if missing:
            raise ValueError(f"Missing required variables: {missing}")
        
        return Template(self.template).safe_substitute(**kwargs)
    
    def __call__(self, **kwargs) -> str:
        """Allow calling template directly."""
        return self.format(**kwargs)


# =============================================================================
# Concept Extraction Prompts
# =============================================================================

CONCEPT_EXTRACTION_PROMPT = PromptTemplate(
    template="""You are an expert at extracting and summarizing key concepts from technical documents.

Given the following text from a book or paper, extract the main concepts with clear, concise definitions and intuitions.

TEXT:
$context

INSTRUCTIONS:
1. Identify the most important concepts (definitions, terms, ideas)
2. For each concept, provide:
   - Name: The concept name
   - Definition: A precise, formal definition
   - Intuition: A simple, intuitive explanation
   - Related concepts: Other concepts it connects to
3. Focus on concepts that are essential for understanding the material
4. Use LaTeX notation for any mathematical symbols: $symbol$

Format your response as a structured list:

CONCEPT: [Name]
DEFINITION: [Formal definition]
INTUITION: [Simple explanation]
RELATED: [Comma-separated related concepts]

---

Extract the key concepts:""",
    input_variables=["context"],
    description="Extract key concepts from technical text"
)


CONCEPT_SUMMARY_PROMPT = PromptTemplate(
    template="""Summarize the following concepts into a cohesive overview for a cheat sheet.

CONCEPTS:
$concepts

INSTRUCTIONS:
1. Group related concepts together
2. Create a logical flow from basic to advanced
3. Keep definitions precise but concise
4. Use bullet points for clarity
5. Include LaTeX notation for math: $symbol$
6. Maximum length: $max_words words

Write a well-organized concept summary:""",
    input_variables=["concepts", "max_words"],
    description="Summarize extracted concepts"
)


# =============================================================================
# Formula Extraction Prompts
# =============================================================================

FORMULA_EXTRACTION_PROMPT = PromptTemplate(
    template="""You are an expert at extracting and organizing mathematical formulas from technical documents.

Given the following text, extract all important mathematical formulas and equations.

TEXT:
$context

INSTRUCTIONS:
1. Identify all significant formulas and equations
2. For each formula, provide:
   - Name/Label: What the formula represents
   - Formula: The formula in LaTeX notation
   - Variables: Define each variable
   - Context: When/why this formula is used
   - Dependencies: Other formulas it relies on
3. Normalize all formulas to consistent LaTeX syntax
4. Group formulas by topic/purpose

Format your response:

FORMULA: [Name/Label]
LATEX: $$ [LaTeX formula] $$
VARIABLES:
- [variable]: [description]
CONTEXT: [Brief explanation of usage]
DEPENDS_ON: [Related formula names]

---

Extract the formulas:""",
    input_variables=["context"],
    description="Extract mathematical formulas"
)


FORMULA_ORGANIZATION_PROMPT = PromptTemplate(
    template="""Organize the following formulas into a coherent reference section for a cheat sheet.

FORMULAS:
$formulas

TOPIC: $topic

INSTRUCTIONS:
1. Group formulas by subtopic
2. Order from fundamental to derived
3. Show relationships between formulas
4. Keep LaTeX notation clean and consistent
5. Add brief explanatory notes where helpful
6. Maximum $max_formulas formulas

Create an organized formula reference:""",
    input_variables=["formulas", "topic", "max_formulas"],
    description="Organize formulas for cheat sheet"
)


# =============================================================================
# Model/Theory Summarization Prompts
# =============================================================================

MODEL_EXTRACTION_PROMPT = PromptTemplate(
    template="""You are an expert at summarizing machine learning and mathematical models.

Given the following text about a model or theoretical framework, create a comprehensive summary.

TEXT:
$context

MODEL/TOPIC: $model_name

INSTRUCTIONS:
1. Summarize the model's key components:
   - Assumptions: What does the model assume?
   - Objective: What is the model trying to achieve?
   - Formulation: Key equations and structure
   - Training/Optimization: How is it learned/solved?
   - Properties: Key theoretical properties
2. Use precise mathematical notation (LaTeX)
3. Be concise but complete
4. Highlight practical implications

Format your response:

## $model_name

### Assumptions
[List key assumptions]

### Objective
[Main goal and loss function/objective]

### Formulation
[Key equations with LaTeX]

### Training/Optimization
[How to train or solve]

### Key Properties
[Important theoretical results]

### Practical Notes
[Implementation tips, common pitfalls]

---

Summarize the model:""",
    input_variables=["context", "model_name"],
    description="Extract and summarize model descriptions"
)


# =============================================================================
# Algorithm Synthesis Prompts
# =============================================================================

ALGORITHM_EXTRACTION_PROMPT = PromptTemplate(
    template="""You are an expert at extracting and formalizing algorithms from technical documents.

Given the following text, extract and formalize the algorithm(s) described.

TEXT:
$context

INSTRUCTIONS:
1. Identify the main algorithm(s)
2. For each algorithm, provide:
   - Name: Algorithm name
   - Purpose: What problem it solves
   - Input: Required inputs with types
   - Output: What it produces
   - Pseudocode: Clear, numbered steps
   - Complexity: Time and space complexity
   - Key insight: The main idea that makes it work
3. Use clear pseudocode conventions
4. Note any important edge cases

Format your response:

ALGORITHM: [Name]
PURPOSE: [What it solves]
INPUT: [Input specification]
OUTPUT: [Output specification]

PSEUDOCODE:
```
1. [Step 1]
2. [Step 2]
   2.1 [Substep]
   ...
```

COMPLEXITY:
- Time: O(...)
- Space: O(...)

KEY INSIGHT: [Main algorithmic insight]
EDGE CASES: [Important edge cases to handle]

---

Extract the algorithm(s):""",
    input_variables=["context"],
    description="Extract algorithm specifications"
)


ALGORITHM_SYNTHESIS_PROMPT = PromptTemplate(
    template="""Create clean, well-documented pseudocode for the following algorithm description.

DESCRIPTION:
$description

ALGORITHM NAME: $algorithm_name

INSTRUCTIONS:
1. Write clear, numbered pseudocode
2. Use consistent indentation
3. Include comments for complex steps
4. Define any helper functions needed
5. Add complexity analysis
6. Include a brief example if helpful

Create the pseudocode:

```
Algorithm: $algorithm_name
Input: [specify]
Output: [specify]

[Pseudocode here]
```

Complexity Analysis:
- Time: 
- Space:

Example (optional):
[Brief example of execution]

---

Generate the pseudocode:""",
    input_variables=["description", "algorithm_name"],
    description="Synthesize pseudocode from description"
)


# =============================================================================
# General RAG Prompts
# =============================================================================

RAG_QA_PROMPT = PromptTemplate(
    template="""Answer the question based on the provided context. If the answer cannot be found in the context, say "I cannot find this information in the provided context."

CONTEXT:
$context

QUESTION: $question

INSTRUCTIONS:
1. Answer based only on the provided context
2. Be precise and concise
3. Use LaTeX for any math: $symbol$
4. Cite specific parts of the context if relevant
5. If uncertain, express the uncertainty

ANSWER:""",
    input_variables=["context", "question"],
    description="RAG question answering"
)


SUMMARIZATION_PROMPT = PromptTemplate(
    template="""Summarize the following text for inclusion in a technical cheat sheet.

TEXT:
$context

SECTION: $section_name
MAX LENGTH: $max_words words

INSTRUCTIONS:
1. Focus on the most important points
2. Use bullet points for clarity
3. Preserve technical accuracy
4. Include key formulas in LaTeX
5. Be concise but complete

SUMMARY:""",
    input_variables=["context", "section_name", "max_words"],
    description="Summarization for cheat sheet"
)


# =============================================================================
# Cheat Sheet Assembly Prompts
# =============================================================================

OVERVIEW_PROMPT = PromptTemplate(
    template="""Create a brief overview and roadmap for a cheat sheet about the following topic.

TOPIC: $topic
SUBTOPICS: $subtopics

INSTRUCTIONS:
1. Write a 2-3 sentence overview
2. Explain what the reader will learn
3. List the main sections with one-line descriptions
4. Keep it concise and engaging

OVERVIEW:""",
    input_variables=["topic", "subtopics"],
    description="Generate cheat sheet overview"
)


RECAP_PROMPT = PromptTemplate(
    template="""Create a "If you remember one page..." recap section for a cheat sheet.

MAIN CONTENT:
$content

INSTRUCTIONS:
1. Distill the absolute essentials
2. Include the 3-5 most important formulas
3. List the 5-7 key concepts
4. Add 2-3 practical tips
5. This should fit on one page
6. Use bullet points

RECAP:

## If You Remember One Page...

### Core Ideas
[Most important concepts]

### Essential Formulas
[Most important equations in LaTeX]

### Quick Tips
[Practical advice]

---

Generate the recap:""",
    input_variables=["content"],
    description="Generate one-page recap"
)


CONNECTIONS_PROMPT = PromptTemplate(
    template="""Identify connections between the main topic and related areas.

MAIN TOPIC: $topic
CONTENT SUMMARY: $summary
RELATED AREAS TO CONSIDER: $related_areas

INSTRUCTIONS:
1. Identify meaningful connections
2. Explain how concepts transfer
3. Note prerequisites and follow-up topics
4. Keep entries brief

FORMAT:
- [Related topic]: [Brief description of connection]

CONNECTIONS:""",
    input_variables=["topic", "summary", "related_areas"],
    description="Generate topic connections"
)


# =============================================================================
# Prompt Registry
# =============================================================================

PROMPT_REGISTRY = {
    # Concept extraction
    "concept_extraction": CONCEPT_EXTRACTION_PROMPT,
    "concept_summary": CONCEPT_SUMMARY_PROMPT,
    
    # Formula extraction
    "formula_extraction": FORMULA_EXTRACTION_PROMPT,
    "formula_organization": FORMULA_ORGANIZATION_PROMPT,
    
    # Model summarization
    "model_extraction": MODEL_EXTRACTION_PROMPT,
    
    # Algorithm synthesis
    "algorithm_extraction": ALGORITHM_EXTRACTION_PROMPT,
    "algorithm_synthesis": ALGORITHM_SYNTHESIS_PROMPT,
    
    # General RAG
    "rag_qa": RAG_QA_PROMPT,
    "summarization": SUMMARIZATION_PROMPT,
    
    # Cheat sheet assembly
    "overview": OVERVIEW_PROMPT,
    "recap": RECAP_PROMPT,
    "connections": CONNECTIONS_PROMPT,
}


def get_prompt(name: str) -> PromptTemplate:
    """Get a prompt template by name."""
    if name not in PROMPT_REGISTRY:
        raise ValueError(f"Unknown prompt: {name}. Available: {list(PROMPT_REGISTRY.keys())}")
    return PROMPT_REGISTRY[name]


def list_prompts() -> List[str]:
    """List available prompt names."""
    return list(PROMPT_REGISTRY.keys())
