prompeta 0.1.0

//! # prompeta 🎺
//!
//! Declarative prompt engineering with semantic roles and pluggable rendering strategies.
//!
//! Instead of building prompts as strings, you declare content with semantic roles
//! (instruction, context, focus, constraint, example, schema) and a [`PromptStrategy`]
//! renders the final text using best-practice prompt engineering.
//!
//! ## Quick start
//!
//! ```rust
//! use prompeta::{Prompt, ContentRole};
//!
//! let prompt = Prompt::new()
//!     .instruction("You are a helpful assistant.")
//!     .context("User profile", "Name: Alice, Role: Engineer")
//!     .focus("Current task", "Debug the authentication module")
//!     .constraint("Do not modify production database")
//!     .schema(serde_json::json!({
//!         "diagnosis": "string",
//!         "fix": "string",
//!         "confidence": "float 0-1"
//!     }));
//!
//! // Renders using the default ClaudeStrategy
//! let text = prompt.to_string();
//! ```
//!
//! ## Key concepts
//!
//! - **Content roles** — Each block of content has a semantic role that tells the
//!   renderer how to format and order it.
//! - **Pluggable strategies** — The [`PromptStrategy`] trait lets you swap rendering
//!   logic per model family. A [`ClaudeStrategy`] is provided by default.
//! - **Budget trimming** — [`Prompt::trim_to_budget`] removes lowest-priority blocks
//!   until the output fits within a character budget.
//! - **Tag-based manipulation** — Block-level tags let you surgically remove or replace
//!   content without position-dependent code.

use std::fmt;

use serde::{Deserialize, Serialize};

// ---------------------------------------------------------------------------
// ContentRole
// ---------------------------------------------------------------------------

/// Semantic role of a content block within a prompt.
///
/// Strategies use roles to decide ordering, formatting, and trimming behavior.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum ContentRole {
    /// System-level directives the model must follow.
    Instruction,
    /// Background information the model should be aware of.
    Context,
    /// High-attention content the model should focus on.
    Focus,
    /// Negative instructions — things the model must NOT do.
    Constraint,
    /// Demonstrations of desired input/output behavior.
    Example,
    /// Expected output structure (JSON schema, format spec).
    Schema,
}

impl ContentRole {
    /// Default priority for this role (higher = harder to trim).
    fn default_priority(self) -> u8 {
        match self {
            Self::Instruction => 255,
            Self::Schema => 240,
            Self::Focus => 192,
            Self::Constraint => 200,
            Self::Example => 100,
            Self::Context => 128,
        }
    }

    /// Ordering rank used by [`ClaudeStrategy`]. Lower = rendered earlier.
    fn claude_order(self) -> u8 {
        match self {
            Self::Instruction => 0,
            Self::Focus => 1,
            Self::Context => 2,
            Self::Example => 3,
            Self::Constraint => 4,
            Self::Schema => 5,
        }
    }
}

// ---------------------------------------------------------------------------
// Block
// ---------------------------------------------------------------------------

/// A content block with semantic metadata.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Block {
    /// Semantic role of this block.
    pub role: ContentRole,
    /// The content text.
    pub content: String,
    /// Human-readable label (rendered as a heading by most strategies).
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub label: Option<String>,
    /// Programmatic tag for removal/replacement.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub tag: Option<String>,
    /// Trim priority: 0 = trim first, 255 = never trim.
    #[serde(default = "Block::default_priority_value")]
    pub priority: u8,
}

impl Block {
    fn default_priority_value() -> u8 {
        128
    }
}

// ---------------------------------------------------------------------------
// Prompt
// ---------------------------------------------------------------------------

/// A declarative prompt: a collection of semantically-typed blocks.
///
/// Build prompts by declaring content with roles. The rendering strategy
/// (default: [`ClaudeStrategy`]) decides the final text layout.
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct Prompt {
    /// Ordered content blocks.
    pub blocks: Vec<Block>,
}

impl Prompt {
    /// Create an empty prompt.
    pub fn new() -> Self {
        Self::default()
    }

    // --- Role-specific builders ---

    /// Add an instruction block (highest priority, rendered first).
    pub fn instruction(self, content: impl Into<String>) -> Self {
        self.push_block(ContentRole::Instruction, content, None)
    }

