use serde::de::DeserializeOwned;
use serde_json::Value;

use crate::bracket::{extract_all_json_entities, extract_json_entity};
use crate::error::{ExtractError, Extracted};
use crate::fence::{strip_fences, strip_think_blocks};
use crate::schema::{Extract, RepairSchema, SchemaKind};
use fuzzy_parser::{repair_fields_with_list, Algorithm, Correction, FuzzyOptions};

/// Default maximum input length (1 MB).
const DEFAULT_MAX_INPUT_LEN: usize = 1024 * 1024;

/// Repair mode configuration.
#[derive(Debug, Clone)]
pub enum RepairMode {
    /// No repair. Layer 1 output is deserialized directly.
    Strict,
    /// Fuzzy matching repair with configurable threshold.
    Fuzzy {
        min_similarity: f64,
        algorithm: Algorithm,
    },
}

impl Default for RepairMode {
    fn default() -> Self {
        RepairMode::Fuzzy {
            min_similarity: 0.7,
            algorithm: Algorithm::JaroWinkler,
        }
    }
}

/// Configurable extraction pipeline.
pub struct Extractor {
    max_input_len: usize,
    repair_mode: RepairMode,
}

/// Builder for [`Extractor`].
pub struct ExtractorBuilder {
    max_input_len: usize,
    repair_mode: RepairMode,
}

impl ExtractorBuilder {
    /// Set maximum input length.
    pub fn max_input_len(mut self, len: usize) -> Self {
        self.max_input_len = len;
        self
    }

    /// Set repair mode.
    pub fn repair(mut self, mode: RepairMode) -> Self {
        self.repair_mode = mode;
        self
    }

    /// Build the extractor.
    pub fn build(self) -> Extractor {
        Extractor {
            max_input_len: self.max_input_len,
            repair_mode: self.repair_mode,
        }
    }
}

impl Extractor {
    /// Create a builder with default settings.
    pub fn builder() -> ExtractorBuilder {
        ExtractorBuilder {
            max_input_len: DEFAULT_MAX_INPUT_LEN,
            repair_mode: RepairMode::default(),
        }
    }

    /// Layer 1: Extract JSON as `serde_json::Value`.
    pub fn extract_json(&self, text: &str) -> Result<Value, ExtractError> {
        extract_json_with_limit(text, self.max_input_len)
    }

    /// Layer 2: Extract and deserialize to `T` with repair.
    pub fn extract<T: Extract>(&self, text: &str) -> Result<T, ExtractError> {
        let value = self.extract_json(text)?;
        deserialize_with_repair::<T>(value, &self.repair_mode)
    }

    /// Layer 1 + deserialize without repair.
    pub fn extract_strict<T: DeserializeOwned>(&self, text: &str) -> Result<T, ExtractError> {
        let value = self.extract_json(text)?;
        deserialize_strict(value)
    }

    /// Layer 1: Extract all JSON entities as `Vec<Value>`.
    pub fn extract_json_all(&self, text: &str) -> Result<Vec<Value>, ExtractError> {
        extract_json_all_with_limit(text, self.max_input_len)
    }

    /// Layer 2: Extract all JSON entities and deserialize each to `T`.
    pub fn extract_all<T: Extract>(&self, text: &str) -> Result<Vec<T>, ExtractError> {
        let values = self.extract_json_all(text)?;
        values
            .into_iter()
            .map(|v| deserialize_with_repair::<T>(v, &self.repair_mode))
            .collect()
    }

    /// Layer 2 with corrections: returns `Extracted<T>` including repair metadata.
    pub fn extract_verbose<T: Extract>(&self, text: &str) -> Result<Extracted<T>, ExtractError> {
        let value = self.extract_json(text)?;
        deserialize_with_repair_verbose::<T>(value, &self.repair_mode)
    }
}

// ---------------------------------------------------------------------------
// Free functions (convenience API using defaults)
// ---------------------------------------------------------------------------

/// Layer 1: Extract JSON `Value` from LLM response text.
///
/// Pipeline: strip think blocks -> strip fences -> extract entity -> jsonrepair -> parse.
pub fn extract_json(text: &str) -> Result<Value, ExtractError> {
    extract_json_with_limit(text, DEFAULT_MAX_INPUT_LEN)
}

