use super::*;

use crate::client::{Message, MessageContent, MessageRole, Model};

use anyhow::Result;
use fancy_regex::Regex;
use rust_embed::Embed;
use serde::{Deserialize, Serialize};
use serde_json::Value;
use std::sync::LazyLock;

pub const SHELL_ROLE: &str = "shell";
pub const EXPLAIN_SHELL_ROLE: &str = "explain-shell";
pub const CODE_ROLE: &str = "code";
pub const CREATE_TITLE_ROLE: &str = "create-title";

pub const INPUT_PLACEHOLDER: &str = "__INPUT__";

#[derive(Embed)]
#[folder = "assets/roles/"]
struct RolesAsset;

static RE_METADATA: LazyLock<Regex> =
    LazyLock::new(|| Regex::new(r"(?s)-{3,}\s*(.*?)\s*-{3,}\s*(.*)").unwrap());

pub trait RoleLike {
    fn to_role(&self) -> Role;
    fn model(&self) -> &Model;
    fn temperature(&self) -> Option<f64>;
    fn top_p(&self) -> Option<f64>;
    fn enabled_tools(&self) -> Option<Vec<String>>;
    fn enabled_mcp_servers(&self) -> Option<Vec<String>>;
    fn set_model(&mut self, model: Model);
    fn set_temperature(&mut self, value: Option<f64>);
    fn reasoning_effort(&self) -> Option<String>;
    fn set_top_p(&mut self, value: Option<f64>);
    fn set_reasoning_effort(&mut self, value: Option<String>);
    fn set_enabled_tools(&mut self, value: Option<Vec<String>>);
    fn set_enabled_mcp_servers(&mut self, value: Option<Vec<String>>);
}

#[derive(Debug, Clone, Default, Deserialize, Serialize)]
pub struct Role {
    name: String,
    #[serde(default)]
    prompt: String,
    #[serde(
        rename(serialize = "model", deserialize = "model"),
        skip_serializing_if = "Option::is_none"
    )]
    model_id: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    temperature: Option<f64>,
    #[serde(skip_serializing_if = "Option::is_none")]
    top_p: Option<f64>,
    #[serde(skip_serializing_if = "Option::is_none")]
    reasoning_effort: Option<String>,
    #[serde(
        default,
        skip_serializing_if = "Option::is_none",
        deserialize_with = "super::deserialize_csv_or_vec"
    )]
    enabled_tools: Option<Vec<String>>,
    #[serde(
        default,
        skip_serializing_if = "Option::is_none",
        deserialize_with = "super::deserialize_csv_or_vec"
    )]
    enabled_mcp_servers: Option<Vec<String>>,
    #[serde(skip_serializing_if = "Option::is_none")]
    skills_enabled: Option<bool>,
    #[serde(
        default,
        skip_serializing_if = "Option::is_none",
        deserialize_with = "super::deserialize_csv_or_vec"
    )]
    enabled_skills: Option<Vec<String>>,
    #[serde(skip_serializing_if = "Option::is_none")]
    auto_continue: Option<bool>,
    #[serde(skip_serializing_if = "Option::is_none")]
    max_auto_continues: Option<usize>,
    #[serde(skip_serializing_if = "Option::is_none")]
    inject_todo_instructions: Option<bool>,
    #[serde(skip_serializing_if = "Option::is_none")]
    continuation_prompt: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    inject_skill_instructions: Option<bool>,
    #[serde(skip_serializing_if = "Option::is_none")]
    skill_instructions: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    memory: Option<bool>,

    #[serde(skip)]
    model: Model,
}