    /// Add a context block with a label.
    pub fn context(self, label: impl Into<String>, content: impl Into<String>) -> Self {
        self.push_block(ContentRole::Context, content, Some(label.into()))
    }

    /// Add a focus block (high attention, rendered early).
    pub fn focus(self, label: impl Into<String>, content: impl Into<String>) -> Self {
        self.push_block(ContentRole::Focus, content, Some(label.into()))
    }

    /// Add a constraint (negative instruction).
    pub fn constraint(self, content: impl Into<String>) -> Self {
        self.push_block(ContentRole::Constraint, content, None)
    }

    /// Add an example with a label.
    pub fn example(self, label: impl Into<String>, content: impl Into<String>) -> Self {
        self.push_block(ContentRole::Example, content, Some(label.into()))
    }

    /// Add a JSON schema block from a `serde_json::Value`.
    ///
    /// The strategy will prepend "Return ONLY valid JSON" (or equivalent).
    pub fn schema(self, schema: serde_json::Value) -> Self {
        let json = serde_json::to_string_pretty(&schema).unwrap_or_default();
        self.push_block(ContentRole::Schema, json, None)
    }

    /// Add a JSON schema block with a custom note/instruction.
    ///
    /// The note replaces the default "Return ONLY valid JSON" header.
    pub fn schema_with_note(self, schema: serde_json::Value, note: impl Into<String>) -> Self {
        let json = serde_json::to_string_pretty(&schema).unwrap_or_default();
        // Store the note in the label field; renderer uses it as the header.
        let mut this = self.push_block(ContentRole::Schema, json, None);
        if let Some(last) = this.blocks.last_mut() {
            last.label = Some(note.into());
        }
        this
    }

    // --- Generic builder ---

    /// Add a block with an explicit role.
    pub fn block(self, role: ContentRole, content: impl Into<String>) -> Self {
        self.push_block(role, content, None)
    }

    /// Add a labeled block with an explicit role.
    pub fn labeled_block(
        self,
        role: ContentRole,
        label: impl Into<String>,
        content: impl Into<String>,
    ) -> Self {
        self.push_block(role, content, Some(label.into()))
    }

    // --- Metadata on last block ---

    /// Set the tag on the last block (for removal/replacement).
    pub fn tag(mut self, tag: impl Into<String>) -> Self {
        if let Some(last) = self.blocks.last_mut() {
            last.tag = Some(tag.into());
        }
        self
    }

    /// Set the priority on the last block.
    pub fn priority(mut self, priority: u8) -> Self {
        if let Some(last) = self.blocks.last_mut() {
            last.priority = priority;
        }
        self
    }

    // --- Manipulation ---

    /// Remove all blocks with the given tag.
    pub fn remove(mut self, tag: &str) -> Self {
        self.blocks.retain(|b| b.tag.as_deref() != Some(tag));
        self
    }

    /// Append all blocks from another prompt.
    pub fn merge(mut self, other: Prompt) -> Self {
        self.blocks.extend(other.blocks);
        self
    }

    // --- Budget ---

    /// Remove lowest-priority blocks until the rendered output fits within `max_chars`.
    ///
    /// Ties are broken by position (later blocks trimmed first).
    /// Instructions and Schema blocks are never trimmed regardless of priority.
    pub fn trim_to_budget(&mut self, max_chars: usize) {
        while self.estimated_len() > max_chars {
            // Find the trimmable block with the lowest priority (last among ties).
            let candidate = self
                .blocks
                .iter()
                .enumerate()
                .filter(|(_, b)| {
                    !b.content.is_empty()
                        && b.role != ContentRole::Instruction
                        && b.role != ContentRole::Schema
                })
                .min_by(|(i_a, a), (i_b, b)| a.priority.cmp(&b.priority).then(i_b.cmp(i_a)))
                .map(|(i, _)| i);

            match candidate {
                Some(idx) => {
                    self.blocks.remove(idx);
                }
                None => break, // nothing left to trim
            }
        }
    }

    /// Estimate the rendered character length without allocating the full string.
    pub fn estimated_len(&self) -> usize {
        let mut len = 0usize;
        for block in &self.blocks {
            if block.content.is_empty() {
                continue;
            }
            len += block.content.len();
            if let Some(ref label) = block.label {
                len += label.len() + 6; // "## " + label + "\n\n"
            }
            len += 2; // block separator "\n\n"
        }
        len
    }

    // --- Rendering ---

