"""
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="""Extract the key concepts from this technical text. Be precise and structured.

TEXT:
$context

OUTPUT FORMAT (follow exactly):
CONCEPT: [concept name]
DEFINITION: [one clear sentence definition]
INTUITION: [simple explanation in plain words]
RELATED: [comma-separated related terms]
---

Rules:
- Extract 3-5 most important concepts only
- Keep definitions under 30 words
- Use $...$ for math symbols
- One concept per block, separated by ---

Extract concepts now:""",
    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="""Extract mathematical formulas from this text. Output in structured format.

TEXT:
$context

OUTPUT FORMAT (follow exactly):
FORMULA: [name]
LATEX: [formula using LaTeX, e.g., $E = mc^2$ or $$F = ma$$]
VARIABLES: [var1: meaning, var2: meaning]
CONTEXT: [when to use this formula]
---

Rules:
- Extract 3-5 most important formulas only
- Use proper LaTeX syntax: $ for inline, $$ for display
- Keep variable descriptions brief
- One formula per block, separated by ---

Extract formulas now:""",
    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="""Summarize the model or theory described in this text.

TEXT:
$context

MODEL NAME: $model_name

OUTPUT FORMAT (follow exactly):
NAME: [model name]
DESCRIPTION: [2-3 sentence summary]
ASSUMPTIONS: [bullet list of key assumptions]
FORMULATION: [main equation in LaTeX]
PROPERTIES: [key theoretical properties]
---

Rules:
- Be concise, max 100 words total
- Use proper LaTeX for equations
- List 2-4 assumptions maximum
- Focus on practical understanding

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


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

ALGORITHM_EXTRACTION_PROMPT = PromptTemplate(
    template="""Extract the algorithm from this text and write pseudocode.

TEXT:
$context

OUTPUT FORMAT (follow exactly):
ALGORITHM: [name]
PURPOSE: [what problem it solves, one sentence]
INPUT: [input parameters]
OUTPUT: [what it returns]
PSEUDOCODE:
1. [step 1]
2. [step 2]
3. [step 3]
COMPLEXITY: Time O(...), Space O(...)
---

Rules:
- Maximum 10 steps in pseudocode
- Use simple, clear language
- No code syntax, just pseudocode
- One algorithm per block

Extract algorithm now:""",
    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="""Write a brief overview for a cheat sheet on this topic.

TOPIC: $topic
SUBTOPICS: $subtopics

OUTPUT FORMAT:
Write exactly 3 sentences:
1. What this topic is about
2. Why it matters
3. What the reader will learn

Keep it under 50 words total. No bullet points.

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


RECAP_PROMPT = PromptTemplate(
    template="""Create a quick reference summary from this content.

CONTENT:
$content

OUTPUT FORMAT:
## Quick Reference

### Key Concepts (5 max)
- [concept 1]: [brief definition]
- [concept 2]: [brief definition]

### Essential Formulas (3 max)
- [name]: $formula$

### Tips
- [practical tip 1]
- [practical tip 2]

Rules:
- Maximum 100 words total
- Use LaTeX for formulas
- Be extremely concise

Write recap now:""",
    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())
