// src/training/websearch.rs
#![allow(dead_code)]

use crate::data_ls::schema::EnglishSample;
use anyhow::{anyhow, Result};
use ollama_rs::generation::chat::request::ChatMessageRequest;
use ollama_rs::generation::chat::ChatMessage;
use ollama_rs::generation::completion::request::GenerationRequest;
use ollama_rs::generation::parameters::{FormatType, KeepAlive};
use ollama_rs::generation::embeddings::request::{EmbeddingsInput, GenerateEmbeddingsRequest};
use ollama_rs::models::ModelOptions;
use ollama_rs::Ollama;
use reqwest::{header, Client};
use scraper::{Html, Selector};
use std::fs::OpenOptions;
use std::io::Write;
use std::path::Path;
use urlencoding::encode;

fn extract_json_from_response(text: &str) -> Option<String> {
    let trimmed = text.trim();
    let without_fence = trimmed
        .strip_prefix("```json")
        .or_else(|| trimmed.strip_prefix("```json "))
        .or_else(|| trimmed.strip_prefix("```json\n"))
        .or_else(|| trimmed.strip_prefix("```"))
        .unwrap_or(trimmed)
        .trim_start();

    let without_trailing = without_fence
        .strip_suffix("```")
        .or_else(|| without_fence.strip_suffix("```\n"))
        .or_else(|| without_fence.strip_suffix("``` "))
        .unwrap_or(without_fence)
        .trim();

    let start = without_trailing.find('{')?;
    let end = without_trailing.rfind('}')?;
    if start >= end {
        return None;
    }
    let json_str = &without_trailing[start..=end];
    if json_str.starts_with('{') && json_str.ends_with('}') {
        Some(json_str.to_string())
    } else {
        None
    }
}

fn is_json_braces_balanced(text: &str) -> bool {
    let mut balance = 0;
    for c in text.chars() {
        if c == '{' {
            balance += 1;
        } else if c == '}' {
            balance -= 1;
            if balance < 0 {
                return false;
            }
        }
    }
    balance == 0
}

fn repair_json_braces(json: &str) -> String {
    let mut repaired = json.trim_end().to_string();
    let mut balance = 0;
    for c in repaired.chars() {
        if c == '{' {
            balance += 1;
        } else if c == '}' {
            balance -= 1;
        }
    }
    while balance > 0 {
        repaired.push('}');
        balance -= 1;
    }
    repaired
}

fn clean_output(output: &str) -> String {
    let without_url_spaces = output
        .replace("http ://", "http://")
        .replace("https ://", "https://")
        .replace(" : / /", "://");

    without_url_spaces
        .split_whitespace()
        .collect::<Vec<_>>()
        .join(" ")
}

fn validate_sample(sample: &EnglishSample) -> Result<()> {
    let instruction_lower = sample.instruction.to_lowercase();
    let forbidden_phrases = ["structured learning sample", "valid json format"];
    for forbidden in forbidden_phrases {
        if instruction_lower.contains(forbidden) {
            return Err(anyhow!(
                "Parsed instruction contains forbidden meta-language: '{}'",
                forbidden
            ));
        }
    }

    if sample.output.trim().is_empty() {
        return Err(anyhow!(
            "Parsed sample output is empty. The model must produce a direct answer or summary."
        ));
    }

    Ok(())
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum SearchEngine {
    Dolphin,
    DuckDuckGo,
    Deepseek,
    Qwen,
    Grok,
}

impl std::fmt::Display for SearchEngine {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(
            f,
            "{}",
            match self {
                SearchEngine::Dolphin => "Dolphin Chat",
                SearchEngine::DuckDuckGo => "DuckDuckGo",
                SearchEngine::Deepseek => "Deepseek",
                SearchEngine::Qwen => "Qwen",
                SearchEngine::Grok => "Grok",
            }
        )
    }
}

