use crate::domain::{Chunk, Role};
use crate::query::Query;
use anyhow::{anyhow, Result};

// Legacy single-message template (backward compatible)
const SYSTEM_PROMPT: &str = include_str!("../../../templates/gru-mem.txt");

// Cache-aligned templates: split into stable prefix + varying suffix
const CACHE_SYSTEM: &str = include_str!("../../../templates/gru-mem-system.txt");
const CACHE_QUERY: &str = include_str!("../../../templates/gru-mem-query.txt");
const CACHE_TURN: &str = include_str!("../../../templates/gru-mem-turn.txt");

const BUDGET_TOTAL: usize = 32768;
const BUDGET_RESPONSE: usize = 2048;
const BUDGET_SYSTEM: usize = 400;
const BUDGET_QUESTION: usize = 150;
const BUDGET_MEMORY_MAX: usize = 1024;
const BUDGET_CHUNK_MAX: usize = 5000;

/// Prompt messages for the LLM API, supporting both legacy and cache-aligned modes.
///
/// Cache-aligned mode splits the prompt into separate messages so that the
/// stable prefix (system instructions + query) can be cached by the LLM provider,
/// while only the varying suffix (memory + chunk) changes per call.
#[derive(Debug, Clone)]
pub struct PromptMessages {
    /// System message: instructions (stable across all calls)
    pub system: String,
    /// Messages to send as user turns.
    /// In cache-aligned mode: [query (stable per run), turn (varies per call)]
    /// In legacy mode: [single combined message]
    pub user_messages: Vec<String>,
    /// Whether cache alignment is active
    pub cache_aligned: bool,
}

impl PromptMessages {
    /// The stable prefix length (system + query) that should be cache-marked.
    /// Returns 0 if not cache-aligned.
    pub fn cache_prefix_tokens(&self) -> usize {
        if !self.cache_aligned {
            return 0;
        }
        // Rough estimate: 4 chars ≈ 1 token
        let prefix_chars = self.system.len()
            + self.user_messages.first().map(|s| s.len()).unwrap_or(0);
        prefix_chars / 4
    }

    /// Total estimated tokens across all messages.
    pub fn total_tokens(&self) -> usize {
        let total_chars = self.system.len()
            + self.user_messages.iter().map(|s| s.len()).sum::<usize>();
        total_chars / 4
    }

    /// Headroom: tokens available for the response.
    pub fn headroom(&self) -> usize {
        let used = self.total_tokens();
        if used + BUDGET_RESPONSE > BUDGET_TOTAL {
            0
        } else {
            BUDGET_TOTAL - used - BUDGET_RESPONSE
        }
    }
}

/// Cache alignment metrics for observability and header generation.
#[derive(Debug, Clone)]
pub struct CacheMetrics {
    /// Size of stable prefix (bytes) that can be cached by LLM provider
    pub stable_prefix_bytes: usize,
    /// Size of dynamic tail (bytes) that varies per call
    pub dynamic_tail_bytes: usize,
    /// Drift metric: ratio of dynamic content (0.0 = stable, 1.0 = all dynamic)
    pub drift_metric: f32,
    /// Whether this chunk is cache-eligible (drift < 0.3)
    pub cache_eligible: bool,
    /// Estimated tokens in compressed form
    pub compressed_tokens: usize,
    /// Estimated tokens in original form
    pub original_tokens: usize,
}

impl CacheMetrics {
    /// Compression ratio as percentage
    pub fn compression_ratio(&self) -> f32 {
        if self.original_tokens == 0 {
            0.0
        } else {
            (self.compressed_tokens as f32 / self.original_tokens as f32) * 100.0
        }
    }

    /// Generate HTTP header value for drift
    pub fn header_drift(&self) -> String {
        format!("{:.2}", self.drift_metric)
    }

    /// Generate HTTP header value for eligible status
    pub fn header_eligible(&self) -> String {
        if self.cache_eligible { "true" } else { "false" }.to_string()
    }
}

