//! Advanced prompt template system with variable substitution, conditionals, and inheritance

use crate::error::{Result, TemplateError};
use dashmap::DashMap;
use indexmap::IndexMap;
use parking_lot::RwLock;
use regex::Regex;
use serde_json::Value;
use sha2::{Digest, Sha256};
use std::collections::HashMap;
use std::sync::Arc;
use std::time::Instant;
use uuid::Uuid;

/// Function type for custom template functions
pub type TemplateFunction = Box<dyn Fn() -> Result<Value> + Send + Sync>;

/// Filter type for custom template filters
pub type TemplateFilter = Box<dyn Fn(&Value, &str) -> Result<Value> + Send + Sync>;

/// Template context for variable substitution and function execution
#[derive(Clone)]
pub struct TemplateContext {
    variables: IndexMap<String, Value>,
    functions: Arc<RwLock<HashMap<String, TemplateFunction>>>,
    filters: Arc<RwLock<HashMap<String, TemplateFilter>>>,
}

impl std::fmt::Debug for TemplateContext {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("TemplateContext")
            .field("variables", &self.variables)
            .field(
                "functions",
                &format!("{} functions", self.functions.read().len()),
            )
            .field("filters", &format!("{} filters", self.filters.read().len()))
            .finish()
    }
}

impl TemplateContext {
    /// Create a new template context
    pub fn new() -> Self {
        Self {
            variables: IndexMap::new(),
            functions: Arc::new(RwLock::new(HashMap::new())),
            filters: Arc::new(RwLock::new(HashMap::new())),
        }
    }

    /// Insert a variable into the context
    pub fn insert(&mut self, key: impl Into<String>, value: Value) {
        self.variables.insert(key.into(), value);
    }

    /// Insert nested JSON structure
    pub fn insert_nested(&mut self, value: Value) {
        if let Value::Object(obj) = value {
            for (key, val) in obj {
                self.insert(key, val);
            }
        }
    }

    /// Get a variable from the context
    pub fn get(&self, key: &str) -> Option<&Value> {
        // Support nested access like "user.profile.name"
        let parts: Vec<&str> = key.split('.').collect();
        if parts.len() == 1 {
            return self.variables.get(key);
        }

        let mut current = self.variables.get(parts[0])?;
        for &part in &parts[1..] {
            current = current.get(part)?;
        }
        Some(current)
    }

    /// Register a custom function
    pub fn register_function(&mut self, name: impl Into<String>, func: TemplateFunction) {
        self.functions.write().insert(name.into(), func);
    }

    /// Register a custom filter
    pub fn register_filter(&mut self, name: impl Into<String>, filter: TemplateFilter) {
        self.filters.write().insert(name.into(), filter);
    }

    /// Call a registered function
    pub fn call_function(&self, name: &str) -> Result<Value> {
        let functions = self.functions.read();
        let func = functions
            .get(name)
            .ok_or_else(|| TemplateError::FunctionNotFound(name.to_string()))?;
        func()
    }

    /// Apply a registered filter
    pub fn apply_filter(&self, value: &Value, filter_name: &str, args: &str) -> Result<Value> {
        let filters = self.filters.read();
        let filter = filters
            .get(filter_name)
            .ok_or_else(|| TemplateError::FilterNotFound(filter_name.to_string()))?;
        filter(value, args)
    }
}

impl Default for TemplateContext {
    fn default() -> Self {
        Self::new()
    }
}

/// Compiled template node
#[derive(Debug, Clone)]
pub enum TemplateNode {
    Text(String),
    Variable {
        name: String,
        filters: Vec<(String, String)>, // (filter_name, args)
    },
    Conditional {
        condition: String,
        if_block: Vec<Self>,
        else_block: Option<Vec<Self>>,
    },
    Loop {
        variable: String,
        iterable: String,
        body: Vec<Self>,
    },
    Function {
        name: String,
        filters: Vec<(String, String)>,
    },
    Block {
        name: String,
        content: Vec<Self>,
    },
    Extends {
        parent: String,
    },
}

