use std::fs::{self, File};
use std::io::Write;
use std::path::Path;
use std::collections::HashSet;
use reqwest::Client;
use scraper::{Html, Selector};
use tokio::task;
use rayon::prelude::*;
use url::{Url};
use futures;
use std::pin::Pin;
use futures::Future;
use std::sync::{Arc, Mutex};
use std::thread;
use tui::Terminal;
use tui::backend::CrosstermBackend;
use tui::widgets::{Block, Borders, Paragraph};
use tui::layout::{Layout, Constraint, Direction};
use tui::text::{Span, Spans};
use crossterm::event::{self, DisableMouseCapture, EnableMouseCapture, Event, KeyCode};
use crossterm::terminal::{disable_raw_mode, enable_raw_mode, EnterAlternateScreen, LeaveAlternateScreen};
use crossterm::execute;

const SPINNER_FRAMES: &[&str] = &["⠋", "⠙", "⠹", "⠸", "⠼", "⠴", "⠦", "⠧", "⠇", "⠏"];

#[derive(Default)]
struct Status {
    current_url: String,
    pages_visited: usize,
    done: bool,
    spinner_idx: usize,
}

fn run_tui(status: Arc<Mutex<Status>>) {
    enable_raw_mode().unwrap();
    let mut stdout = std::io::stdout();
    execute!(stdout, EnterAlternateScreen, EnableMouseCapture).unwrap();
    let backend = CrosstermBackend::new(stdout);
    let mut terminal = Terminal::new(backend).unwrap();

    let mut frame_count = 0;

    loop {
        let mut s = status.lock().unwrap();
        // Animate spinner
        if !s.done {
            s.spinner_idx = (s.spinner_idx + 1) % SPINNER_FRAMES.len();
        }
        let spinner = SPINNER_FRAMES[s.spinner_idx];
        terminal.draw(|f| {
            let size = f.size();
            // Layout: center box and status bar at bottom
            let layout = Layout::default()
                .direction(Direction::Vertical)
                .margin(2)
                .constraints([
                    Constraint::Percentage(40),
                    Constraint::Length(9), // main box
                    Constraint::Percentage(40),
                    Constraint::Length(1), // status bar
                ]);
            let chunks = layout.split(size);

            // Centered main status box
            let block = Block::default().title(Span::styled(
                format!(" {} Prawn Web Scraper {} ", spinner, spinner),
                tui::style::Style::default().fg(tui::style::Color::Cyan).add_modifier(tui::style::Modifier::BOLD),
            )).borders(Borders::ALL);
            let text = vec![
                Spans::from(vec![Span::styled(
                    format!("Pages visited: {}", s.pages_visited),
                    tui::style::Style::default().fg(tui::style::Color::Green).add_modifier(tui::style::Modifier::BOLD),
                )]),
                Spans::from(vec![Span::raw("")]),
                Spans::from(vec![Span::styled(
                    "Current URL:",
                    tui::style::Style::default().fg(tui::style::Color::Yellow),
                )]),
                Spans::from(vec![Span::styled(
                    format!("{}", s.current_url),
                    tui::style::Style::default().fg(tui::style::Color::White),
                )]),
                Spans::from(vec![Span::raw("")]),
                Spans::from(vec![Span::styled(
                    if s.done {"Done! Press Ctrl+C to exit."} else {"Scraping..."},
                    tui::style::Style::default().fg(if s.done {tui::style::Color::Green} else {tui::style::Color::Blue}).add_modifier(tui::style::Modifier::BOLD),
                )]),
            ];
            let paragraph = Paragraph::new(text).block(block).alignment(tui::layout::Alignment::Center);
            f.render_widget(paragraph, chunks[1]);

            // Status bar at the bottom
            let status_bar = Paragraph::new(Spans::from(vec![
                Span::styled(
                    format!(" {} Prawn scraping in progress...  ", spinner),
                    tui::style::Style::default().fg(tui::style::Color::Magenta).add_modifier(tui::style::Modifier::BOLD),
                ),
                Span::raw(" | Press Ctrl+C to quit "),
            ]))
            .block(Block::default().borders(Borders::NONE))
            .alignment(tui::layout::Alignment::Center);
            f.render_widget(status_bar, chunks[3]);
        }).unwrap();
        if s.done {
            break;
        }
        drop(s);
        std::thread::sleep(std::time::Duration::from_millis(80));
    }
    disable_raw_mode().unwrap();
    let mut stdout = std::io::stdout();
    execute!(stdout, LeaveAlternateScreen, DisableMouseCapture).unwrap();
}

#[tokio::main]
async fn main() {
    let args: Vec<String> = std::env::args().collect();
    if args.len() != 2 {
        eprintln!("Usage: cargo run -- <URL>");
        std::process::exit(1);
    }
    let mut url = args[1].clone();
    if !url.starts_with("http://") && !url.starts_with("https://") {
        url = format!("https://{}", url);
    }
    let base_url = Url::parse(&url).expect("Invalid URL");
    let client = Client::builder().build().unwrap();

    // Prepare output directories
    fs::create_dir_all("saved_site/css").unwrap();
    fs::create_dir_all("saved_site/js").unwrap();

    // Shared status for TUI
    let status = Arc::new(Mutex::new(Status::default()));
    let status_tui = status.clone();

    // Start TUI in a separate thread
    let tui_handle = thread::spawn(move || {
        run_tui(status_tui);
    });

    // Start recursive scraping
    let mut visited = HashSet::new();
    scrape_page(&client, &base_url, &base_url, "saved_site", &mut visited, 0, 2, &status).await;

    // Wait for TUI to finish (send quit signal)
    {
        let mut s = status.lock().unwrap();
        s.done = true;
    }
    tui_handle.join().unwrap();
    println!("Site saved to ./saved_site/");
}