/// Builds GRU-Mem prompts for the update gate.
///
/// Supports two modes:
/// - **Legacy** (`build`): Single user message with everything inlined.
///   Compatible with existing callers.
/// - **Cache-aligned** (`build_cache_aligned`): Splits into separate messages
///   so LLM providers can cache the stable prefix (instructions + query).
///
/// # Cache Alignment Strategy
///
/// ```text
/// ┌─────────────────────────────────────────────┐
/// │ SYSTEM MESSAGE (stable across ALL calls)    │ ← cached by provider
/// │   Instructions, rules, output format        │
/// ├─────────────────────────────────────────────┤
/// │ USER MSG 1: Query (stable per run)          │ ← cached by provider
/// │   <problem> What arch decisions? </problem> │
/// ├─────────────────────────────────────────────┤
/// │ USER MSG 2: Turn data (varies each call)    │ ← NOT cached (changes)
/// │   <memory> ... </memory>                    │
/// │   <section> ... </section>                  │
/// └─────────────────────────────────────────────┘
/// ```
///
/// For a run processing 50 chunks with the same query:
/// - Legacy: 0% cache hits (entire message changes every call)
/// - Cache-aligned: ~60-70% cache hits (system + query prefix reused)
///
/// At typical Anthropic pricing, this saves ~50% on input token costs
/// for multi-chunk ingestion runs.
pub struct PromptBuilder;

impl PromptBuilder {
    /// Calculate cache alignment metrics for a query and chunk.
    pub fn cache_metrics(query: &Query, chunk: &Chunk) -> Result<CacheMetrics> {
        use crate::optimizer::{ContextOptimizer, CacheAligner};

        let chunk_text = Self::render_chunk(chunk)?;
        let aligned = CacheAligner::align(&chunk_text);

        // Get compression metrics if optimizer is available
        let (original_tokens, compressed_tokens) = if let Ok(optimizer) = ContextOptimizer::from_env() {
            if let Ok(optimized) = optimizer.optimize(&chunk_text) {
                (optimized.original_tokens, optimized.compressed_tokens)
            } else {
                let tokens = estimate_tokens(&chunk_text);
                (tokens, tokens)
            }
        } else {
            let tokens = estimate_tokens(&chunk_text);
            (tokens, tokens)
        };

        Ok(CacheMetrics {
            stable_prefix_bytes: aligned.stable_prefix.len(),
            dynamic_tail_bytes: aligned.dynamic_tail.len(),
            drift_metric: aligned.drift_metric,
            cache_eligible: aligned.drift_metric < 0.3,
            original_tokens,
            compressed_tokens,
        })
    }

    /// Legacy build: single user message (backward compatible).
    ///
    /// Returns `(system_prompt, user_message)` tuple.
    pub fn build(query: &Query, previous_memory: Option<&str>, chunk: &Chunk) -> Result<(String, String)> {
        let chunk_text = Self::render_chunk(chunk)?;
        let chunk_bytes = chunk_text.len();
        let memory_text = previous_memory.unwrap_or("No previous memory");

        Self::check_budgets(memory_text, chunk_bytes)?;

        let user_message = SYSTEM_PROMPT
            .replace("{prompt}", &query.question)
            .replace("{memory}", memory_text)
            .replace("{chunk}", &chunk_text);

        let total_tokens = (SYSTEM_PROMPT.len() + query.question.len() + memory_text.len() + chunk_bytes) / 4;
        if total_tokens + BUDGET_RESPONSE > BUDGET_TOTAL {
            return Err(anyhow!(
                "Total prompt budget exceeded: {} + {} (response) > {} tokens",
                total_tokens, BUDGET_RESPONSE, BUDGET_TOTAL
            ));
        }

        Ok((SYSTEM_PROMPT.to_string(), user_message))
    }