impl SearchEngine {
    pub fn search_url(&self, query: &str) -> String {
        match self {
            SearchEngine::Dolphin => format!("https://chat.dphn.ai/?q={}", encode(query)),
            SearchEngine::DuckDuckGo => {
                format!("https://html.duckduckgo.com/html/?q={}", encode(query))
            }
            SearchEngine::Deepseek => format!("https://deepseek.ai/search?q={}", encode(query)),
            SearchEngine::Qwen => format!("https://qwen.ai/search?q={}", encode(query)),
            SearchEngine::Grok => format!("https://grok.com/search?q={}", encode(query)),
        }
    }
}

fn create_browser_client() -> Result<Client, reqwest::Error> {
    let mut headers = header::HeaderMap::new();
    headers.insert(
        header::ACCEPT,
        header::HeaderValue::from_static(
            "text/html,application/xhtml+xml,application/xml;q=0.9,*/*;q=0.8",
        ),
    );
    headers.insert(
        header::ACCEPT_LANGUAGE,
        header::HeaderValue::from_static("en-US,en;q=0.9"),
    );
    headers.insert(
        header::CONNECTION,
        header::HeaderValue::from_static("keep-alive"),
    );
    headers.insert(
        header::USER_AGENT,
        header::HeaderValue::from_static("Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/125.0.0.0 Safari/537.36"),
    );
    Client::builder().default_headers(headers).build()
}

pub async fn web_search(query: &str, engine: SearchEngine) -> Result<String> {
    let url = engine.search_url(query);
    let client = create_browser_client()?;
    let body = client
        .get(&url)
        .header(header::REFERER, "https://chat.dphn.ai/")
        .send()
        .await?
        .error_for_status()?
        .text()
        .await?;
    let doc = Html::parse_document(&body);

    let snippets: Vec<String> = match engine {
        SearchEngine::Dolphin => {
            let selector = Selector::parse("p, div, span").unwrap();
            let mut snippets: Vec<String> = doc
                .select(&selector)
                .filter_map(|el| {
                    let text = el.text().collect::<Vec<_>>().join(" ").trim().to_string();
                    if text.len() >= 20 && text.len() <= 300 {
                        Some(text)
                    } else {
                        None
                    }
                })
                .take(5)
                .collect();

            if snippets.is_empty() {
                let body_text = doc.root_element().text().collect::<Vec<_>>().join(" ");
                snippets = body_text
                    .split_terminator(|c| c == '.' || c == '!' || c == '?')
                    .map(str::trim)
                    .filter(|text| text.len() >= 20 && text.len() <= 300)
                    .take(5)
                    .map(String::from)
                    .collect();
            }

            snippets
        }
        SearchEngine::DuckDuckGo => {
            let selector = Selector::parse(".result__snippet").unwrap();
            doc.select(&selector)
                .take(5)
                .map(|el| el.inner_html())
                .collect()
        }
        SearchEngine::Deepseek => {
            let selector = Selector::parse(".snippet, .search-result, .result, p").unwrap();
            doc.select(&selector)
                .filter_map(|el| {
                    let text = el.text().collect::<Vec<_>>().join(" ").trim().to_string();
                    if text.len() >= 20 && text.len() <= 300 {
                        Some(text)
                    } else {
                        None
                    }
                })
                .take(5)
                .collect()
        }
        SearchEngine::Qwen => {
            let selector = Selector::parse(".search-result, .result, p").unwrap();
            doc.select(&selector)
                .filter_map(|el| {
                    let text = el.text().collect::<Vec<_>>().join(" ").trim().to_string();
                    if text.len() >= 20 && text.len() <= 300 {
                        Some(text)
                    } else {
                        None
                    }
                })
                .take(5)
                .collect()
        }
        SearchEngine::Grok => {
            let selector = Selector::parse(".search-result, .result, p").unwrap();
            doc.select(&selector)
                .filter_map(|el| {
                    let text = el.text().collect::<Vec<_>>().join(" ").trim().to_string();
                    if text.len() >= 20 && text.len() <= 300 {
                        Some(text)
                    } else {
                        None
                    }
                })
                .take(5)
                .collect()
        }
    };

    if snippets.is_empty() {
        return Err(anyhow!(
            "{} returned no usable results from {}",
            engine,
            url
        ));
    }

    Ok(format!("{}\n\nSearch page: {}", snippets.join(" "), url))
}

