"""
NCS-Specific Prompt Templates

Specialized prompts for Nuclear Criticality Safety domain,
designed for accurate, conservative, and well-cited responses.
"""

from sage.rag.prompts.base import (
    PromptTemplate,
    PromptVariable,
    SystemPrompt,
    register_system_prompt,
    register_template,
)

# =============================================================================
# System Prompts
# =============================================================================

NCS_SYSTEM_PROMPT = register_system_prompt(
    SystemPrompt(
        name="ncs_expert",
        description="Primary system prompt for NCS question answering",
        prompt="""You are SAGE-NCS, an expert assistant for Nuclear Criticality Safety (NCS).

ROLE:
You help nuclear engineers, criticality safety analysts, and researchers answer technical questions about criticality safety using authoritative sources.

CITATION RULES:
1. ALWAYS cite sources using [DocID, Section] format (e.g., [ANSI/ANS-8.1-2014, Table 1])
2. Every factual claim MUST have a citation from the provided context
3. When citing multiple sources, list them separately: [Source1], [Source2]
4. If information comes from a specific page, include it: [NUREG-1520, Section 3.4, p.42]

SAFETY RULES:
1. Be CONSERVATIVE - when uncertain, recommend more restrictive limits
2. NEVER fabricate references, data values, or technical parameters
3. If critical data is missing from context, explicitly state this
4. For calculations, recommend verification with SCALE/MCNP or other validated codes
5. When regulatory requirements conflict with best practices, cite both and recommend the more conservative approach

RESPONSE RULES:
1. If the context doesn't contain relevant information, say so explicitly
2. Distinguish between regulatory requirements (must) and guidance (should)
3. Include units for all numerical values
4. Note when standards have been superseded by newer versions
5. For safety limits, always specify the assumptions (e.g., water reflection, full density)

DOMAIN KNOWLEDGE:
- Familiar with NRC regulations (10 CFR), NUREGs, and regulatory guides
- Familiar with ANSI/ANS standards (8-series for criticality safety)
- Familiar with DOE orders and standards (DOE-STD-3007, etc.)
- Familiar with SCALE/MCNP codes and validation reports
- Familiar with ICSBEP benchmarks and criticality experiments""",
        max_input_tokens=100000,
        max_output_tokens=4096,
        recommended_temperature=0.3,
        require_citations=True,
        conservative_mode=True,
    )
)

NCS_CALCULATION_PROMPT = register_system_prompt(
    SystemPrompt(
        name="ncs_calculation",
        description="System prompt for calculation-related queries",
        prompt="""You are SAGE-NCS, assisting with Nuclear Criticality Safety calculations.

CALCULATION RULES:
1. Always show your work and assumptions
2. Use standard NCS notation and units
3. Cite the source of any formulas, correlations, or data values
4. Express uncertainties and safety margins clearly
5. Recommend validation with Monte Carlo codes for critical calculations

VALIDATION REQUIREMENTS:
1. Identify applicable benchmark experiments from ICSBEP if relevant
2. Note bias and bias uncertainty from validation studies
3. Recommend appropriate upper subcritical limit (USL) methodology
4. Cite applicable validation reports or technical basis documents

CONSERVATIVE APPROACH:
1. When parameters have ranges, use bounding values
2. Account for process variations and measurement uncertainties
3. Explicitly state the margin to criticality
4. Recommend additional contingencies for novel configurations

Remember: Never provide final safety limits without recommending independent verification.""",
        max_input_tokens=100000,
        max_output_tokens=8192,
        recommended_temperature=0.2,
        require_citations=True,
        conservative_mode=True,
    )
)

NCS_REGULATORY_PROMPT = register_system_prompt(
    SystemPrompt(
        name="ncs_regulatory",
        description="System prompt for regulatory interpretation queries",
        prompt="""You are SAGE-NCS, assisting with NCS regulatory interpretation.

REGULATORY INTERPRETATION RULES:
1. Cite specific regulatory text verbatim when answering
2. Distinguish between:
   - Regulations (legally binding, 10 CFR)
   - Regulatory guides (acceptable methods)
   - NUREGs (technical information)
   - Standards (industry consensus)
3. Note when interpretations may vary between licensees/facilities
4. Reference relevant precedents or enforcement actions when applicable

HIERARCHY OF REQUIREMENTS:
1. Federal regulations (10 CFR) take precedence
2. License conditions are facility-specific requirements
3. Regulatory guides describe acceptable implementation approaches
4. Standards endorsed by regulation have regulatory weight

When unsure about interpretation, recommend consultation with the facility's licensing basis or direct communication with NRC.""",
        max_input_tokens=100000,
        max_output_tokens=4096,
        recommended_temperature=0.2,
        require_citations=True,
        conservative_mode=True,
    )
)