/// Prompt template with advanced features
#[derive(Debug, Clone)]
pub struct PromptTemplate {
    id: String,
    source: String,
    nodes: Vec<TemplateNode>,
    parent_template: Option<String>,
    blocks: HashMap<String, Vec<TemplateNode>>,
    creation_time: Instant,
}

impl PromptTemplate {
    /// Create template from string
    pub fn from_str(template: &str) -> Result<Self> {
        let id = format!("template_{}", Uuid::new_v4());
        let parser = TemplateParser::new();
        let nodes = parser.parse(template)?;

        let mut blocks = HashMap::new();
        let mut parent_template = None;

        // Extract blocks and parent template
        for node in &nodes {
            match node {
                TemplateNode::Block { name, content } => {
                    blocks.insert(name.clone(), content.clone());
                }
                TemplateNode::Extends { parent } => {
                    parent_template = Some(parent.clone());
                }
                _ => {}
            }
        }

        Ok(Self {
            id,
            source: template.to_string(),
            nodes,
            parent_template,
            blocks,
            creation_time: Instant::now(),
        })
    }

    /// Get template ID
    pub fn template_id(&self) -> &str {
        &self.id
    }

    /// Render template with context
    pub async fn render(&self, context: &TemplateContext) -> Result<String> {
        self.render_nodes(&self.nodes, context).await
    }

    /// Bind partial context to create a new template
    pub async fn bind_partial(&self, partial_context: &TemplateContext) -> Result<Self> {
        // Create a new template with some variables pre-filled
        let mut new_nodes = Vec::new();

        for node in &self.nodes {
            let processed_node = self.process_partial_binding(node, partial_context).await?;
            new_nodes.push(processed_node);
        }

        Ok(Self {
            id: format!("{}_partial", self.id),
            source: self.source.clone(),
            nodes: new_nodes,
            parent_template: self.parent_template.clone(),
            blocks: self.blocks.clone(),
            creation_time: Instant::now(),
        })
    }

