use anyhow::{Context, Result};
use serde::{Deserialize, Serialize};
use tracing::{debug, instrument, warn};

/// A search plan produced by an LLM that describes which grep commands to run
/// and which files to read in order to satisfy a user query.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SearchPlan {
    /// Grep patterns/commands to execute against the workspace.
    pub grep_commands: Vec<String>,
    /// Specific files to read (absolute or workspace-relative paths).
    pub files_to_read: Vec<String>,
}

impl SearchPlan {
    pub fn is_empty(&self) -> bool {
        self.grep_commands.is_empty() && self.files_to_read.is_empty()
    }
}

/// Parses structured search plans out of LLM responses.
///
/// The parser expects the LLM to return a JSON object (possibly wrapped in a
/// markdown code fence) that conforms to the [`SearchPlan`] schema:
///
/// ```json
/// {
///   "grep_commands": ["pattern1", "pattern2"],
///   "files_to_read": ["src/main.rs"]
/// }
/// ```
pub struct PlanParser;

impl PlanParser {
    pub fn new() -> Self {
        Self
    }

    /// Attempt to parse a [`SearchPlan`] from the raw text returned by an LLM.
    ///
    /// The method is intentionally lenient:
    /// - It strips markdown code fences (```json ... ```) if present.
    /// - It searches for the first `{` and last `}` to isolate the JSON blob.
    /// - Missing fields default to empty vectors.
    #[instrument(skip(self, llm_response), fields(response_len = llm_response.len()))]
    pub fn parse(&self, llm_response: &str) -> Result<SearchPlan> {
        let json_str = Self::extract_json(llm_response)
            .context("failed to locate JSON object in LLM response")?;

        debug!(json_len = json_str.len(), "extracted JSON from LLM response");

        // Try strict parse first.
        if let Ok(plan) = serde_json::from_str::<SearchPlan>(&json_str) {
            debug!(
                grep_count = plan.grep_commands.len(),
                file_count = plan.files_to_read.len(),
                "parsed search plan (strict)"
            );
            return Ok(plan);
        }

        // Fallback: parse as a generic Value so we can handle missing fields.
        let value: serde_json::Value = serde_json::from_str(&json_str)
            .context("LLM response does not contain valid JSON")?;

        let grep_commands = Self::extract_string_array(&value, "grep_commands");
        let files_to_read = Self::extract_string_array(&value, "files_to_read");

        if grep_commands.is_empty() && files_to_read.is_empty() {
            warn!("parsed plan is empty -- LLM may not have followed the schema");
        }

        let plan = SearchPlan {
            grep_commands,
            files_to_read,
        };

        debug!(
            grep_count = plan.grep_commands.len(),
            file_count = plan.files_to_read.len(),
            "parsed search plan (lenient)"
        );

        Ok(plan)
    }

    /// Parse additional grep commands from an LLM response that is expected to
    /// contain a JSON array of strings (or an object with a `grep_commands`
    /// field).
    #[instrument(skip(self, llm_response))]
    pub fn parse_additional_greps(&self, llm_response: &str) -> Result<Vec<String>> {
        let json_str = Self::extract_json(llm_response)
            .context("failed to locate JSON in LLM response for additional greps")?;

        // Try as a plain array of strings first.
        if let Ok(commands) = serde_json::from_str::<Vec<String>>(&json_str) {
            return Ok(commands);
        }

        // Try as an object with a grep_commands field.
        if let Ok(value) = serde_json::from_str::<serde_json::Value>(&json_str) {
            let commands = Self::extract_string_array(&value, "grep_commands");
            if !commands.is_empty() {
                return Ok(commands);
            }
            // Also try "patterns" or "commands" as alternative field names.
            for key in &["patterns", "commands", "searches"] {
                let commands = Self::extract_string_array(&value, key);
                if !commands.is_empty() {
                    return Ok(commands);
                }
            }
        }

        anyhow::bail!("could not parse grep commands from LLM response")
    }

    /// Parse a sufficiency judgement from the LLM. Returns `true` when the LLM
    /// indicates that the existing results are sufficient.
    #[instrument(skip(self, llm_response))]
    pub fn parse_sufficiency(&self, llm_response: &str) -> bool {
        let lower = llm_response.to_lowercase();

        // Look for explicit JSON boolean first.
        if let Ok(json_str) = Self::extract_json(llm_response) {
            if let Ok(value) = serde_json::from_str::<serde_json::Value>(&json_str) {
                if let Some(sufficient) = value.get("sufficient").and_then(|v| v.as_bool()) {
                    return sufficient;
                }
                if let Some(sufficient) = value.get("is_sufficient").and_then(|v| v.as_bool()) {
                    return sufficient;
                }
            }
        }

        // Heuristic fallback: look for keywords.
        if lower.contains("sufficient") && !lower.contains("not sufficient") && !lower.contains("insufficient") {
            return true;
        }
        if lower.contains("yes") && !lower.contains("no") {
            return true;
        }

        false
    }