# =============================================================================
# QA Templates
# =============================================================================

QA_WITH_CONTEXT = register_template(
    PromptTemplate(
        name="qa_with_context",
        category="ncs",
        description="Standard Q&A template with retrieved context",
        template="""CONTEXT FROM KNOWLEDGE BASE:
{{context}}

---

USER QUESTION:
{{question}}

---

Provide a comprehensive answer based on the context above. Follow all citation rules.""",
        variables=[
            PromptVariable(
                name="context",
                description="Retrieved chunks from the knowledge base with citations",
                required=True,
            ),
            PromptVariable(
                name="question",
                description="User's question about NCS",
                required=True,
            ),
        ],
    )
)

QA_WITH_CONTEXT_AND_HISTORY = register_template(
    PromptTemplate(
        name="qa_with_context_and_history",
        category="ncs",
        description="Q&A template with conversation history",
        template="""CONTEXT FROM KNOWLEDGE BASE:
{{context}}

---

CONVERSATION HISTORY:
{{history}}

---

USER QUESTION:
{{question}}

---

Provide a comprehensive answer based on the context and conversation history. Follow all citation rules.""",
        variables=[
            PromptVariable(
                name="context",
                description="Retrieved chunks from the knowledge base",
                required=True,
            ),
            PromptVariable(
                name="history",
                description="Previous conversation turns",
                required=True,
            ),
            PromptVariable(
                name="question",
                description="Current user question",
                required=True,
            ),
        ],
    )
)

QA_INSUFFICIENT_CONTEXT = register_template(
    PromptTemplate(
        name="qa_insufficient_context",
        category="ncs",
        description="Template when context is insufficient",
        template="""AVAILABLE CONTEXT (may be incomplete):
{{context}}

---

USER QUESTION:
{{question}}

---

The available context may not fully address this question. Provide what information you can from the context, clearly noting:
1. What IS addressed by the available sources
2. What is NOT addressed and would require additional references
3. Recommendations for authoritative sources to consult

Do not fabricate information not present in the context.""",
        variables=[
            PromptVariable(
                name="context",
                description="Retrieved chunks (may be incomplete)",
                required=True,
            ),
            PromptVariable(
                name="question",
                description="User's question",
                required=True,
            ),
        ],
    )
)

# =============================================================================
# Context Formatting Templates
# =============================================================================

CONTEXT_CHUNK_FORMAT = register_template(
    PromptTemplate(
        name="context_chunk_format",
        category="formatting",
        description="Format for individual context chunks",
        template="""[Source: {{doc_id}}{{section}}{{page}}]
{{content}}
---""",
        variables=[
            PromptVariable(
                name="doc_id",
                description="Document identifier",
                required=True,
            ),
            PromptVariable(
                name="section",
                description="Section reference",
                required=False,
                default="",
            ),
            PromptVariable(
                name="page",
                description="Page number",
                required=False,
                default="",
            ),
            PromptVariable(
                name="content",
                description="Chunk content",
                required=True,
            ),
        ],
    )
)

# =============================================================================
# Specialized Templates
# =============================================================================

LIMIT_COMPARISON = register_template(
    PromptTemplate(
        name="limit_comparison",
        category="ncs",
        description="Compare limits from multiple sources",
        template="""CONTEXT FROM MULTIPLE SOURCES:
{{context}}

---

REQUEST: Compare the {{limit_type}} limits for {{material}} from the available sources.

For each source, provide:
1. The limit value with units
2. Conditions/assumptions (reflection, moderation, form)
3. Regulatory status (requirement vs guidance)
4. Any notable differences between sources

Conclude with a recommendation for which limit to apply and why.""",
        variables=[
            PromptVariable("context", "Retrieved chunks about limits", required=True),
            PromptVariable("limit_type", "Type of limit (mass, concentration, etc.)", required=True),
            PromptVariable("material", "Fissile material", required=True),
        ],
    )
)

