use std::collections::HashSet;

use chrono::Utc;

use crate::catalog::{
    benchmark_fixture_uri, eval_case_key, eval_case_ordering, eval_golden_uri,
    normalize_eval_case_source, parse_benchmark_fixture_document,
};
use crate::error::{AxiomError, Result};
use crate::models::{
    BenchmarkRunOptions, EvalGoldenDocument, EvalQueryCase, FindResult, RetrievalTrace,
    SearchOptions, TraceIndexEntry,
};
use crate::uri::{AxiomUri, Scope};

use super::AxiomSync;

const TRACE_CASE_EXPANSION_FACTOR: usize = 4;

impl AxiomSync {
    pub(crate) fn persist_trace_result(&self, result: &mut FindResult) -> Result<()> {
        let Some(trace) = result.trace.as_ref() else {
            result.trace_uri = None;
            return Ok(());
        };

        let trace_uri = self.persist_trace(trace)?;
        result.trace_uri = Some(trace_uri);
        Ok(())
    }

    pub(super) fn persist_trace(&self, trace: &RetrievalTrace) -> Result<String> {
        let trace_uri = AxiomUri::root(Scope::Queue)
            .join("traces")?
            .join(&format!("{}.json", trace.trace_id))?;
        let serialized = serde_json::to_string_pretty(trace)?;
        self.fs.write(&trace_uri, &serialized, true)?;

        self.state.upsert_trace_index(&TraceIndexEntry {
            trace_id: trace.trace_id.clone(),
            uri: trace_uri.to_string(),
            request_type: trace.request_type.clone(),
            query: trace.query.clone(),
            target_uri: trace.target_uri.clone(),
            created_at: Utc::now().to_rfc3339(),
        })?;

        Ok(trace_uri.to_string())
    }

    pub(crate) fn collect_trace_eval_cases(
        &self,
        trace_limit: usize,
    ) -> Result<(Vec<EvalQueryCase>, usize)> {
        self.collect_trace_eval_cases_with_expectations(trace_limit, true)
    }

    pub(crate) fn collect_trace_eval_cases_with_expectations(
        &self,
        trace_limit: usize,
        include_expectations: bool,
    ) -> Result<(Vec<EvalQueryCase>, usize)> {
        let trace_entries = self.list_traces(trace_limit)?;
        let traces_scanned = trace_entries.len();
        let mut cases = Vec::<EvalQueryCase>::new();

        for entry in trace_entries {
            let Some(trace) = self.get_trace(&entry.trace_id)? else {
                continue;
            };
            cases.push(EvalQueryCase {
                source_trace_id: trace.trace_id,
                query: trace.query,
                target_uri: trace.target_uri,
                expected_top_uri: if include_expectations {
                    trace.final_topk.first().map(|x| x.uri.clone())
                } else {
                    None
                },
                source: if include_expectations {
                    "trace".to_string()
                } else {
                    "trace-unlabeled".to_string()
                },
            });
        }

        Ok((cases, traces_scanned))
    }

    pub(crate) fn collect_benchmark_query_cases(
        &self,
        options: &BenchmarkRunOptions,
        query_limit: usize,
    ) -> Result<Vec<EvalQueryCase>> {
        let query_limit = query_limit.max(1);
        if let Some(fixture_name) = options.fixture_name.as_deref() {
            let fixture_uri = benchmark_fixture_uri(fixture_name)?;
            let raw = self.fs.read(&fixture_uri)?;
            let mut doc = parse_benchmark_fixture_document(&raw)?;
            for case in &mut doc.cases {
                if case.source.trim().is_empty() {
                    case.source = "fixture".to_string();
                }
            }
            if options.include_stress {
                let stress_cases = build_stress_cases(&doc.cases, query_limit);
                doc.cases.extend(stress_cases);
            }
            doc.cases.sort_by(eval_case_ordering);
            doc.cases.truncate(query_limit);
            return Ok(doc.cases);
        }

        let mut seen = HashSet::<(String, Option<String>)>::new();
        let mut query_cases = Vec::<EvalQueryCase>::new();
        let mut golden_seed_cases = Vec::<EvalQueryCase>::new();
        if options.include_golden {
            for mut case in self.list_eval_golden_queries()? {
                if query_cases.len() >= query_limit {
                    break;
                }
                normalize_eval_case_source(&mut case, "golden");
                if !seen.insert(eval_case_key(&case)) {
                    continue;
                }
                golden_seed_cases.push(case.clone());
                query_cases.push(case);
            }
        }
        if options.include_stress && query_cases.len() < query_limit {
            let stress_cases = build_stress_cases(
                &golden_seed_cases,
                query_limit.saturating_sub(query_cases.len()),
            );
            for mut case in stress_cases {
                if query_cases.len() >= query_limit {
                    break;
                }
                normalize_eval_case_source(&mut case, "stress");
                if !seen.insert(eval_case_key(&case)) {
                    continue;
                }
                query_cases.push(case);
            }
        }
        if options.include_trace && query_cases.len() < query_limit {
            let trace_limit = query_limit
                .saturating_mul(TRACE_CASE_EXPANSION_FACTOR)
                .max(query_limit);
            let (trace_cases, _) = self.collect_trace_eval_cases_with_expectations(
                trace_limit,
                options.trace_expectations,
            )?;
            for mut case in trace_cases {
                if query_cases.len() >= query_limit {
                    break;
                }
                let fallback = if options.trace_expectations {
                    "trace"
                } else {
                    "trace-unlabeled"
                };
                normalize_eval_case_source(&mut case, fallback);
                if !seen.insert(eval_case_key(&case)) {
                    continue;
                }
                query_cases.push(case);
            }
        }
        Ok(query_cases)
    }