    /// Render using a specific strategy.
    pub fn render(&self, strategy: &dyn PromptStrategy) -> String {
        strategy.render(self)
    }

    /// Render using the default [`ClaudeStrategy`].
    pub fn render_default(&self) -> String {
        CLAUDE_STRATEGY.render(self)
    }

    // --- Internal ---

    fn push_block(
        mut self,
        role: ContentRole,
        content: impl Into<String>,
        label: Option<String>,
    ) -> Self {
        let content = content.into();
        if content.trim().is_empty() {
            return self; // skip empty blocks
        }
        let priority = role.default_priority();
        self.blocks.push(Block {
            role,
            content,
            label,
            tag: None,
            priority,
        });
        self
    }
}

impl fmt::Display for Prompt {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{}", self.render_default())
    }
}

// ---------------------------------------------------------------------------
// PromptStrategy
// ---------------------------------------------------------------------------

/// Pluggable rendering algorithm for converting a [`Prompt`] to text.
///
/// Implement this trait to provide custom prompt rendering logic — different
/// models may benefit from different ordering, formatting, or attention patterns.
pub trait PromptStrategy: Send + Sync {
    /// Render the prompt to a string.
    fn render(&self, prompt: &Prompt) -> String;
}

// ---------------------------------------------------------------------------
// ClaudeStrategy
// ---------------------------------------------------------------------------

/// Default strategy optimized for Claude models.
///
/// Applies research-backed practices:
///
/// 1. **Role-aware ordering** (primacy-recency):
///    Instructions → Focus → Context → Examples → Constraints → Schema.
///    Critical directives at the start; output format at the end (closest to generation).
///
/// 2. **Formatting per role**:
///    - `Instruction` → plain text, no heading (direct authoritative voice)
///    - `Context` → `## {label}\n\n{content}`
///    - `Focus` → `## {label}\n\n{content}`
///    - `Constraint` → collected into a bullet list
///    - `Example` → `### Example: {label}\n\n{content}`
///    - `Schema` → "Return ONLY valid JSON (no markdown fences):\n{json}" at end
///
/// 3. **Empty blocks** are silently skipped.
pub struct ClaudeStrategy;

static CLAUDE_STRATEGY: ClaudeStrategy = ClaudeStrategy;

impl PromptStrategy for ClaudeStrategy {
    fn render(&self, prompt: &Prompt) -> String {
        if prompt.blocks.is_empty() {
            return String::new();
        }

        // Sort blocks by role order (stable sort preserves insertion order within role).
        let mut sorted: Vec<&Block> = prompt
            .blocks
            .iter()
            .filter(|b| !b.content.trim().is_empty())
            .collect();
        sorted.sort_by_key(|b| b.role.claude_order());

        let mut parts: Vec<String> = Vec::new();
        let mut constraints: Vec<&str> = Vec::new();

        for block in &sorted {
            match block.role {
                ContentRole::Instruction => {
                    parts.push(block.content.clone());
                }
                ContentRole::Context => {
                    if let Some(ref label) = block.label {
                        parts.push(format!("## {}\n\n{}", label, block.content));
                    } else {
                        parts.push(block.content.clone());
                    }
                }
                ContentRole::Focus => {
                    if let Some(ref label) = block.label {
                        parts.push(format!("## {}\n\n{}", label, block.content));
                    } else {
                        parts.push(block.content.clone());
                    }
                }
                ContentRole::Example => {
                    if let Some(ref label) = block.label {
                        parts.push(format!("### Example: {}\n\n{}", label, block.content));
                    } else {
                        parts.push(format!("### Example\n\n{}", block.content));
                    }
                }
                ContentRole::Constraint => {
                    // Collect constraints; render as a single block later.
                    constraints.push(&block.content);
                }
                ContentRole::Schema => {
                    // Schema blocks are collected separately and rendered last.
                    // (handled after the loop)
                }
            }
        }

        // Render collected constraints as a bullet list.
        if !constraints.is_empty() {
            let list = constraints
                .iter()
                .map(|c| format!("- {}", c))
                .collect::<Vec<_>>()
                .join("\n");
            parts.push(format!("Rules:\n{}", list));
        }

        // Render schema blocks last.
        for block in &sorted {
            if block.role == ContentRole::Schema {
                let header = block
                    .label
                    .as_deref()
                    .unwrap_or("Return ONLY valid JSON (no markdown fences):");
                parts.push(format!("{}\n{}", header, block.content));
            }
        }

        parts.join("\n\n")
    }
}

// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn empty_prompt_renders_empty() {
        assert_eq!(Prompt::new().to_string(), "");
    }

    #[test]
    fn instruction_only() {
        let p = Prompt::new().instruction("You are a helpful assistant.");
        assert_eq!(p.to_string(), "You are a helpful assistant.");
    }

    #[test]
    fn ordering_instructions_first_schema_last() {
        let p = Prompt::new()
            .schema(serde_json::json!({"answer": "string"}))
            .context("Data", "Some data here")
            .instruction("Be concise.");

        let rendered = p.to_string();
        let inst_pos = rendered.find("Be concise.").unwrap();
        let ctx_pos = rendered.find("Some data here").unwrap();
        let schema_pos = rendered.find("\"answer\"").unwrap();

        assert!(inst_pos < ctx_pos, "instruction should come before context");
        assert!(ctx_pos < schema_pos, "context should come before schema");
    }

    #[test]
    fn constraints_collected_into_rules_block() {
        let p = Prompt::new()
            .constraint("Do not hallucinate")
            .constraint("Be concise");

        let rendered = p.to_string();
        assert!(rendered.contains("Rules:"));
        assert!(rendered.contains("- Do not hallucinate"));
        assert!(rendered.contains("- Be concise"));
    }

    #[test]
    fn context_with_label() {
        let p = Prompt::new().context("User Profile", "Name: Alice");
        assert!(p.to_string().contains("## User Profile\n\nName: Alice"));
    }

    #[test]
    fn focus_with_label() {
        let p = Prompt::new().focus("Current Task", "Debug the auth module");
        assert!(p
            .to_string()
            .contains("## Current Task\n\nDebug the auth module"));
    }

    #[test]
    fn example_with_label() {
        let p = Prompt::new().example("Good response", "Here is how...");
        assert!(p
            .to_string()
            .contains("### Example: Good response\n\nHere is how..."));
    }

    #[test]
    fn schema_renders_json() {
        let p = Prompt::new().schema(serde_json::json!({"score": "float"}));
        let rendered = p.to_string();
        assert!(rendered.contains("Return ONLY valid JSON"));
        assert!(rendered.contains("\"score\": \"float\""));
    }

    #[test]
    fn schema_with_note() {
        let p = Prompt::new().schema_with_note(serde_json::json!({"x": 1}), "Output format:");
        let rendered = p.to_string();
        assert!(rendered.starts_with("Output format:"));
        assert!(rendered.contains("\"x\": 1"));
    }

    #[test]
    fn empty_content_skipped() {
        let p = Prompt::new()
            .instruction("Hello")
            .context("Empty", "")
            .context("Also empty", "   ");
        assert_eq!(p.to_string(), "Hello");
    }

    #[test]
    fn tag_and_remove() {
        let p = Prompt::new()
            .context("Keep", "important")
            .tag("keep")
            .context("Drop", "unimportant")
            .tag("drop");
        let p = p.remove("drop");
        let rendered = p.to_string();
        assert!(rendered.contains("important"));
        assert!(!rendered.contains("unimportant"));
    }

    #[test]
    fn priority_override() {
        let p = Prompt::new()
            .context("Low", "low priority data")
            .priority(10)
            .context("High", "high priority data")
            .priority(250);

        // Trim until only one context fits
        let mut p = p;
        // Force a small budget
        p.trim_to_budget(40);
        let rendered = p.to_string();
        // Low priority should be trimmed first
        assert!(!rendered.contains("low priority data") || rendered.contains("high priority data"));
    }

    #[test]
    fn trim_to_budget_never_trims_instructions() {
        let mut p = Prompt::new()
            .instruction("Critical instruction")
            .context("Verbose context", &"x".repeat(1000));
        p.trim_to_budget(100);
        let rendered = p.to_string();
        assert!(rendered.contains("Critical instruction"));
        assert!(!rendered.contains(&"x".repeat(1000)));
    }

    #[test]
    fn trim_to_budget_never_trims_schema() {
        let mut p = Prompt::new()
            .schema(serde_json::json!({"required": "field"}))
            .context("Verbose", &"y".repeat(1000));
        p.trim_to_budget(100);
        let rendered = p.to_string();
        assert!(rendered.contains("required"));
        assert!(!rendered.contains(&"y".repeat(1000)));
    }

    #[test]
    fn merge_combines_blocks() {
        let a = Prompt::new().instruction("A");
        let b = Prompt::new().instruction("B");
        let merged = a.merge(b);
        assert_eq!(merged.blocks.len(), 2);
    }

    #[test]
    fn serde_roundtrip() {
        let p = Prompt::new()
            .instruction("Test")
            .context("Data", "value")
            .constraint("No bad things");

        let json = serde_json::to_string(&p).unwrap();
        let p2: Prompt = serde_json::from_str(&json).unwrap();

        assert_eq!(p.blocks.len(), p2.blocks.len());
        assert_eq!(p.to_string(), p2.to_string());
    }

    #[test]
    fn plain_string_fallback() {
        // When deserializing a plain string, it should fail gracefully
        let result = serde_json::from_str::<Prompt>("\"just a string\"");
        assert!(result.is_err());
    }

    #[test]
    fn python_json_roundtrip() {
        // Exact JSON that Python's Prompt.to_json() produces
        let python_json = r#"{"blocks": [{"role": "instruction", "content": "You are a test system.", "priority": 255}, {"role": "context", "content": "Some context data here", "label": "Data", "priority": 128}, {"role": "focus", "content": "Do the thing", "label": "Task", "priority": 192}, {"role": "constraint", "content": "Be concise", "priority": 200}, {"role": "schema", "content": "{\n  \"result\": \"string\"\n}", "priority": 240}]}"#;

        let prompt: Prompt =
            serde_json::from_str(python_json).expect("Rust must deserialize Python's Prompt JSON");

        assert_eq!(prompt.blocks.len(), 5);

        let rendered = prompt.to_string();
        assert!(
            rendered.contains("You are a test system."),
            "instruction missing"
        );
        assert!(
            rendered.contains("Some context data here"),
            "context missing"
        );
        assert!(rendered.contains("Do the thing"), "focus missing");
        assert!(rendered.contains("Be concise"), "constraint missing");
        assert!(
            rendered.contains("\"result\": \"string\""),
            "schema missing"
        );

        // Verify ordering: instruction < focus < context < constraint < schema
        let inst = rendered.find("You are a test system.").unwrap();
        let focus = rendered.find("Do the thing").unwrap();
        let ctx = rendered.find("Some context data here").unwrap();
        let rules = rendered.find("Rules:").unwrap();
        let schema = rendered.find("Return ONLY valid JSON").unwrap();
        assert!(inst < focus, "instruction should precede focus");
        assert!(focus < ctx, "focus should precede context");
        assert!(ctx < rules, "context should precede rules");
        assert!(rules < schema, "rules should precede schema");

        // Verify Python and Rust render identically
        let expected = "You are a test system.\n\n\
            ## Task\n\n\
            Do the thing\n\n\
            ## Data\n\n\
            Some context data here\n\n\
            Rules:\n\
            - Be concise\n\n\
            Return ONLY valid JSON (no markdown fences):\n\
            {\n  \"result\": \"string\"\n}";
        assert_eq!(
            rendered, expected,
            "Rust rendering must match Python fallback"
        );
    }

    #[test]
    fn full_prompt_example() {
        let p = Prompt::new()
            .instruction("You are a code reviewer.")
            .context("File", "src/main.rs: fn main() { ... }")
            .focus("Issue", "Potential null pointer dereference on line 42")
            .constraint("Do not suggest rewriting the entire function")
            .constraint("Keep suggestions under 3 lines")
            .example(
                "Good review",
                "Line 42: Add a null check before dereferencing",
            )
            .schema(serde_json::json!({
                "severity": "low|medium|high",
                "line": "integer",
                "suggestion": "string",
                "reasoning": "string"
            }));

        let rendered = p.to_string();

        // Verify ordering
        let positions: Vec<usize> = [
            "You are a code reviewer",
            "Potential null pointer",
            "src/main.rs",
            "Good review",
            "Rules:",
            "Return ONLY valid JSON",
        ]
        .iter()
        .map(|s| rendered.find(s).expect(&format!("missing: {}", s)))
        .collect();

        for i in 1..positions.len() {
            assert!(
                positions[i - 1] < positions[i],
                "ordering violated at index {}",
                i
            );
        }
    }
}