/// Use LLM to convert search result into a structured EnglishSample
pub async fn derive_sample_from_search(
    query: &str,
    snippets: &str,
    ollama: &Ollama,
) -> Result<EnglishSample> {
    let prompt = format!(
        "Convert the following search result into a structured learning sample in JSON format.\n\
         Use this exact schema:\n\
         {{\"type\": \"instruction_following\", \"instruction\": \"{}\", \"input\": \"\", \"reasoning\": \"...\", \"output\": \"...\"}}\n\n\
         Query: {}\nSearch result: {}\n\nOnly output valid JSON, no extra text. Fill in the reasoning and output based on the search result. The instruction should be the query itself or a rephrased version. The type is always \"instruction_following\".",
        query, query, snippets
    );
    let req = ChatMessageRequest::new("sadiq-bd/llama3.2-3b-uncensored".to_string(), vec![ChatMessage::user(prompt)])
        .options(
            ModelOptions::default()
                .num_ctx(8192)
                .num_thread(8)
                .num_predict(1024)
                .temperature(0.85),
        )
        .keep_alive(KeepAlive::Indefinitely);
    let resp = ollama.send_chat_messages(req).await?;
    let json_str = resp.message.content;
    let sample: EnglishSample = serde_json::from_str(&json_str)?;
    Ok(sample)
}