impl Role {
    pub fn new(name: &str, content: &str) -> Self {
        let mut metadata = "";
        let mut prompt = content.trim();
        if let Ok(Some(caps)) = RE_METADATA.captures(content)
            && let (Some(metadata_value), Some(prompt_value)) = (caps.get(1), caps.get(2))
        {
            metadata = metadata_value.as_str().trim();
            prompt = prompt_value.as_str().trim();
        }
        let mut prompt = prompt.to_string();
        interpolate_variables(&mut prompt);
        let mut role = Self {
            name: name.to_string(),
            prompt,
            ..Default::default()
        };
        if !metadata.is_empty()
            && let Ok(value) = serde_yaml::from_str::<Value>(metadata)
            && let Some(value) = value.as_object()
        {
            for (key, value) in value {
                match key.as_str() {
                    "model" => role.model_id = value.as_str().map(|v| v.to_string()),
                    "temperature" => role.temperature = value.as_f64(),
                    "top_p" => role.top_p = value.as_f64(),
                    "reasoning_effort" => {
                        role.reasoning_effort = value.as_str().map(|v| v.to_string())
                    }
                    "enabled_tools" => role.enabled_tools = parse_string_or_array(value),
                    "enabled_mcp_servers" => {
                        role.enabled_mcp_servers = parse_string_or_array(value)
                    }
                    "skills_enabled" => role.skills_enabled = value.as_bool(),
                    "enabled_skills" => role.enabled_skills = parse_string_or_array(value),
                    "auto_continue" => role.auto_continue = value.as_bool(),
                    "max_auto_continues" => {
                        role.max_auto_continues = value.as_u64().map(|v| v as usize)
                    }
                    "inject_todo_instructions" => role.inject_todo_instructions = value.as_bool(),
                    "continuation_prompt" => {
                        role.continuation_prompt = value.as_str().map(|v| v.to_string())
                    }
                    "inject_skill_instructions" => role.inject_skill_instructions = value.as_bool(),
                    "skill_instructions" => {
                        role.skill_instructions = value.as_str().map(|v| v.to_string())
                    }
                    "memory" => role.memory = value.as_bool(),
                    _ => (),
                }
            }
        }
        role
    }

    pub fn builtin(name: &str) -> Result<Self> {
        let content = RolesAsset::get(&format!("{name}.md"))
            .ok_or_else(|| anyhow!("Unknown role `{name}`"))?;
        let content = unsafe { std::str::from_utf8_unchecked(&content.data) };
        Ok(Role::new(name, content))
    }

    pub fn list_builtin_role_names() -> Vec<String> {
        RolesAsset::iter()
            .filter_map(|v| v.strip_suffix(".md").map(|v| v.to_string()))
            .collect()
    }

    pub fn has_args(&self) -> bool {
        self.name.contains('#')
    }

    pub fn export(&self) -> String {
        let mut metadata = vec![];
        if let Some(model) = self.model_id() {
            metadata.push(format!("model: {model}"));
        }
        if let Some(temperature) = self.temperature() {
            metadata.push(format!("temperature: {temperature}"));
        }
        if let Some(top_p) = self.top_p() {
            metadata.push(format!("top_p: {top_p}"));
        }
        if let Some(reasoning_effort) = self.reasoning_effort() {
            metadata.push(format!("reasoning_effort: {reasoning_effort}"));
        }
        if let Some(enabled_tools) = &self.enabled_tools {
            let inline = serde_json::to_string(enabled_tools).unwrap_or_else(|_| "[]".to_string());
            metadata.push(format!("enabled_tools: {inline}"));
        }
        if let Some(enabled_mcp_servers) = &self.enabled_mcp_servers {
            let inline =
                serde_json::to_string(enabled_mcp_servers).unwrap_or_else(|_| "[]".to_string());
            metadata.push(format!("enabled_mcp_servers: {inline}"));
        }
        if let Some(skills_enabled) = self.skills_enabled {
            metadata.push(format!("skills_enabled: {skills_enabled}"));
        }
        if let Some(enabled_skills) = &self.enabled_skills {
            let inline = serde_json::to_string(enabled_skills).unwrap_or_else(|_| "[]".to_string());
            metadata.push(format!("enabled_skills: {inline}"));
        }
        if let Some(auto_continue) = self.auto_continue {
            metadata.push(format!("auto_continue: {auto_continue}"));
        }
        if let Some(max_auto_continues) = self.max_auto_continues {
            metadata.push(format!("max_auto_continues: {max_auto_continues}"));
        }
        if let Some(inject_todo_instructions) = self.inject_todo_instructions {
            metadata.push(format!(
                "inject_todo_instructions: {inject_todo_instructions}"
            ));
        }
        if let Some(continuation_prompt) = &self.continuation_prompt {
            metadata.push(format!("continuation_prompt: {continuation_prompt}"));
        }
        if let Some(inject_skill_instructions) = self.inject_skill_instructions {
            metadata.push(format!(
                "inject_skill_instructions: {inject_skill_instructions}"
            ));
        }
        if let Some(skill_instructions) = &self.skill_instructions {
            metadata.push(format!("skill_instructions: {skill_instructions}"));
        }
        if let Some(memory) = self.memory {
            metadata.push(format!("memory: {memory}"));
        }
        if metadata.is_empty() {
            format!("{}\n", self.prompt)
        } else if self.prompt.is_empty() {
            format!("---\n{}\n---\n", metadata.join("\n"))
        } else {
            format!("---\n{}\n---\n\n{}\n", metadata.join("\n"), self.prompt)
        }
    }