    /// Cache-aligned build: splits prompt into cacheable prefix + varying suffix.
    ///
    /// The system message and query are stable across all chunks in a run,
    /// enabling LLM provider prompt caching. Only the turn message (memory +
    /// chunk) changes per call.
    ///
    /// # Returns
    /// `PromptMessages` with cache alignment metadata.
    pub fn build_cache_aligned(
        query: &Query,
        previous_memory: Option<&str>,
        chunk: &Chunk,
    ) -> Result<PromptMessages> {
        let chunk_text = Self::render_chunk(chunk)?;
        let chunk_bytes = chunk_text.len();
        let memory_text = previous_memory.unwrap_or("No previous memory");

        Self::check_budgets(memory_text, chunk_bytes)?;

        // System: stable instructions (same every call, every run)
        let system = CACHE_SYSTEM.to_string();

        // User message 1: query (stable per run — same across all chunks)
        let query_msg = CACHE_QUERY.replace("{prompt}", &query.question);

        // User message 2: turn data (varies every call)
        let turn_msg = CACHE_TURN
            .replace("{memory}", memory_text)
            .replace("{chunk}", &chunk_text);

        let messages = PromptMessages {
            system: system.clone(),
            user_messages: vec![query_msg, turn_msg],
            cache_aligned: true,
        };

        // Check total budget
        let total = messages.total_tokens();
        if total + BUDGET_RESPONSE > BUDGET_TOTAL {
            return Err(anyhow!(
                "Total prompt budget exceeded: {} + {} (response) > {} tokens",
                total, BUDGET_RESPONSE, BUDGET_TOTAL
            ));
        }

        Ok(messages)
    }

    /// Check memory and chunk budgets.
    fn check_budgets(memory_text: &str, chunk_bytes: usize) -> Result<()> {
        if memory_text.len() > BUDGET_MEMORY_MAX {
            return Err(anyhow!(
                "Memory budget exceeded: {} > {} tokens",
                memory_text.len() / 4,
                BUDGET_MEMORY_MAX / 4
            ));
        }
        if chunk_bytes > BUDGET_CHUNK_MAX {
            return Err(anyhow!(
                "Chunk budget exceeded: {} > {} bytes",
                chunk_bytes, BUDGET_CHUNK_MAX
            ));
        }
        Ok(())
    }

    /// Render a chunk as formatted text with role labels.
    fn render_chunk(chunk: &Chunk) -> Result<String> {
        let mut lines = Vec::new();

        for record in &chunk.records {
            let role_label = match record.role {
                Role::User => "[User]",
                Role::Assistant => "[Assistant]",
                Role::ToolResult => "[ToolResult]",
                Role::System => "[System]",
            };

            let text = format!("{} {}", role_label, record.text);
            lines.push(text);
        }

        Ok(lines.join("\n\n"))
    }
}