    pub(crate) fn eval_top_result_uri(
        &self,
        query: &str,
        target_uri: Option<&str>,
        search_limit: usize,
    ) -> Result<Option<String>> {
        let uris = self.eval_result_uris(query, target_uri, search_limit, "eval")?;
        Ok(uris.first().cloned())
    }

    pub(crate) fn eval_result_uris(
        &self,
        query: &str,
        target_uri: Option<&str>,
        search_limit: usize,
        request_type: &str,
    ) -> Result<Vec<String>> {
        let target = target_uri.map(AxiomUri::parse).transpose()?;
        let options = SearchOptions {
            query: query.to_string(),
            target_uri: target,
            session: None,
            session_hints: Vec::new(),
            budget: None,
            limit: search_limit.max(1),
            score_threshold: None,
            min_match_tokens: None,
            filter: None,
            request_type: request_type.to_string(),
        };
        let result = {
            let index = self
                .index
                .read()
                .map_err(|_| AxiomError::lock_poisoned("index"))?;
            self.drr.run(&index, &options)
        };
        Ok(result.query_results.into_iter().map(|x| x.uri).collect())
    }

    pub(crate) fn persist_eval_golden_queries(&self, cases: &[EvalQueryCase]) -> Result<String> {
        let golden_uri = eval_golden_uri()?;
        let document = EvalGoldenDocument {
            version: 1,
            updated_at: Utc::now().to_rfc3339(),
            cases: cases.to_vec(),
        };
        self.fs
            .write(&golden_uri, &serde_json::to_string_pretty(&document)?, true)?;
        Ok(golden_uri.to_string())
    }
}

fn build_stress_cases(seed_cases: &[EvalQueryCase], limit: usize) -> Vec<EvalQueryCase> {
    if seed_cases.is_empty() || limit == 0 {
        return Vec::new();
    }

    let mut out = Vec::<EvalQueryCase>::new();
    let mut seen = HashSet::<(String, Option<String>)>::new();
    for seed in seed_cases {
        for (variant_name, query) in build_stress_query_variants(&seed.query) {
            if out.len() >= limit {
                return out;
            }
            let case = EvalQueryCase {
                source_trace_id: seed.source_trace_id.clone(),
                query,
                target_uri: seed.target_uri.clone(),
                expected_top_uri: seed.expected_top_uri.clone(),
                source: format!("stress:{variant_name}"),
            };
            if !seen.insert(eval_case_key(&case)) {
                continue;
            }
            out.push(case);
        }
    }
    out
}

fn build_stress_query_variants(query: &str) -> Vec<(&'static str, String)> {
    let trimmed = query.trim();
    if trimmed.is_empty() {
        return Vec::new();
    }

    let mut variants = Vec::<(&'static str, String)>::new();
    if let Some(variant) = stress_typo_variant(trimmed) {
        variants.push(("typo", variant));
    }
    if let Some(variant) = stress_alias_variant(trimmed) {
        variants.push(("alias", variant));
    }
    if let Some(variant) = stress_paraphrase_variant(trimmed) {
        variants.push(("paraphrase", variant));
    }

    let mut unique = HashSet::<String>::new();
    let mut out = Vec::<(&'static str, String)>::new();
    for (name, variant) in variants {
        if variant.trim().is_empty() || variant.eq_ignore_ascii_case(trimmed) {
            continue;
        }
        let key = variant.trim().to_lowercase();
        if unique.insert(key) {
            out.push((name, variant));
        }
    }
    out
}

fn stress_typo_variant(query: &str) -> Option<String> {
    let mut tokens = query
        .split_whitespace()
        .map(ToString::to_string)
        .collect::<Vec<_>>();
    let idx = tokens
        .iter()
        .enumerate()
        .max_by_key(|(_, token)| token.chars().count())
        .map(|(idx, _)| idx)?;

    let chars = tokens[idx].chars().collect::<Vec<_>>();
    if chars.len() < 5 {
        return None;
    }
    let mut typo = chars;
    typo.remove(typo.len() / 2);
    tokens[idx] = typo.into_iter().collect::<String>();
    Some(tokens.join(" "))
}

fn stress_alias_variant(query: &str) -> Option<String> {
    let mut changed = false;
    let tokens = query
        .to_lowercase()
        .split_whitespace()
        .map(|token| {
            let normalized = token.trim_matches(|c: char| !c.is_alphanumeric());
            let replacement = match normalized {
                "auth" | "oauth" | "authentication" | "authorization" | "login" | "signin" => {
                    Some("인증")
                }
                "token" => Some("토큰"),
                "session" => Some("세션"),
                "memory" => Some("메모리"),
                "queue" => Some("큐"),
                "benchmark" => Some("벤치마크"),
                "인증" => Some("auth"),
                "토큰" => Some("token"),
                "세션" => Some("session"),
                "메모리" => Some("memory"),
                "큐" => Some("queue"),
                "벤치마크" => Some("benchmark"),
                _ => None,
            };
            replacement.map_or_else(
                || token.to_string(),
                |alias| {
                    changed = true;
                    alias.to_string()
                },
            )
        })
        .collect::<Vec<_>>();

    if changed {
        Some(tokens.join(" "))
    } else {
        None
    }
}

fn stress_paraphrase_variant(query: &str) -> Option<String> {
    let lower = query.trim().to_ascii_lowercase();
    if lower.starts_with("how to ") {
        None
    } else if query.split_whitespace().count() >= 2 {
        Some(format!("how to {query}"))
    } else {
        None
    }
}