    pub fn save(&mut self, role_name: &str, role_path: &Path, is_repl: bool) -> Result<()> {
        ensure_parent_exists(role_path)?;

        let content = self.export();
        std::fs::write(role_path, content).with_context(|| {
            format!(
                "Failed to write role {} to {}",
                self.name,
                role_path.display()
            )
        })?;

        if is_repl {
            println!("✓ Saved role to '{}'.", role_path.display());
        }

        if role_name != self.name {
            self.name = role_name.to_string();
        }

        Ok(())
    }

    pub fn sync<T: RoleLike>(&mut self, role_like: &T) {
        let model = role_like.model();
        let reasoning_effort = role_like.reasoning_effort();
        let temperature = role_like.temperature();
        let top_p = role_like.top_p();
        let enabled_tools = role_like.enabled_tools();
        let enabled_mcp_servers = role_like.enabled_mcp_servers();
        self.batch_set(
            model,
            reasoning_effort,
            temperature,
            top_p,
            enabled_tools,
            enabled_mcp_servers,
        );
    }

    pub fn batch_set(
        &mut self,
        model: &Model,
        reasoning_effort: Option<String>,
        temperature: Option<f64>,
        top_p: Option<f64>,
        enabled_tools: Option<Vec<String>>,
        enabled_mcp_servers: Option<Vec<String>>,
    ) {
        self.set_model(model.clone());
        if reasoning_effort.is_some() {
            self.set_reasoning_effort(reasoning_effort.clone());
        }
        if temperature.is_some() {
            self.set_temperature(temperature);
        }
        if top_p.is_some() {
            self.set_top_p(top_p);
        }
        if enabled_tools.is_some() {
            self.set_enabled_tools(enabled_tools);
        }
        if enabled_mcp_servers.is_some() {
            self.set_enabled_mcp_servers(enabled_mcp_servers);
        }
    }

    pub fn is_derived(&self) -> bool {
        self.name.is_empty()
    }

    pub fn name(&self) -> &str {
        &self.name
    }

    pub fn model_id(&self) -> Option<&str> {
        self.model_id.as_deref()
    }

    pub fn prompt(&self) -> &str {
        &self.prompt
    }

    pub fn is_empty_prompt(&self) -> bool {
        self.prompt.is_empty()
    }

    pub fn is_embedded_prompt(&self) -> bool {
        self.prompt.contains(INPUT_PLACEHOLDER)
    }

    pub fn auto_continue(&self) -> Option<bool> {
        self.auto_continue
    }

    pub fn max_auto_continues(&self) -> Option<usize> {
        self.max_auto_continues
    }

    pub fn inject_todo_instructions(&self) -> Option<bool> {
        self.inject_todo_instructions
    }

    pub fn continuation_prompt(&self) -> Option<&str> {
        self.continuation_prompt.as_deref()
    }

    pub fn inject_skill_instructions(&self) -> Option<bool> {
        self.inject_skill_instructions
    }

    pub fn skill_instructions(&self) -> Option<&str> {
        self.skill_instructions.as_deref()
    }

    pub fn memory(&self) -> Option<bool> {
        self.memory
    }

    pub fn skills_enabled(&self) -> Option<bool> {
        self.skills_enabled
    }

    pub fn enabled_skills(&self) -> Option<&[String]> {
        self.enabled_skills.as_deref()
    }

    pub fn append_to_prompt(&mut self, text: &str) {
        self.prompt.push_str(text);
    }

    pub fn echo_messages(&self, input: &Input) -> String {
        let input_markdown = input.render();
        if self.is_empty_prompt() {
            input_markdown
        } else if self.is_embedded_prompt() {
            self.prompt.replace(INPUT_PLACEHOLDER, &input_markdown)
        } else {
            format!("{}\n\n{}", self.prompt, input_markdown)
        }
    }