/// Layer 2: Extract and deserialize to `T` with fuzzy repair.
///
/// Requires `T: Extract` (via `#[derive(Extract)]`).
pub fn extract<T: Extract>(text: &str) -> Result<T, ExtractError> {
    let value = extract_json(text)?;
    deserialize_with_repair::<T>(value, &RepairMode::default())
}

/// Layer 1 + deserialize without repair. Requires `T: DeserializeOwned`.
pub fn extract_strict<T: DeserializeOwned>(text: &str) -> Result<T, ExtractError> {
    let value = extract_json(text)?;
    deserialize_strict(value)
}

/// Layer 1: Extract all JSON entities as `Vec<Value>`.
///
/// Returns all balanced JSON objects/arrays found in the text.
/// Each is independently repaired via jsonrepair.
pub fn extract_json_all(text: &str) -> Result<Vec<Value>, ExtractError> {
    extract_json_all_with_limit(text, DEFAULT_MAX_INPUT_LEN)
}

/// Layer 2: Extract all JSON entities and deserialize each to `T`.
pub fn extract_all<T: Extract>(text: &str) -> Result<Vec<T>, ExtractError> {
    let values = extract_json_all(text)?;
    values
        .into_iter()
        .map(|v| deserialize_with_repair::<T>(v, &RepairMode::default()))
        .collect()
}

/// Layer 2 with corrections: returns `Extracted<T>` including repair metadata.
///
/// Use this when you need to know what was repaired (for logging, monitoring, etc.).
pub fn extract_verbose<T: Extract>(text: &str) -> Result<Extracted<T>, ExtractError> {
    let value = extract_json(text)?;
    deserialize_with_repair_verbose::<T>(value, &RepairMode::default())
}

// ---------------------------------------------------------------------------
// Internal: Layer 1
// ---------------------------------------------------------------------------

/// jsonrepair options tuned for LLM output.
fn jsonrepair_opts() -> jsonrepair::Options {
    jsonrepair::Options {
        fenced_code_blocks: true,
        allow_python_keywords: true,
        tolerate_hash_comments: true,
        repair_undefined: true,
        normalize_js_nonfinite: true,
        ensure_ascii: false,
        ..jsonrepair::Options::default()
    }
}

fn extract_json_with_limit(text: &str, max_len: usize) -> Result<Value, ExtractError> {
    if text.len() > max_len {
        return Err(ExtractError::InputTooLarge {
            len: text.len(),
            max: max_len,
        });
    }

    // Step 0: Strip <think> blocks (CoT reasoning)
    let without_think = strip_think_blocks(text);

    // Step 1: Strip code fences + surrounding prose
    let stripped = strip_fences(&without_think);

    // Step 2: Try to parse directly first (fast path)
    if let Ok(value) = serde_json::from_str::<Value>(stripped.trim()) {
        return Ok(value);
    }

    // Step 3: Try jsonrepair on the stripped text (handles fence, single quotes,
    // Python keywords, unquoted keys, trailing commas, etc.)
    let opts = jsonrepair_opts();
    if let Ok(repaired) = jsonrepair::repair_to_string(stripped, &opts) {
        if let Ok(value) = serde_json::from_str::<Value>(&repaired) {
            return Ok(value);
        }
    }

    // Step 4: Extract JSON entity via bracket matching, then repair
    if let Some(entity) = extract_json_entity(stripped) {
        if let Ok(repaired) = jsonrepair::repair_to_string(entity, &opts) {
            match serde_json::from_str::<Value>(&repaired) {
                Ok(value) => return Ok(value),
                Err(e) => {
                    return Err(ExtractError::Parse {
                        extracted: repaired,
                        source: e,
                    });
                }
            }
        }
    }

    // Step 5: Last resort — try jsonrepair on original text
    if let Ok(repaired) = jsonrepair::repair_to_string(text, &opts) {
        if let Ok(value) = serde_json::from_str::<Value>(&repaired) {
            return Ok(value);
        }
    }

    Err(ExtractError::NotFound)
}