REGULATION_LOOKUP = register_template(
    PromptTemplate(
        name="regulation_lookup",
        category="ncs",
        description="Look up specific regulation requirements",
        template="""REGULATORY CONTEXT:
{{context}}

---

REGULATION QUERY:
{{question}}

---

Provide the relevant regulatory text and interpretation. Include:
1. Exact citation (CFR section, paragraph)
2. Key requirements or criteria
3. Applicable guidance documents
4. Common implementation approaches
5. Any relevant exemptions or exceptions""",
        variables=[
            PromptVariable("context", "Retrieved regulatory content", required=True),
            PromptVariable("question", "Regulatory question", required=True),
        ],
    )
)

CALCULATION_REVIEW = register_template(
    PromptTemplate(
        name="calculation_review",
        category="ncs",
        description="Review calculation approach",
        template="""TECHNICAL CONTEXT:
{{context}}

---

CALCULATION DESCRIPTION:
{{calculation}}

---

Review this calculation approach against the technical context. Address:
1. Are the methods consistent with accepted practice?
2. Are the assumptions appropriate and conservative?
3. Are validation references appropriate?
4. What additional considerations are needed?
5. Cite relevant standards or guidance for methodology.""",
        variables=[
            PromptVariable("context", "Technical references", required=True),
            PromptVariable("calculation", "Calculation approach to review", required=True),
        ],
    )
)

DOCUMENT_SUMMARY = register_template(
    PromptTemplate(
        name="document_summary",
        category="ncs",
        description="Summarize key points from a document",
        template="""DOCUMENT CONTENT:
{{context}}

---

Provide a technical summary of this document for NCS purposes. Include:
1. Document scope and applicability
2. Key technical requirements or limits
3. Important assumptions or conditions
4. Relationship to other standards/regulations
5. Implementation considerations

Cite specific sections for each point.""",
        variables=[
            PromptVariable("context", "Document content chunks", required=True),
        ],
    )
)

# =============================================================================
# Citation Format Templates
# =============================================================================

CITATION_FORMAT_INLINE = register_template(
    PromptTemplate(
        name="citation_format_inline",
        category="formatting",
        description="Format for inline citations",
        template="[{{doc_id}}{{section}}]",
        variables=[
            PromptVariable("doc_id", "Document ID", required=True),
            PromptVariable("section", "Section reference", required=False, default=""),
        ],
    )
)

CITATION_FORMAT_FULL = register_template(
    PromptTemplate(
        name="citation_format_full",
        category="formatting",
        description="Format for full citations",
        template="{{doc_id}}, {{title}}{{section}}{{page}}",
        variables=[
            PromptVariable("doc_id", "Document ID", required=True),
            PromptVariable("title", "Document title", required=True),
            PromptVariable("section", "Section reference", required=False, default=""),
            PromptVariable("page", "Page number", required=False, default=""),
        ],
    )
)


# =============================================================================
# Helper Functions
# =============================================================================

def get_ncs_system_prompt(task_type: str = "general") -> SystemPrompt:
    """
    Get the appropriate system prompt for an NCS task.

    Args:
        task_type: Type of task ("general", "calculation", "regulatory")

    Returns:
        Appropriate SystemPrompt
    """
    prompts = {
        "general": NCS_SYSTEM_PROMPT,
        "calculation": NCS_CALCULATION_PROMPT,
        "regulatory": NCS_REGULATORY_PROMPT,
    }
    return prompts.get(task_type, NCS_SYSTEM_PROMPT)


def get_qa_template(has_history: bool = False, low_confidence: bool = False) -> PromptTemplate:
    """
    Get the appropriate QA template.

    Args:
        has_history: Whether conversation history is available
        low_confidence: Whether retrieval confidence is low

    Returns:
        Appropriate PromptTemplate
    """
    if low_confidence:
        return QA_INSUFFICIENT_CONTEXT
    if has_history:
        return QA_WITH_CONTEXT_AND_HISTORY
    return QA_WITH_CONTEXT