    pub fn build_messages(&self, input: &Input) -> Vec<Message> {
        let mut content = input.message_content();
        let mut messages = if self.is_empty_prompt() {
            vec![Message::new(MessageRole::User, content)]
        } else if self.is_embedded_prompt() {
            content.merge_prompt(|v: &str| self.prompt.replace(INPUT_PLACEHOLDER, v));
            vec![Message::new(MessageRole::User, content)]
        } else {
            let mut messages = vec![];
            let (system, cases) = parse_structure_prompt(&self.prompt);
            if !system.is_empty() {
                messages.push(Message::new(
                    MessageRole::System,
                    MessageContent::Text(system.to_string()),
                ));
            }
            if !cases.is_empty() {
                messages.extend(cases.into_iter().flat_map(|(i, o)| {
                    vec![
                        Message::new(MessageRole::User, MessageContent::Text(i.to_string())),
                        Message::new(MessageRole::Assistant, MessageContent::Text(o.to_string())),
                    ]
                }));
            }
            messages.push(Message::new(MessageRole::User, content));
            messages
        };
        if let Some(text) = input.continue_output() {
            messages.push(Message::new(
                MessageRole::Assistant,
                MessageContent::Text(text.into()),
            ));
        }
        messages
    }
}

impl RoleLike for Role {
    fn to_role(&self) -> Role {
        self.clone()
    }

    fn model(&self) -> &Model {
        &self.model
    }

    fn temperature(&self) -> Option<f64> {
        self.temperature
    }

    fn top_p(&self) -> Option<f64> {
        self.top_p
    }

    fn reasoning_effort(&self) -> Option<String> {
        self.reasoning_effort.clone()
    }

    fn enabled_tools(&self) -> Option<Vec<String>> {
        self.enabled_tools.clone()
    }

    fn enabled_mcp_servers(&self) -> Option<Vec<String>> {
        self.enabled_mcp_servers.clone()
    }

    fn set_model(&mut self, model: Model) {
        if !self.model().id().is_empty() {
            self.model_id = Some(model.id().to_string());
        }
        self.model = model;
    }

    fn set_temperature(&mut self, value: Option<f64>) {
        self.temperature = value;
    }

    fn set_top_p(&mut self, value: Option<f64>) {
        self.top_p = value;
    }

    fn set_reasoning_effort(&mut self, value: Option<String>) {
        self.reasoning_effort = value;
    }

    fn set_enabled_tools(&mut self, value: Option<Vec<String>>) {
        self.enabled_tools = value;
    }

    fn set_enabled_mcp_servers(&mut self, value: Option<Vec<String>>) {
        self.enabled_mcp_servers = value;
    }
}

fn parse_string_or_array(value: &Value) -> Option<Vec<String>> {
    if value.is_null() {
        return None;
    }

    if let Some(s) = value.as_str() {
        return Some(csv_to_vec(s));
    }

    if let Some(arr) = value.as_array() {
        let items: Vec<String> = arr
            .iter()
            .filter_map(|v| v.as_str().map(|s| s.trim().to_string()))
            .filter(|s| !s.is_empty())
            .collect();

        return Some(items);
    }

    None
}