    fn render_nodes<'a>(
        &'a self,
        nodes: &'a [TemplateNode],
        context: &'a TemplateContext,
    ) -> std::pin::Pin<Box<dyn std::future::Future<Output = Result<String>> + Send + 'a>> {
        Box::pin(async move {
            let mut result = String::new();

            for node in nodes {
                match node {
                    TemplateNode::Text(text) => {
                        result.push_str(text);
                    }
                    TemplateNode::Variable { name, filters } => {
                        let value = context
                            .get(name)
                            .ok_or_else(|| TemplateError::MissingVariable(name.clone()))?;

                        let mut processed_value = value.clone();
                        for (filter_name, args) in filters {
                            processed_value =
                                context.apply_filter(&processed_value, filter_name, args)?;
                        }

                        let text = self.value_to_html_escaped_string(&processed_value);
                        result.push_str(&text);
                    }
                    TemplateNode::Function { name, filters } => {
                        let mut value = context.call_function(name)?;
                        for (filter_name, args) in filters {
                            value = context.apply_filter(&value, filter_name, args)?;
                        }

                        let text = self.value_to_string(&value);
                        result.push_str(&text);
                    }
                    TemplateNode::Conditional {
                        condition,
                        if_block,
                        else_block,
                    } => {
                        let should_render = self.evaluate_condition(condition, context)?;
                        if should_render {
                            result.push_str(&self.render_nodes(if_block, context).await?);
                        } else if let Some(else_block) = else_block {
                            result.push_str(&self.render_nodes(else_block, context).await?);
                        }
                    }
                    TemplateNode::Loop {
                        variable,
                        iterable,
                        body,
                    } => {
                        let items = context
                            .get(iterable)
                            .ok_or_else(|| TemplateError::MissingVariable(iterable.clone()))?;

                        if let Value::Array(arr) = items {
                            for item in arr {
                                let mut loop_context = context.clone();
                                loop_context.insert(variable, item.clone());
                                result.push_str(&self.render_nodes(body, &loop_context).await?);
                            }
                        }
                    }
                    TemplateNode::Block { name: _, content } => {
                        result.push_str(&self.render_nodes(content, context).await?);
                    }
                    TemplateNode::Extends { .. } => {
                        // Template inheritance is handled at the engine level
                    }
                }
            }

            Ok(result)
        })
    }

    fn evaluate_condition(&self, condition: &str, context: &TemplateContext) -> Result<bool> {
        // Simple boolean evaluation
        if let Some(value) = context.get(condition) {
            return Ok(match value {
                Value::Bool(b) => *b,
                Value::Null => false,
                Value::Number(n) => n.as_f64().unwrap_or(0.0) != 0.0,
                Value::String(s) => !s.is_empty(),
                Value::Array(arr) => !arr.is_empty(),
                Value::Object(obj) => !obj.is_empty(),
            });
        }
        Ok(false)
    }

    fn value_to_string(&self, value: &Value) -> String {
        match value {
            Value::String(s) => s.clone(),
            Value::Number(n) => n.to_string(),
            Value::Bool(b) => b.to_string(),
            Value::Null => String::new(),
            _ => serde_json::to_string(value).unwrap_or_default(),
        }
    }

    fn value_to_html_escaped_string(&self, value: &Value) -> String {
        let text = self.value_to_string(value);
        // Basic HTML escaping for security
        text.replace('&', "&amp;")
            .replace('<', "&lt;")
            .replace('>', "&gt;")
            .replace('"', "&quot;")
            .replace('\'', "&#39;")
    }

    async fn process_partial_binding(
        &self,
        node: &TemplateNode,
        context: &TemplateContext,
    ) -> Result<TemplateNode> {
        match node {
            TemplateNode::Variable { name, filters: _ } => {
                if let Some(value) = context.get(name) {
                    // Replace with text node if variable is bound
                    let text = self.value_to_string(value);
                    Ok(TemplateNode::Text(text))
                } else {
                    // Keep as variable if not bound
                    Ok(node.clone())
                }
            }
            _ => Ok(node.clone()), // For simplicity, only handle variables in partial binding
        }
    }
}

/// Template parser for converting template strings into AST
struct TemplateParser {
    variable_regex: Regex,
    conditional_regex: Regex,
    loop_regex: Regex,
    function_regex: Regex,
    block_regex: Regex,
}

impl TemplateParser {
    fn new() -> Self {
        Self {
            variable_regex: Regex::new(r"\{([^}]+)\}").unwrap(),
            conditional_regex: Regex::new(r"\{%-?\s*if\s+([^%}]+)\s*-?%\}").unwrap(),
            loop_regex: Regex::new(r"\{%-?\s*for\s+(\w+)\s+in\s+(\w+)\s*-?%\}").unwrap(),
            function_regex: Regex::new(r"\{([^}]+\(\)(?:\s*\|\s*[^}]+)*)\}").unwrap(),
            block_regex: Regex::new(r"\{%-?\s*block\s+(\w+)\s*-?%\}").unwrap(),
        }
    }

    fn parse(&self, template: &str) -> Result<Vec<TemplateNode>> {
        let mut nodes = Vec::new();
        let mut remaining = template;

        while !remaining.is_empty() {
            if let Some(next_node) = self.parse_next_node(&mut remaining)? {
                nodes.push(next_node);
            } else {
                // If no special syntax found, treat rest as text
                if !remaining.is_empty() {
                    nodes.push(TemplateNode::Text(remaining.to_string()));
                    break;
                }
            }
        }

        Ok(nodes)
    }