/// Convert a large PDF text into a sequence of structured EnglishSample objects.
pub async fn derive_samples_from_pdf_text(
    pdf_path: &Path,
    text_chunks: &[String],
    ollama: &Ollama,
    model_name: &str,
) -> Result<Vec<EnglishSample>> {
    let mut samples = Vec::new();
    let total = text_chunks.len();
    for (index, chunk) in text_chunks.iter().enumerate() {
        let prompt = format!(
            "Extract ONE learning sample from the following text in {}.\n\
             Use this exact schema: {{\"type\": \"instruction_following\", \"instruction\": \"...\", \"input\": \"\", \"reasoning\": \"...\", \"output\": \"...\"}}.\n\
             The instruction must be a specific technical question or task based *only* on the text.\n\
             The output must be a direct answer or summary (no empty strings unless truly absent).\n\
             Do NOT include any meta-dialogue, do NOT describe the format, and do NOT output anything except the JSON object.\n\nText:\n{}",
            pdf_path.file_name().unwrap_or_default().to_string_lossy(),
            chunk
        );

        let opts = ModelOptions::default()
            .num_ctx(8192)
            .num_thread(8)
            .num_predict(2048)
            .temperature(0.85);
        let request = GenerationRequest::new(model_name.to_string(), prompt.clone())
            .format(FormatType::Json)
            .options(opts.clone());

        let resp = match ollama.generate(request).await {
            Ok(r) => r,
            Err(e) if model_name == "deepseek-r1:7b" => {
                eprintln!(
                    "⚠️ deepseek-r1:7b failed, falling back to qwen2.5-coder:1.5b. Error: {}",
                    e
                );
                let fallback_req = GenerationRequest::new("qwen2.5-coder:1.5b".to_string(), prompt.clone())
                    .format(FormatType::Json)
                    .options(opts.clone());
                match ollama.generate(fallback_req).await {
                    Ok(r) => r,
                    Err(_) => {
                        eprintln!(
                            "⚠️ qwen2.5-coder:1.5b also failed; falling back to sadiq-bd/llama3.2-3b-uncensored."
                        );
                        let fallback_req = GenerationRequest::new("sadiq-bd/llama3.2-3b-uncensored".to_string(), prompt)
                            .format(FormatType::Json)
                            .options(opts);
                        ollama.generate(fallback_req).await?
                    }
                }
            }
            Err(e) if model_name == "qwen2.5-coder:1.5b" => {
                eprintln!(
                    "⚠️ qwen2.5-coder:1.5b failed, falling back to sadiq-bd/llama3.2-3b-uncensored. Error: {}",
                    e
                );
                let fallback_req = GenerationRequest::new("sadiq-bd/llama3.2-3b-uncensored".to_string(), prompt)
                    .format(FormatType::Json)
                    .options(opts);
                ollama.generate(fallback_req).await?
            }
            Err(e) => return Err(e.into()),
        };

        let raw_response = resp.response;
        let mut json_candidate = extract_json_from_response(&raw_response)
            .unwrap_or_else(|| raw_response.trim().to_string());
        if !is_json_braces_balanced(&json_candidate) {
            json_candidate = repair_json_braces(&json_candidate);
        }

        let sample: EnglishSample = match serde_json::from_str::<EnglishSample>(&json_candidate) {
            Ok(mut s) => {
                s.output = clean_output(&s.output);
                s
            }
            Err(err) => {
                let repaired_candidate = repair_json_braces(&json_candidate);
                if repaired_candidate != json_candidate {
                    serde_json::from_str(&repaired_candidate).map_err(|second_err| {
                        anyhow::anyhow!(
                            "Failed to parse JSON from PDF chunk {} of {}: {}\nRetry parse error: {}\nCandidate: {}\nRepaired: {}\nRaw response: {}",
                            index + 1,
                            total,
                            err,
                            second_err,
                            json_candidate,
                            repaired_candidate,
                            raw_response
                        )
                    })?
                } else if let Some(cleaned) = extract_json_from_response(&raw_response) {
                    serde_json::from_str(&cleaned).map_err(|second_err| {
                        anyhow::anyhow!(
                            "Failed to parse JSON from PDF chunk {} of {}: {}\nRetry parse error: {}\nCleaned response: {}\nRaw response: {}",
                            index + 1,
                            total,
                            err,
                            second_err,
                            cleaned,
                            raw_response
                        )
                    })?
                } else {
                    return Err(anyhow::anyhow!(
                        "Failed to parse JSON from PDF chunk {} of {}.\nRaw response:\n{}",
                        index + 1,
                        total,
                        raw_response
                    ));
                }
            }
        };

        validate_sample(&sample).map_err(|validation_err| {
            anyhow::anyhow!(
                "Validation failed for PDF chunk {} of {}: {}\nCandidate: {}\nRaw response: {}",
                index + 1,
                total,
                validation_err,
                json_candidate,
                raw_response
            )
        })?;

        samples.push(sample);
    }
    Ok(samples)
}

/// Append a new sample to the dynamic knowledge file
pub fn append_sample_to_jsonl(sample: &EnglishSample, path: &str) -> Result<()> {
    let mut file = OpenOptions::new().create(true).append(true).open(path)?;
    let line = serde_json::to_string(sample)?;
    writeln!(file, "{}", line)?;
    Ok(())
}

/// Compute embedding for a sample (so we can add it to in‑memory vectors)
pub async fn compute_sample_embedding(
    sample: &EnglishSample,
    embed_model: &str,
    ollama: &Ollama,
) -> Result<Vec<f32>> {
    let text = format!(
        "Instruction: {}\nOutput: {}",
        sample.instruction, sample.output
    );
    let req =
        GenerateEmbeddingsRequest::new(embed_model.to_string(), EmbeddingsInput::Single(text))
        .keep_alive(KeepAlive::Indefinitely);
    let resp = ollama.generate_embeddings(req).await?;
    Ok(resp.embeddings.into_iter().next().unwrap_or_default())
}