use anyhow::Result;
use headless_chrome::Browser;
use reqwest::Url;
use scraper::{Html, Selector};
use std::collections::HashMap;

use super::ContextItem;

/**
 * Module for gathering context from the web.
 */

// Eventually support dynamic sites with a headless browser
#[derive(Clone, Debug)]
pub struct WebContext {
    static_sites_content: HashMap<Url, String>,
    dynamic_sites_content: HashMap<Url, String>,
    api_content: HashMap<Url, String>,
}

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

impl WebContext {
    pub fn new() -> Self {
        Self {
            static_sites_content: HashMap::new(),
            api_content: HashMap::new(),
            dynamic_sites_content: HashMap::new(),
        }
    }

    pub fn get_sites_content(&self) -> &HashMap<Url, String> {
        &self.static_sites_content
    }

    pub fn get_dynamic_sites_content(&self) -> &HashMap<Url, String> {
        &self.dynamic_sites_content
    }

    pub fn get_api_content(&self) -> &HashMap<Url, String> {
        &self.api_content
    }

    pub fn get_api_content_for_url(&self, url: &Url) -> Option<&String> {
        self.api_content.get(url)
    }

    pub fn get_static_site_content_for_url(&self, url: &Url) -> Option<&String> {
        self.static_sites_content.get(url)
    }

    pub fn get_dynamic_site_content_for_url(&self, url: &Url) -> Option<&String> {
        self.dynamic_sites_content.get(url)
    }

    // Collect static site into context items
    pub fn collect_static_site_into_context_items(&self) -> Result<Vec<ContextItem>> {
        let mut context_items = Vec::new();
        for (url, content) in self.static_sites_content.iter() {
            context_items.push(ContextItem::new(super::ContextContent::Webpage {
                url: url.to_string(),
                content: content.clone(),
            }));
        }
        Ok(context_items)
    }

    pub fn collect_dynamic_sites_into_context_items(&self) -> Result<Vec<ContextItem>> {
        let mut context_items = Vec::new();
        for (url, content) in self.dynamic_sites_content.iter() {
            context_items.push(ContextItem::new(super::ContextContent::Webpage {
                url: url.to_string(),
                content: content.clone(),
            }));
        }
        Ok(context_items)
    }

    // Collect api content into context items
    pub fn collect_api_into_context_items(&self) -> Result<Vec<ContextItem>> {
        let mut context_items = Vec::new();
        for (url, content) in self.api_content.iter() {
            context_items.push(ContextItem::new(super::ContextContent::Api {
                url: url.to_string(),
                content: content.clone(),
            }));
        }
        Ok(context_items)
    }

    pub fn add_content_from_api_url(&mut self, url: &Url) -> Result<()> {
        let content = get_content_from_api_url(url.clone())?;
        self.api_content.insert(url.clone(), content);
        Ok(())
    }

    pub fn add_content_from_dynamic_site_url(&mut self, url: &Url) -> Result<()> {
        let content = get_content_from_dynamic_site_url(url.clone())?;
        self.dynamic_sites_content.insert(url.clone(), content);
        Ok(())
    }

    pub fn add_content_from_static_site_url(&mut self, url: &Url) -> Result<()> {
        let content = get_content_from_static_site_url(url.clone())?;
        self.static_sites_content.insert(url.clone(), content);
        Ok(())
    }
}

fn get_content_from_dynamic_site_url(url: Url) -> Result<String> {
    let browser = Browser::default()?;
    let tab = browser.new_tab()?;
    tab.navigate_to(url.as_ref())?;
    tab.wait_for_element("body")?; // Wait for content to load
    let content = tab.get_content()?;
    let relevant_content = get_relevant_content(content)?;
    Ok(relevant_content)
}

fn get_content_from_api_url(url: Url) -> Result<String> {
    let response = reqwest::blocking::get(url)?;
    match response.json::<serde_json::Value>() {
        Ok(json) => Ok(serde_json::to_string_pretty(&json).unwrap()),
        Err(e) => Err(anyhow::anyhow!("Error parsing JSON: {}", e)),
    }
}