    fn parse_next_node(&self, remaining: &mut &str) -> Result<Option<TemplateNode>> {
        if remaining.is_empty() {
            return Ok(None);
        }

        // Find the earliest position of any special syntax
        let variable_pos = self.variable_regex.find(remaining).map(|m| m.start());
        let conditional_pos = self.conditional_regex.find(remaining).map(|m| m.start());
        let loop_pos = self.loop_regex.find(remaining).map(|m| m.start());
        let block_pos = self.block_regex.find(remaining).map(|m| m.start());
        let function_pos = self.function_regex.find(remaining).map(|m| m.start());
        let extends_pos = if remaining.starts_with("{%- extends") {
            Some(0)
        } else {
            None
        };

        // Find the earliest special syntax
        let earliest_pos = [
            variable_pos,
            conditional_pos,
            loop_pos,
            block_pos,
            function_pos,
            extends_pos,
        ]
        .iter()
        .filter_map(|&pos| pos)
        .min();

        if let Some(pos) = earliest_pos {
            // If there's text before the special syntax, return that first
            if pos > 0 {
                let text = remaining[..pos].to_string();
                *remaining = &remaining[pos..];
                return Ok(Some(TemplateNode::Text(text)));
            }

            // Otherwise, parse the special syntax
            if extends_pos == Some(0) {
                return self.parse_extends(remaining);
            } else if conditional_pos == Some(0) {
                return self.parse_conditional(remaining);
            } else if loop_pos == Some(0) {
                return self.parse_loop(remaining);
            } else if block_pos == Some(0) {
                return self.parse_block(remaining);
            } else if function_pos == Some(0) {
                return self.parse_function(remaining);
            } else if variable_pos == Some(0) {
                return self.parse_variable(remaining);
            }
        } else {
            // No special syntax found, treat remaining as text
            if !remaining.is_empty() {
                let text = remaining.to_string();
                *remaining = "";
                return Ok(Some(TemplateNode::Text(text)));
            }
        }

        Ok(None)
    }

    fn parse_variable(&self, remaining: &mut &str) -> Result<Option<TemplateNode>> {
        if let Some(caps) = self.variable_regex.captures(remaining) {
            let full_match = caps.get(0).unwrap();
            let variable_content = caps.get(1).unwrap().as_str().trim();

            // Parse filters
            let parts: Vec<&str> = variable_content.split('|').collect();
            let name = parts[0].trim().to_string();
            let mut filters = Vec::new();

            for filter_part in parts.iter().skip(1) {
                let filter_parts: Vec<&str> = filter_part.trim().splitn(2, '(').collect();
                let filter_name = filter_parts[0].trim().to_string();
                let args = if filter_parts.len() > 1 {
                    filter_parts[1].trim_end_matches(')').to_string()
                } else {
                    String::new()
                };
                filters.push((filter_name, args));
            }

            *remaining = &remaining[full_match.end()..];
            return Ok(Some(TemplateNode::Variable { name, filters }));
        }
        Ok(None)
    }

    fn parse_conditional(&self, remaining: &mut &str) -> Result<Option<TemplateNode>> {
        if let Some(caps) = self.conditional_regex.captures(remaining) {
            let condition = caps.get(1).unwrap().as_str().trim().to_string();
            let if_start = caps.get(0).unwrap().end();

            // Find endif or else
            let if_content_start = if_start;
            let (else_pos, endif_pos) = self.find_conditional_boundaries(&remaining[if_start..])?;

            let if_content =
                &remaining[if_content_start..if_content_start + else_pos.unwrap_or(endif_pos)];
            let if_block = self.parse(if_content)?;

            let else_block = if let Some(else_start) = else_pos {
                let else_content =
                    &remaining[if_content_start + else_start..if_content_start + endif_pos];
                // Skip the else tag
                let else_tag_end = else_content.find("%}").unwrap_or(0) + 2;
                Some(self.parse(&else_content[else_tag_end..])?)
            } else {
                None
            };

            *remaining = &remaining[if_content_start + endif_pos + self.find_endif_tag_length()..];

            return Ok(Some(TemplateNode::Conditional {
                condition,
                if_block,
                else_block,
            }));
        }
        Ok(None)
    }