    // ---- helpers ----

    /// Extract the first JSON blob (object or array) from a potentially noisy
    /// LLM response. Handles markdown code fences.
    fn extract_json(text: &str) -> Result<String> {
        // Strip markdown code fences.
        let stripped = Self::strip_code_fence(text);

        // Find the first `{` or `[` and the matching last `}` or `]`.
        let start_brace = stripped.find('{');
        let start_bracket = stripped.find('[');

        let (start, end_char) = match (start_brace, start_bracket) {
            (Some(b), Some(k)) => {
                if b <= k {
                    (b, '}')
                } else {
                    (k, ']')
                }
            }
            (Some(b), None) => (b, '}'),
            (None, Some(k)) => (k, ']'),
            (None, None) => anyhow::bail!("no JSON object or array found in text"),
        };

        let end = stripped
            .rfind(end_char)
            .context("no matching closing delimiter found")?;

        if end <= start {
            anyhow::bail!("malformed JSON: closing delimiter before opening");
        }

        Ok(stripped[start..=end].to_string())
    }

    /// Strip ```json ... ``` style code fences.
    fn strip_code_fence(text: &str) -> String {
        let mut result = text.to_string();

        // Remove opening fence: ```json or ``` at the start
        if let Some(start) = result.find("```") {
            let fence_end = result[start + 3..]
                .find('\n')
                .map(|i| start + 3 + i + 1)
                .unwrap_or(start + 3);
            result = format!("{}{}", &result[..start], &result[fence_end..]);
        }

        // Remove closing fence: ``` at the end
        if let Some(end) = result.rfind("```") {
            result = result[..end].to_string();
        }

        result
    }

    /// Extract a `Vec<String>` from a JSON value's field, returning an empty
    /// vec when the field is missing or malformed.
    fn extract_string_array(value: &serde_json::Value, field: &str) -> Vec<String> {
        value
            .get(field)
            .and_then(|v| v.as_array())
            .map(|arr| {
                arr.iter()
                    .filter_map(|item| item.as_str().map(String::from))
                    .collect()
            })
            .unwrap_or_default()
    }
}

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

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

    #[test]
    fn test_parse_clean_json() {
        let parser = PlanParser::new();
        let input = r#"{"grep_commands": ["fn main", "struct Config"], "files_to_read": ["src/main.rs"]}"#;
        let plan = parser.parse(input).unwrap();
        assert_eq!(plan.grep_commands.len(), 2);
        assert_eq!(plan.files_to_read.len(), 1);
    }

    #[test]
    fn test_parse_with_code_fence() {
        let parser = PlanParser::new();
        let input = r#"Here is the plan:
```json
{
  "grep_commands": ["TODO"],
  "files_to_read": []
}
```
"#;
        let plan = parser.parse(input).unwrap();
        assert_eq!(plan.grep_commands, vec!["TODO"]);
        assert!(plan.files_to_read.is_empty());
    }

    #[test]
    fn test_parse_missing_fields() {
        let parser = PlanParser::new();
        let input = r#"{"grep_commands": ["hello"]}"#;
        let plan = parser.parse(input).unwrap();
        assert_eq!(plan.grep_commands.len(), 1);
        assert!(plan.files_to_read.is_empty());
    }

    #[test]
    fn test_parse_sufficiency_true() {
        let parser = PlanParser::new();
        assert!(parser.parse_sufficiency(r#"{"sufficient": true}"#));
        assert!(parser.parse_sufficiency("The results are sufficient."));
        assert!(parser.parse_sufficiency("Yes, these results cover the query."));
    }

    #[test]
    fn test_parse_sufficiency_false() {
        let parser = PlanParser::new();
        assert!(!parser.parse_sufficiency(r#"{"sufficient": false}"#));
        assert!(!parser.parse_sufficiency("The results are not sufficient."));
    }

    #[test]
    fn test_parse_additional_greps_array() {
        let parser = PlanParser::new();
        let input = r#"["fn search", "impl Trait"]"#;
        let greps = parser.parse_additional_greps(input).unwrap();
        assert_eq!(greps, vec!["fn search", "impl Trait"]);
    }

    #[test]
    fn test_parse_additional_greps_object() {
        let parser = PlanParser::new();
        let input = r#"{"grep_commands": ["fn search"]}"#;
        let greps = parser.parse_additional_greps(input).unwrap();
        assert_eq!(greps, vec!["fn search"]);
    }

    #[test]
    fn test_no_json_returns_error() {
        let parser = PlanParser::new();
        let result = parser.parse("just some plain text with no json");
        assert!(result.is_err());
    }
}