fn get_relevant_content(html_content_string: String) -> Result<String> {
    let document = Html::parse_document(&html_content_string);
    let mut content = String::new();
    // Get main content (prefer main, article, or div with main content classes)
    let main_selectors = [
        "main",
        "article",
        "[role='main']",
        ".main-content",
        ".content",
        "#main-content",
        "body", // fallback
    ];

    for selector in main_selectors {
        if let Ok(sel) = Selector::parse(selector) {
            if let Some(main_content) = document.select(&sel).next() {
                content = extract_clean_text(&main_content);
                break;
            }
        }
    }

    // If no content was found, extract from body
    if content.is_empty() {
        if let Ok(body_selector) = Selector::parse("body") {
            if let Some(body) = document.select(&body_selector).next() {
                content = extract_clean_text(&body);
            }
        }
    }

    Ok(content)
}

fn get_content_from_static_site_url(url: Url) -> Result<String> {
    let response = reqwest::blocking::get(url)?;
    let html = response.text()?;

    get_relevant_content(html)
}

fn extract_clean_text(element: &scraper::element_ref::ElementRef) -> String {
    // Elements to skip
    let skip_selectors = [
        "script",
        "style",
        "nav",
        "header",
        "footer",
        ".navigation",
        ".menu",
        ".sidebar",
        ".ads",
        ".advertisement",
        "#navigation",
        "[role='navigation']",
        "[role='banner']",
        "[role='complementary']",
    ];

    let mut skip_elements = Vec::new();
    for selector in skip_selectors {
        if let Ok(sel) = Selector::parse(selector) {
            skip_elements.push(sel);
        }
    }

    let mut text = String::new();
    let mut stack = vec![*element];

    while let Some(current) = stack.pop() {
        // Skip if element matches any of the skip selectors
        if skip_elements.iter().any(|sel| sel.matches(&current)) {
            continue;
        }

        // Get all text nodes for this element
        for node in current.children() {
            if node.value().is_text() {
                let node_text = node.value().as_text().unwrap().trim();
                if !node_text.is_empty() {
                    text.push_str(node_text);
                    text.push('\n');
                }
            }
        }

        // Add children to stack in reverse order
        for child in current.children().rev() {
            if let Some(child_ref) = scraper::element_ref::ElementRef::wrap(child) {
                stack.push(child_ref);
            }
        }
    }

    // Clean up the text
    clean_text(&text)
}

fn clean_text(text: &str) -> String {
    text.lines()
        .map(|line| line.trim())
        .filter(|line| !line.is_empty())
        .collect::<Vec<_>>()
        .join("\n")
}

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

    #[test]
    fn test_clean_text() {
        let input = "  Hello  \n\n  World  \n\n\n  Test  ";
        let expected = "Hello\nWorld\nTest";
        assert_eq!(clean_text(input), expected);
    }

    #[test]
    fn test_extract_from_html() -> Result<()> {
        let html = r#"
            <html>
                <head><title>Test</title></head>
                <body>
                    <nav>Skip this</nav>
                    <main>
                        <h1>Main Content</h1>
                        <p>This is important content.</p>
                        <script>Skip this too</script>
                        <div class="ads">Skip ads</div>
                    </main>
                    <footer>Skip footer</footer>
                </body>
            </html>
        "#;

        let document = Html::parse_document(html);
        if let Ok(main_selector) = Selector::parse("main") {
            if let Some(main_content) = document.select(&main_selector).next() {
                println!("{}", main_content.html());
                let text = extract_clean_text(&main_content);
                println!("{}", text);
                assert!(text.contains("Main Content"));
                assert!(text.contains("This is important content"));
                assert!(!text.contains("Skip this"));
                assert!(!text.contains("Skip ads"));
                assert!(!text.contains("Skip footer"));
            }
        }

        Ok(())
    }
}