/// Simple token estimation (1 token ~= 4 chars or 1 word)
fn estimate_tokens(text: &str) -> usize {
    (text.len() / 4).max(text.split_whitespace().count())
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::domain::{Chunk, Record, Role, Provenance, Level};
    use time::OffsetDateTime;

    fn make_test_chunk(text: &str) -> Chunk {
        Chunk::new(
            1,
            vec![Record {
                role: Role::User,
                text: text.to_string(),
                timestamp: OffsetDateTime::now_utc(),
                provenance: Provenance {
                    source_id: "test".to_string(),
                    offset: 0,
                },
            }],
            10,
        )
    }

    #[test]
    fn test_render_chunk_single_record() {
        let chunk = make_test_chunk("Hello");
        let rendered = PromptBuilder::render_chunk(&chunk).unwrap();
        assert!(rendered.contains("[User]"));
        assert!(rendered.contains("Hello"));
    }

    #[test]
    fn test_render_chunk_multiple_roles() {
        let chunk = Chunk::new(
            1,
            vec![
                Record {
                    role: Role::User,
                    text: "What is 2+2?".to_string(),
                    timestamp: OffsetDateTime::now_utc(),
                    provenance: Provenance {
                        source_id: "test".to_string(),
                        offset: 0,
                    },
                },
                Record {
                    role: Role::Assistant,
                    text: "The answer is 4".to_string(),
                    timestamp: OffsetDateTime::now_utc(),
                    provenance: Provenance {
                        source_id: "test".to_string(),
                        offset: 0,
                    },
                },
                Record {
                    role: Role::ToolResult,
                    text: "Tool confirmed: 4".to_string(),
                    timestamp: OffsetDateTime::now_utc(),
                    provenance: Provenance {
                        source_id: "test".to_string(),
                        offset: 0,
                    },
                },
            ],
            30,
        );

        let rendered = PromptBuilder::render_chunk(&chunk).unwrap();
        assert!(rendered.contains("[User]"));
        assert!(rendered.contains("[Assistant]"));
        assert!(rendered.contains("[ToolResult]"));
        assert!(rendered.contains("\n\n"));
    }

    // ── Cache alignment tests ──────────────────────────────────────

    #[test]
    fn test_cache_aligned_produces_two_user_messages() {
        let query = Query {
            id: "test".to_string(),
            question: "What decisions?".to_string(),
            exit_gate: false,
        };
        let chunk = make_test_chunk("Some evidence");

        let msgs = PromptBuilder::build_cache_aligned(&query, None, &chunk).unwrap();

        assert!(msgs.cache_aligned);
        assert_eq!(msgs.user_messages.len(), 2, "Should have query + turn messages");
    }

    #[test]
    fn test_cache_prefix_is_stable_across_chunks() {
        let query = Query {
            id: "test".to_string(),
            question: "What decisions?".to_string(),
            exit_gate: false,
        };

        let chunk_a = make_test_chunk("Evidence chunk A");
        let chunk_b = make_test_chunk("Evidence chunk B - completely different");

        let msgs_a = PromptBuilder::build_cache_aligned(&query, None, &chunk_a).unwrap();
        let msgs_b = PromptBuilder::build_cache_aligned(&query, None, &chunk_b).unwrap();

        // System messages must be identical
        assert_eq!(msgs_a.system, msgs_b.system, "System prompt must be stable");

        // Query messages (user_messages[0]) must be identical
        assert_eq!(
            msgs_a.user_messages[0], msgs_b.user_messages[0],
            "Query message must be stable across chunks"
        );

        // Turn messages (user_messages[1]) must differ
        assert_ne!(
            msgs_a.user_messages[1], msgs_b.user_messages[1],
            "Turn messages should differ (different chunks)"
        );
    }

    #[test]
    fn test_cache_prefix_is_stable_across_memory_changes() {
        let query = Query {
            id: "test".to_string(),
            question: "What decisions?".to_string(),
            exit_gate: false,
        };
        let chunk = make_test_chunk("Same chunk");

        let msgs_t1 = PromptBuilder::build_cache_aligned(&query, None, &chunk).unwrap();
        let msgs_t2 = PromptBuilder::build_cache_aligned(
            &query,
            Some("Memory from turn 1"),
            &chunk,
        ).unwrap();

        // System + query must be identical even when memory changes
        assert_eq!(msgs_t1.system, msgs_t2.system);
        assert_eq!(msgs_t1.user_messages[0], msgs_t2.user_messages[0]);

        // Turn messages differ (different memory)
        assert_ne!(msgs_t1.user_messages[1], msgs_t2.user_messages[1]);
    }

    #[test]
    fn test_cache_prefix_tokens_positive() {
        let query = Query {
            id: "test".to_string(),
            question: "What decisions?".to_string(),
            exit_gate: false,
        };
        let chunk = make_test_chunk("Evidence");
        let msgs = PromptBuilder::build_cache_aligned(&query, None, &chunk).unwrap();

        assert!(
            msgs.cache_prefix_tokens() > 0,
            "Cache prefix should have positive token count"
        );
    }

    #[test]
    fn test_headroom_positive_under_budget() {
        let query = Query {
            id: "test".to_string(),
            question: "Test?".to_string(),
            exit_gate: false,
        };
        let chunk = make_test_chunk("Short");
        let msgs = PromptBuilder::build_cache_aligned(&query, None, &chunk).unwrap();

        assert!(
            msgs.headroom() > 0,
            "Should have positive headroom for short prompts"
        );
    }

    #[test]
    fn test_legacy_build_still_works() {
        let query = Query {
            id: "test".to_string(),
            question: "What?".to_string(),
            exit_gate: false,
        };
        let chunk = make_test_chunk("Evidence");

        let (system, user) = PromptBuilder::build(&query, None, &chunk).unwrap();
        assert!(!system.is_empty());
        assert!(!user.is_empty());
        assert!(user.contains("What?"));
        assert!(user.contains("Evidence"));
    }

    #[test]
    fn test_cache_aligned_contains_query() {
        let query = Query {
            id: "test".to_string(),
            question: "What architectural decisions?".to_string(),
            exit_gate: false,
        };
        let chunk = make_test_chunk("We use microservices");

        let msgs = PromptBuilder::build_cache_aligned(&query, None, &chunk).unwrap();

        // Query message should contain the problem
        assert!(
            msgs.user_messages[0].contains("What architectural decisions?"),
            "Query message should contain the question"
        );

        // Turn message should contain chunk and memory
        assert!(
            msgs.user_messages[1].contains("We use microservices"),
            "Turn message should contain chunk text"
        );
        assert!(
            msgs.user_messages[1].contains("No previous memory"),
            "Turn message should contain memory placeholder"
        );
    }

    #[test]
    fn test_cache_aligned_memory_budget_exceeded() {
        let query = Query {
            id: "test".to_string(),
            question: "Q?".to_string(),
            exit_gate: false,
        };
        let chunk = make_test_chunk("E");
        let big_memory = "x".repeat(BUDGET_MEMORY_MAX + 1);

        let result = PromptBuilder::build_cache_aligned(&query, Some(&big_memory), &chunk);
        assert!(result.is_err(), "Should reject over-budget memory");
    }

    #[test]
    fn test_cache_aligned_chunk_budget_exceeded() {
        let query = Query {
            id: "test".to_string(),
            question: "Q?".to_string(),
            exit_gate: false,
        };
        let big_chunk = make_test_chunk(&"x".repeat(BUDGET_CHUNK_MAX + 1));

        let result = PromptBuilder::build_cache_aligned(&query, None, &big_chunk);
        assert!(result.is_err(), "Should reject over-budget chunk");
    }

    // ── Cache metrics tests ────────────────────────────────────────

    #[test]
    fn test_cache_metrics_stable_query() {
        let query = Query {
            id: "q1".to_string(),
            question: "What happened?".to_string(),
            exit_gate: false,
        };
        let chunk = make_test_chunk("Plain stable content");

        let result = PromptBuilder::cache_metrics(&query, &chunk);
        assert!(result.is_ok());

        let metrics = result.unwrap();
        assert!(metrics.stable_prefix_bytes > 0);
        assert!(metrics.drift_metric >= 0.0 && metrics.drift_metric <= 1.0);
    }

    #[test]
    fn test_cache_metrics_compression_ratio() {
        let query = Query {
            id: "q2".to_string(),
            question: "Test?".to_string(),
            exit_gate: false,
        };
        let chunk = make_test_chunk("ERROR: failed\nINFO: debug");

        let result = PromptBuilder::cache_metrics(&query, &chunk);
        assert!(result.is_ok());

        let metrics = result.unwrap();
        let ratio = metrics.compression_ratio();
        assert!(ratio >= 0.0 && ratio <= 100.0);
    }

    #[test]
    fn test_cache_metrics_header_drift() {
        let query = Query {
            id: "q3".to_string(),
            question: "Q?".to_string(),
            exit_gate: false,
        };
        let chunk = make_test_chunk("content");

        let result = PromptBuilder::cache_metrics(&query, &chunk);
        assert!(result.is_ok());

        let metrics = result.unwrap();
        let drift_header = metrics.header_drift();
        let eligible_header = metrics.header_eligible();

        assert!(!drift_header.is_empty());
        assert!(eligible_header == "true" || eligible_header == "false");
    }
}