fn extract_json_all_with_limit(text: &str, max_len: usize) -> Result<Vec<Value>, ExtractError> {
    if text.len() > max_len {
        return Err(ExtractError::InputTooLarge {
            len: text.len(),
            max: max_len,
        });
    }

    let without_think = strip_think_blocks(text);
    let stripped = strip_fences(&without_think);
    let opts = jsonrepair_opts();

    let entities = extract_all_json_entities(stripped);
    if entities.is_empty() {
        return Err(ExtractError::NotFound);
    }

    let mut results = Vec::with_capacity(entities.len());
    for entity in entities {
        // Try direct parse first
        if let Ok(value) = serde_json::from_str::<Value>(entity.trim()) {
            results.push(value);
            continue;
        }
        // Try jsonrepair
        if let Ok(repaired) = jsonrepair::repair_to_string(entity, &opts) {
            if let Ok(value) = serde_json::from_str::<Value>(&repaired) {
                results.push(value);
                continue;
            }
        }
        // Skip unparseable entities
    }

    if results.is_empty() {
        return Err(ExtractError::NotFound);
    }

    Ok(results)
}

// ---------------------------------------------------------------------------
// Internal: Deserialization
// ---------------------------------------------------------------------------

fn deserialize_strict<T: DeserializeOwned>(value: Value) -> Result<T, ExtractError> {
    let json_str = value.to_string();
    serde_json::from_value::<T>(value).map_err(|e| ExtractError::Parse {
        extracted: json_str,
        source: e,
    })
}

fn deserialize_with_repair<T: Extract>(value: Value, mode: &RepairMode) -> Result<T, ExtractError> {
    let json_str = value.to_string();

    // First, try without repair (fast path)
    if let Ok(result) = serde_json::from_value::<T>(value.clone()) {
        return Ok(result);
    }

    // If strict mode, fail immediately
    let (min_similarity, algorithm) = match mode {
        RepairMode::Strict => {
            return serde_json::from_value::<T>(value).map_err(|e| ExtractError::Parse {
                extracted: json_str,
                source: e,
            });
        }
        RepairMode::Fuzzy {
            min_similarity,
            algorithm,
        } => (*min_similarity, *algorithm),
    };

    // Apply repair based on schema
    let schema = T::repair_schema();
    let options = FuzzyOptions {
        min_similarity,
        algorithm,
    };

    let (repaired_value, corrections) = apply_repair(value, &schema, &options);
    let repaired_str = repaired_value.to_string();

    serde_json::from_value::<T>(repaired_value).map_err(|e| ExtractError::RepairFailed {
        repaired: repaired_str,
        corrections,
        source: e,
    })
}

fn deserialize_with_repair_verbose<T: Extract>(
    value: Value,
    mode: &RepairMode,
) -> Result<Extracted<T>, ExtractError> {
    let json_str = value.to_string();

    // Fast path: try without repair
    if let Ok(result) = serde_json::from_value::<T>(value.clone()) {
        return Ok(Extracted {
            value: result,
            corrections: Vec::new(),
        });
    }

    let (min_similarity, algorithm) = match mode {
        RepairMode::Strict => {
            return serde_json::from_value::<T>(value)
                .map(|v| Extracted {
                    value: v,
                    corrections: Vec::new(),
                })
                .map_err(|e| ExtractError::Parse {
                    extracted: json_str,
                    source: e,
                });
        }
        RepairMode::Fuzzy {
            min_similarity,
            algorithm,
        } => (*min_similarity, *algorithm),
    };

    let schema = T::repair_schema();
    let options = FuzzyOptions {
        min_similarity,
        algorithm,
    };

    let (repaired_value, corrections) = apply_repair(value, &schema, &options);
    let repaired_str = repaired_value.to_string();

    serde_json::from_value::<T>(repaired_value)
        .map(|v| Extracted {
            value: v,
            corrections,
        })
        .map_err(|e| ExtractError::RepairFailed {
            repaired: repaired_str,
            corrections: Vec::new(),
            source: e,
        })
}