    fn parse_loop(&self, remaining: &mut &str) -> Result<Option<TemplateNode>> {
        if let Some(caps) = self.loop_regex.captures(remaining) {
            let variable = caps.get(1).unwrap().as_str().to_string();
            let iterable = caps.get(2).unwrap().as_str().to_string();
            let loop_start = caps.get(0).unwrap().end();

            let endfor_pos = self.find_endfor_position(&remaining[loop_start..])?;
            let loop_content = &remaining[loop_start..loop_start + endfor_pos];
            let body = self.parse(loop_content)?;

            *remaining = &remaining[loop_start + endfor_pos + self.find_endfor_tag_length()..];

            return Ok(Some(TemplateNode::Loop {
                variable,
                iterable,
                body,
            }));
        }
        Ok(None)
    }

    fn parse_block(&self, remaining: &mut &str) -> Result<Option<TemplateNode>> {
        if let Some(caps) = self.block_regex.captures(remaining) {
            let name = caps.get(1).unwrap().as_str().to_string();
            let block_start = caps.get(0).unwrap().end();

            let endblock_pos = self.find_endblock_position(&remaining[block_start..])?;
            let block_content = &remaining[block_start..block_start + endblock_pos];
            let content = self.parse(block_content)?;

            *remaining = &remaining[block_start + endblock_pos + self.find_endblock_tag_length()..];

            return Ok(Some(TemplateNode::Block { name, content }));
        }
        Ok(None)
    }

    fn parse_function(&self, remaining: &mut &str) -> Result<Option<TemplateNode>> {
        if let Some(caps) = self.function_regex.captures(remaining) {
            let function_content = caps.get(1).unwrap().as_str().trim();

            // Parse function name and filters
            let parts: Vec<&str> = function_content.split('|').collect();
            let func_call = parts[0].trim();
            let name = func_call.trim_end_matches("()").to_string();

            let mut filters = Vec::new();
            for filter_part in parts.iter().skip(1) {
                let filter_parts: Vec<&str> = filter_part.trim().splitn(2, '(').collect();
                let filter_name = filter_parts[0].trim().to_string();
                let args = if filter_parts.len() > 1 {
                    filter_parts[1]
                        .trim_end_matches(')')
                        .trim_matches('\'')
                        .to_string()
                } else {
                    String::new()
                };
                filters.push((filter_name, args));
            }

            *remaining = &remaining[caps.get(0).unwrap().end()..];
            return Ok(Some(TemplateNode::Function { name, filters }));
        }
        Ok(None)
    }

    fn parse_extends(&self, remaining: &mut &str) -> Result<Option<TemplateNode>> {
        if let Some(end) = remaining.find("%}") {
            let extends_content = &remaining[..end];
            if let Some(parent_start) = extends_content.rfind(' ') {
                let parent = extends_content[parent_start + 1..].trim().to_string();
                *remaining = &remaining[end + 2..];
                return Ok(Some(TemplateNode::Extends { parent }));
            }
        }
        Err(TemplateError::SyntaxError("Invalid extends syntax".to_string()).into())
    }

    fn find_conditional_boundaries(&self, content: &str) -> Result<(Option<usize>, usize)> {
        let else_regex = Regex::new(r"\{%-?\s*else\s*-?%\}").unwrap();
        let endif_regex = Regex::new(r"\{%-?\s*endif\s*-?%\}").unwrap();

        let else_pos = else_regex.find(content).map(|m| m.start());
        let endif_pos = endif_regex
            .find(content)
            .ok_or_else(|| TemplateError::SyntaxError("Missing endif".to_string()))?
            .start();

        Ok((else_pos, endif_pos))
    }

    fn find_endfor_position(&self, content: &str) -> Result<usize> {
        let endfor_regex = Regex::new(r"\{%-?\s*endfor\s*-?%\}").unwrap();
        endfor_regex
            .find(content)
            .map(|m| m.start())
            .ok_or_else(|| TemplateError::SyntaxError("Missing endfor".to_string()).into())
    }