fn scrape_page<'a>(
    client: &'a Client,
    base_url: &'a Url,
    page_url: &'a Url,
    save_dir: &'a str,
    visited: &'a mut HashSet<String>,
    depth: usize,
    max_depth: usize,
    status: &'a Arc<Mutex<Status>>,
) -> Pin<Box<dyn Future<Output = ()> + 'a>> {
    Box::pin(async move {
        {
            let mut s = status.lock().unwrap();
            s.current_url = page_url.as_str().to_string();
        }
        if depth > max_depth {
            return;
        }
        let url_str = page_url.as_str().to_string();
        if visited.contains(&url_str) {
            return;
        }
        visited.insert(url_str.clone());

        {
            let mut s = status.lock().unwrap();
            s.pages_visited += 1;
        }

        let html = match client.get(page_url.clone()).send().await {
            Ok(resp) => match resp.text().await {
                Ok(text) => text,
                Err(_) => return,
            },
            Err(_) => return,
        };
        let document = Html::parse_document(&html);

        // Extract CSS and JS links
        let css_selector = Selector::parse("link[rel=stylesheet]").unwrap();
        let js_selector = Selector::parse("script[src]").unwrap();

        let mut css_links = Vec::new();
        for element in document.select(&css_selector) {
            if let Some(href) = element.value().attr("href") {
                if let Ok(abs_url) = page_url.join(href) {
                    css_links.push((href.to_string(), abs_url.to_string()));
                }
            }
        }

        let mut js_links = Vec::new();
        for element in document.select(&js_selector) {
            if let Some(src) = element.value().attr("src") {
                if let Ok(abs_url) = page_url.join(src) {
                    js_links.push((src.to_string(), abs_url.to_string()));
                }
            }
        }

        // Download CSS and JS concurrently
        let css_results: Vec<(String, String)> = css_links.iter().map(|(orig, _abs)| {
            let filename = Path::new(orig).file_name().unwrap_or_else(|| std::ffi::OsStr::new("style.css")).to_string_lossy().to_string();
            (filename, orig.clone())
        }).collect();
        let css_futures: Vec<Pin<Box<dyn Future<Output = ()>>>> = css_links.iter().map(|(orig, abs)| {
            let filename = Path::new(orig).file_name().unwrap_or_else(|| std::ffi::OsStr::new("style.css")).to_string_lossy().to_string();
            let save_path = format!("{}/css/{}", save_dir, filename);
            let abs = abs.clone();
            Box::pin(async move {
                let client = Client::new();
                match client.get(&abs).send().await {
                    Ok(resp) => {
                        let bytes = resp.bytes().await.unwrap_or_default();
                        let mut file = File::create(&save_path).unwrap();
                        file.write_all(&bytes).unwrap();
                    },
                    Err(_) => {}
                }
            }) as Pin<Box<dyn Future<Output = ()>>>
        }).collect();

        let js_results: Vec<(String, String)> = js_links.iter().map(|(orig, _abs)| {
            let filename = Path::new(orig).file_name().unwrap_or_else(|| std::ffi::OsStr::new("script.js")).to_string_lossy().to_string();
            (filename, orig.clone())
        }).collect();
        let js_futures: Vec<Pin<Box<dyn Future<Output = ()>>>> = js_links.iter().map(|(orig, abs)| {
            let filename = Path::new(orig).file_name().unwrap_or_else(|| std::ffi::OsStr::new("script.js")).to_string_lossy().to_string();
            let save_path = format!("{}/js/{}", save_dir, filename);
            let abs = abs.clone();
            Box::pin(async move {
                let client = Client::new();
                match client.get(&abs).send().await {
                    Ok(resp) => {
                        let bytes = resp.bytes().await.unwrap_or_default();
                        let mut file = File::create(&save_path).unwrap();
                        file.write_all(&bytes).unwrap();
                    },
                    Err(_) => {}
                }
            }) as Pin<Box<dyn Future<Output = ()>>>
        }).collect();

        futures::future::join_all(css_futures).await;
        futures::future::join_all(js_futures).await;

        // Rewrite HTML to point to local files
        let mut new_html = html.clone();
        for (filename, orig) in &css_results {
            let local_path = format!("css/{}", filename);
            new_html = new_html.replace(orig, &local_path);
        }
        for (filename, orig) in &js_results {
            let local_path = format!("js/{}", filename);
            new_html = new_html.replace(orig, &local_path);
        }

        // Save rewritten HTML
        let page_path = if page_url.path() == "/" || page_url.path().is_empty() {
            format!("{}/index.html", save_dir)
        } else {
            let mut path = page_url.path().trim_start_matches('/').replace('/', "_");
            if !path.ends_with(".html") {
                path.push_str(".html");
            }
            format!("{}/{}", save_dir, path)
        };
        if let Some(parent) = Path::new(&page_path).parent() {
            fs::create_dir_all(parent).unwrap();
        }
        let mut file = File::create(&page_path).unwrap();
        file.write_all(new_html.as_bytes()).unwrap();

        // Extract and recursively scrape <a href> links
        let a_selector = Selector::parse("a[href]").unwrap();
        for element in document.select(&a_selector) {
            if let Some(href) = element.value().attr("href") {
                if let Ok(link_url) = page_url.join(href) {
                    // Only follow links within the same domain
                    if link_url.domain() == base_url.domain() {
                        scrape_page(client, base_url, &link_url, save_dir, visited, depth + 1, max_depth, status).await;
                    }
                }
            }
        }
    })
}