fn apply_repair(
    mut value: Value,
    schema: &RepairSchema,
    options: &FuzzyOptions,
) -> (Value, Vec<Correction>) {
    let mut corrections = Vec::new();

    match &schema.kind {
        SchemaKind::Struct(struct_schema) => {
            if let Value::Object(ref mut map) = value {
                let field_names: Vec<&str> = struct_schema.fields.iter().map(|f| f.name).collect();
                let new_corrections = repair_fields_with_list(map, &field_names, "", options);
                corrections.extend(new_corrections);

                // Recursive repair for nested fields with child_schema
                for field_def in struct_schema.fields.iter() {
                    if let Some(child_schema_fn) = field_def.child_schema {
                        if let Some(child_value) = map.remove(field_def.name) {
                            let child_schema = child_schema_fn();
                            let (repaired_child, child_corrections) =
                                apply_repair(child_value, &child_schema, options);
                            corrections.extend(child_corrections);
                            map.insert(field_def.name.to_string(), repaired_child);
                        }
                    }
                }
            }
        }
        SchemaKind::TaggedEnum(enum_def) => {
            if let Value::Object(ref mut map) = value {
                // Repair tag value
                if let Some(tag_value) = map.get(enum_def.tag_field) {
                    if let Some(tag_str) = tag_value.as_str() {
                        if !enum_def.variants.contains(&tag_str) {
                            // Fuzzy match against valid variants
                            let candidates = enum_def.variants.iter().copied();
                            let best = fuzzy_parser::distance::find_closest(
                                tag_str,
                                candidates,
                                options.min_similarity,
                                options.algorithm,
                            );
                            if let Some(m) = best {
                                corrections.push(Correction {
                                    original: tag_str.to_string(),
                                    corrected: m.candidate.clone(),
                                    similarity: m.similarity,
                                    field_path: enum_def.tag_field.to_string(),
                                });
                                map.insert(
                                    enum_def.tag_field.to_string(),
                                    Value::String(m.candidate),
                                );
                            }
                        }
                    }
                }
            }
        }
    }

    (value, corrections)
}

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

    #[test]
    fn layer1_bare_json() {
        let v = extract_json(r#"{"a": 1}"#).unwrap();
        assert_eq!(v["a"], 1);
    }

    #[test]
    fn layer1_fenced() {
        let v = extract_json("```json\n{\"a\": 1}\n```").unwrap();
        assert_eq!(v["a"], 1);
    }

    #[test]
    fn layer1_with_prose() {
        let v = extract_json("Here's the result: {\"a\": 1} Hope this helps!").unwrap();
        assert_eq!(v["a"], 1);
    }

    #[test]
    fn layer1_trailing_comma() {
        let v = extract_json(r#"{"a": 1,}"#).unwrap();
        assert_eq!(v["a"], 1);
    }

    #[test]
    fn layer1_unclosed_bracket() {
        let v = extract_json(r#"{"a": 1"#).unwrap();
        assert_eq!(v["a"], 1);
    }

    #[test]
    fn layer1_not_found() {
        // Note: jsonrepair may parse plain text as a JSON string.
        // NotFound is only returned when no JSON value can be produced at all.
        let err = extract_json("").unwrap_err();
        assert!(matches!(
            err,
            ExtractError::NotFound | ExtractError::Parse { .. }
        ));
    }

    #[test]
    fn layer1_input_too_large() {
        let large = "x".repeat(DEFAULT_MAX_INPUT_LEN + 1);
        let err = extract_json(&large).unwrap_err();
        assert!(matches!(err, ExtractError::InputTooLarge { .. }));
    }

    #[test]
    fn strict_deserialize() {
        #[derive(serde::Deserialize, Debug, PartialEq)]
        struct S {
            a: i32,
        }
        let result: S = extract_strict(r#"{"a": 1}"#).unwrap();
        assert_eq!(result, S { a: 1 });
    }

    // jsonrepair-powered tests
    #[test]
    fn layer1_single_quotes() {
        let v = extract_json("{'a': 1}").unwrap();
        assert_eq!(v["a"], 1);
    }

    #[test]
    fn layer1_unquoted_keys() {
        let v = extract_json("{a: 1, b: 2}").unwrap();
        assert_eq!(v["a"], 1);
        assert_eq!(v["b"], 2);
    }

    #[test]
    fn layer1_python_keywords() {
        let v = extract_json(r#"{"active": True, "deleted": False, "value": None}"#).unwrap();
        assert_eq!(v["active"], true);
        assert_eq!(v["deleted"], false);
        assert!(v["value"].is_null());
    }

    #[test]
    fn layer1_comments() {
        let v = extract_json(
            r#"{
            // This is a comment
            "a": 1, /* inline comment */
            "b": 2
        }"#,
        )
        .unwrap();
        assert_eq!(v["a"], 1);
        assert_eq!(v["b"], 2);
    }
}