    fn find_endblock_position(&self, content: &str) -> Result<usize> {
        let endblock_regex = Regex::new(r"\{%-?\s*endblock\s*-?%\}").unwrap();
        endblock_regex
            .find(content)
            .map(|m| m.start())
            .ok_or_else(|| TemplateError::SyntaxError("Missing endblock".to_string()).into())
    }

    fn find_endif_tag_length(&self) -> usize {
        "{%- endif -%}".len() // Approximate, should parse actual tag
    }

    fn find_endfor_tag_length(&self) -> usize {
        "{%- endfor -%}".len()
    }

    fn find_endblock_tag_length(&self) -> usize {
        "{%- endblock -%}".len()
    }
}

/// Template engine with caching and inheritance support
pub struct TemplateEngine {
    templates: Arc<DashMap<String, PromptTemplate>>,
    compiled_cache: Arc<DashMap<String, PromptTemplate>>,
    cache_enabled: bool,
}

impl TemplateEngine {
    /// Create a new template engine
    pub fn new() -> Self {
        Self {
            templates: Arc::new(DashMap::new()),
            compiled_cache: Arc::new(DashMap::new()),
            cache_enabled: false,
        }
    }

    /// Enable or disable template caching
    pub fn enable_caching(&mut self, enabled: bool) {
        self.cache_enabled = enabled;
    }

    /// Register a template with name
    pub async fn register_template(&mut self, name: &str, template: &str) -> Result<()> {
        let compiled_template = PromptTemplate::from_str(template)?;
        self.templates.insert(name.to_string(), compiled_template);
        Ok(())
    }

    /// Compile a template with inheritance support
    pub async fn compile_template(&self, template: &str) -> Result<PromptTemplate> {
        // Check cache first
        if self.cache_enabled {
            let hash = self.compute_template_hash(template);
            if let Some(cached) = self.compiled_cache.get(&hash) {
                return Ok(cached.clone());
            }
        }

        let mut compiled = PromptTemplate::from_str(template)?;

        // Handle template inheritance
        if let Some(parent_name) = &compiled.parent_template
            && let Some(parent_template) = self.templates.get(parent_name)
        {
            compiled = self.merge_with_parent(&compiled, &parent_template)?;
        }

        // Cache the compiled template
        if self.cache_enabled {
            let hash = self.compute_template_hash(template);
            self.compiled_cache.insert(hash, compiled.clone());
        }

        Ok(compiled)
    }

    fn merge_with_parent(
        &self,
        child: &PromptTemplate,
        parent: &PromptTemplate,
    ) -> Result<PromptTemplate> {
        let mut merged_nodes = parent.nodes.clone();

        // Replace blocks with child's version
        for (block_name, child_block) in &child.blocks {
            // Find and replace the corresponding block in parent
            self.replace_block_in_nodes(&mut merged_nodes, block_name, child_block);
        }

        Ok(PromptTemplate {
            id: format!("{}_extended", child.id),
            source: format!("{}\n{}", parent.source, child.source),
            nodes: merged_nodes,
            parent_template: None,
            blocks: child.blocks.clone(),
            creation_time: Instant::now(),
        })
    }

    fn replace_block_in_nodes(
        &self,
        nodes: &mut Vec<TemplateNode>,
        block_name: &str,
        replacement: &[TemplateNode],
    ) {
        for node in nodes.iter_mut() {
            if let TemplateNode::Block { name, content } = node
                && name == block_name
            {
                *content = replacement.to_vec();
            }
        }
    }

    fn compute_template_hash(&self, template: &str) -> String {
        let mut hasher = Sha256::new();
        hasher.update(template.as_bytes());
        format!("{:x}", hasher.finalize())
    }
}

impl Default for TemplateEngine {
    fn default() -> Self {
        Self::new()
    }
}

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

    #[tokio::test]
    async fn test_basic_template_parsing() {
        let template = PromptTemplate::from_str("Hello {name}!").unwrap();
        let mut context = TemplateContext::new();
        context.insert("name", Value::String("World".to_string()));

        let result = template.render(&context).await.unwrap();
        assert_eq!(result, "Hello World!");
    }
}