fn parse_structure_prompt(prompt: &str) -> (&str, Vec<(&str, &str)>) {
    let mut text = prompt;
    let mut search_input = true;
    let mut system = None;
    let mut parts = vec![];
    loop {
        let search = if search_input {
            "### INPUT:"
        } else {
            "### OUTPUT:"
        };
        match text.find(search) {
            Some(idx) => {
                if system.is_none() {
                    system = Some(&text[..idx])
                } else {
                    parts.push(&text[..idx])
                }
                search_input = !search_input;
                text = &text[(idx + search.len())..];
            }
            None => {
                if !text.is_empty() {
                    if system.is_none() {
                        system = Some(text)
                    } else {
                        parts.push(text)
                    }
                }
                break;
            }
        }
    }
    let parts_len = parts.len();
    if parts_len > 0 && parts_len % 2 == 0 {
        let cases: Vec<(&str, &str)> = parts
            .iter()
            .step_by(2)
            .zip(parts.iter().skip(1).step_by(2))
            .map(|(i, o)| (i.trim(), o.trim()))
            .collect();
        let system = system.map(|v| v.trim()).unwrap_or_default();
        return (system, cases);
    }

    (prompt, vec![])
}

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

    #[test]
    fn role_new_parses_prompt() {
        let role = Role::new("test", "You are a helpful assistant");
        assert_eq!(role.name(), "test");
        assert_eq!(role.prompt(), "You are a helpful assistant");
    }

    #[test]
    fn role_new_parses_metadata() {
        let content =
            "---\nmodel: openai:gpt-4\ntemperature: 0.7\ntop_p: 0.9\n---\nYou are helpful";
        let role = Role::new("test", content);
        assert_eq!(role.model_id(), Some("openai:gpt-4"));
        assert_eq!(role.temperature(), Some(0.7));
        assert_eq!(role.top_p(), Some(0.9));
        assert_eq!(role.prompt(), "You are helpful");
    }

    #[test]
    fn role_new_parses_enabled_tools() {
        let content = "---\nenabled_tools: tool1,tool2\n---\nPrompt";
        let role = Role::new("test", content);
        assert_eq!(
            role.enabled_tools(),
            Some(vec!["tool1".to_string(), "tool2".to_string()])
        );
    }

    #[test]
    fn role_new_parses_enabled_mcp_servers() {
        let content = "---\nenabled_mcp_servers: github,jira\n---\nPrompt";
        let role = Role::new("test", content);
        assert_eq!(
            role.enabled_mcp_servers(),
            Some(vec!["github".to_string(), "jira".to_string()])
        );
    }

    #[test]
    fn role_new_no_metadata_has_none_fields() {
        let role = Role::new("test", "Just a prompt");
        assert_eq!(role.model_id(), None);
        assert_eq!(role.temperature(), None);
        assert_eq!(role.top_p(), None);
        assert_eq!(role.enabled_tools(), None);
        assert_eq!(role.enabled_mcp_servers(), None);
    }

    #[test]
    fn role_builtin_shell_loads() {
        let role = Role::builtin("shell").unwrap();
        assert_eq!(role.name(), "shell");
        assert!(!role.prompt().is_empty());
    }

    #[test]
    fn role_builtin_code_loads() {
        let role = Role::builtin("code").unwrap();
        assert_eq!(role.name(), "code");
        assert!(!role.prompt().is_empty());
    }

    #[test]
    fn role_builtin_nonexistent_errors() {
        let result = Role::builtin("nonexistent_role_xyz");
        assert!(result.is_err());
    }

    #[test]
    fn role_default_has_empty_fields() {
        let role = Role::default();
        assert_eq!(role.name(), "");
        assert_eq!(role.prompt(), "");
        assert_eq!(role.model_id(), None);
    }

    #[test]
    fn role_set_model_updates_model() {
        let mut role = Role::new("test", "prompt");
        let model = Model::default();
        role.set_model(model.clone());
        assert_eq!(role.model().id(), model.id());
    }

    #[test]
    fn role_set_temperature_works() {
        let mut role = Role::new("test", "prompt");
        role.set_temperature(Some(0.5));
        assert_eq!(role.temperature(), Some(0.5));
    }

    #[test]
    fn role_export_includes_metadata() {
        let content = "---\ntemperature: 0.8\n---\nMy prompt";
        let role = Role::new("test", content);
        let exported = role.export();
        assert!(exported.contains("temperature"));
        assert!(exported.contains("My prompt"));
    }

    #[test]
    fn test_parse_structure_prompt1() {
        let prompt = r#"
System message
### INPUT:
Input 1
### OUTPUT:
Output 1
"#;
        assert_eq!(
            parse_structure_prompt(prompt),
            ("System message", vec![("Input 1", "Output 1")])
        );
    }

    #[test]
    fn test_parse_structure_prompt2() {
        let prompt = r#"
### INPUT:
Input 1
### OUTPUT:
Output 1
"#;
        assert_eq!(
            parse_structure_prompt(prompt),
            ("", vec![("Input 1", "Output 1")])
        );
    }

    #[test]
    fn test_parse_structure_prompt3() {
        let prompt = r#"
System message
### INPUT:
Input 1
"#;
        assert_eq!(parse_structure_prompt(prompt), (prompt, vec![]));
    }
}
