#[macro_use]
extern crate lazy_static;

mod config;
mod countries;
mod mirror;
mod speed_test;
mod target_configs;
mod targets;

use crate::config::{AppError, Config, FetchMirrors};
use crate::speed_test::{SpeedTestResult, SpeedTestResults, test_speed_by_countries};
use chrono::prelude::*;
use config::LogFormatter;
use itertools::Itertools;
use mirror::Mirror;
use nix::unistd::Uid;
use std::env;
use std::fmt::Display;
use std::fs::File;
use std::io;
use std::io::prelude::*;
use std::sync::Arc;
use std::sync::mpsc;
use std::thread;

struct OutputSink<'a, T: LogFormatter> {
    filename: Option<String>,
    output_lines: Option<Vec<String>>,
    formatter: &'a T,
    comments_enabled: bool,
    comments_in_file_enabled: bool,
    mirror_count: usize,
}

impl<'a, T: LogFormatter> OutputSink<'a, T> {
    pub fn new(
        formatter: &'a T,
        filename: Option<&str>,
        comments_enabled: bool,
        comments_in_file_enabled: bool,
    ) -> Result<Self, io::Error> {
        let output = match filename {
            Some(filename) => Self {
                formatter,
                filename: Some(filename.to_string()),
                output_lines: Some(Vec::new()),
                comments_enabled,
                comments_in_file_enabled,
                mirror_count: 0,
            },
            None => Self {
                formatter,
                filename: None,
                output_lines: None,
                comments_enabled,
                comments_in_file_enabled,
                mirror_count: 0,
            },
        };
        Ok(output)
    }

    fn write_stdout_line(&mut self, line: &str) -> Result<(), AppError> {
        let mut stdout = io::stdout().lock();
        writeln!(stdout, "{}", line).map_err(|err| {
            if err.kind() == io::ErrorKind::BrokenPipe {
                AppError::StdoutBrokenPipe
            } else {
                AppError::IoError(err)
            }
        })
    }

    pub fn display_comment(&mut self, line: impl Display) -> Result<(), AppError> {
        if self.comments_enabled {
            let s = self.formatter.format_comment(line);
            self.write_stdout_line(&s)?;
            if self.comments_in_file_enabled {
                if let Some(output_lines) = &mut self.output_lines {
                    output_lines.push(s);
                }
            }
        }
        Ok(())
    }

    pub fn display_mirror(&mut self, mirror: &Mirror) -> Result<(), AppError> {
        let s = self.formatter.format_mirror(&mirror);
        self.write_stdout_line(&s)?;
        if let Some(output_lines) = &mut self.output_lines {
            output_lines.push(s);
        }
        self.mirror_count += 1;
        Ok(())
    }

    pub fn save_to_file(&mut self) -> Result<(), io::Error> {
        if let Some(output_lines) = &mut self.output_lines {
            if let Some(filename) = self.filename.as_ref() {
                let mut f = File::create(filename)?;
                f.write_all(output_lines.join("\n").as_bytes())?;
                f.write_all("\n".as_bytes())?;
            }
        }
        return Ok(());
    }
}

fn main() -> Result<(), AppError> {
    match run() {
        Err(AppError::StdoutBrokenPipe) => Ok(()),
        result => result,
    }
}

fn run() -> Result<(), AppError> {
    let config = Arc::new(Config::new());
    if !config.allow_root && Uid::effective().is_root() {
        return Err(AppError::Root);
    }
    let max_mirrors_to_output = config.max_mirrors_to_output.clone();
    let disable_untested_fallback = config.disable_untested_fallback;

    let ref formatter = Arc::clone(&config).target;
    let mut output = OutputSink::new(
        formatter,
        config.save_to_file.as_deref(),
        !config.disable_comments,
        !config.disable_comments_in_file,
    )?;

    output.display_comment(format!("STARTED AT: {}", Local::now()))?;
    output.display_comment(format!("VERSION: {}", env!("CARGO_PKG_VERSION")))?;
    output.display_comment(format!("ARGS: {}", env::args().join(" ")))?;

    let (tx_progress, rx_progress) = mpsc::channel::<String>();
    let (tx_results, rx_results) = mpsc::channel::<SpeedTestResults>();
    let (tx_mirrors, rx_mirrors) = mpsc::channel::<Mirror>();

    let thread_handle = thread::spawn(move || -> Result<(), AppError> {
        let mut mirrors = config.target.fetch_mirrors(tx_progress.clone())?;

        // Centralized protocol filtering
        let before_protocol = mirrors.len();
        mirrors.retain(|m| config.is_protocol_allowed_for_url(&m.url));
        if mirrors.len() < before_protocol {
            tx_progress
                .send(format!(
                    "PROTOCOL FILTER: {} -> {} mirrors",
                    before_protocol,
                    mirrors.len()
                ))
                .unwrap();
        }

        // Country filtering before dedup so excluded-country duplicates
        // don't shadow valid mirrors from non-excluded countries
        let before_country = mirrors.len();
        mirrors.retain(|m| {
            m.country
                .map(|c| !config.is_country_excluded(c.code))
                .unwrap_or(!config.excluded_countries_set.contains("zz"))
        });
        if mirrors.len() < before_country {
            tx_progress
                .send(format!(
                    "COUNTRY FILTER: {} -> {} mirrors",
                    before_country,
                    mirrors.len()
                ))
                .unwrap();
        }

        // Prefer https over http when both are available for the same host
        mirrors.sort_by_key(|m| match m.url.scheme() {
            "https" => 0,
            "http" => 1,
            _ => 2,
        });

        // Deduplicate mirrors by host+port+path (keeps first = preferred protocol)
        let before_dedup = mirrors.len();
        let mut seen = std::collections::HashSet::new();
        mirrors.retain(|m| {
            let key = format!(
                "{}{}{}",
                m.url.host_str().unwrap_or(""),
                m.url.port().map(|p| format!(":{}", p)).unwrap_or_default(),
                m.url.path()
            );
            seen.insert(key)
        });
        if mirrors.len() < before_dedup {
            tx_progress
                .send(format!(
                    "DEDUP: {} -> {} mirrors",
                    before_dedup,
                    mirrors.len()
                ))
                .unwrap();
        }

        // sending filtered mirrors back so we have a fallback in case if all tests fail
        for mirror in mirrors.iter().cloned() {
            tx_mirrors.send(mirror).unwrap();
        }

        tx_progress
            .send(format!("MIRRORS LEFT AFTER FILTERING: {}", mirrors.len()))
            .unwrap();

        test_speed_by_countries(mirrors, config, tx_progress, tx_results);
        Ok(())
    });

    for progress in rx_progress.into_iter() {
        output.display_comment(progress)?;
    }

    thread_handle.join().unwrap()?;

    let results: Vec<_> = rx_results.iter().flatten().collect();

    if results.is_empty() {
        let untested_mirrors: Vec<Mirror> = rx_mirrors.into_iter().collect();
        if untested_mirrors.len() == 0 {
            output.display_comment("==== NO MIRRORS AFTER FILTERING ====")?;
            return Err(AppError::NoMirrorsAfterFiltering);
        }
        if disable_untested_fallback {
            output.display_comment("==== ALL SPEED TESTS FAILED ====")?;
            return Err(AppError::SpeedTestsFailed);
        }
        output.display_comment("==== FAILED TO TEST SPEEDS, RETURNING UNTESTED MIRRORS ====")?;
        for mirror in untested_mirrors.into_iter() {
            output.display_mirror(&mirror)?;
        }
    } else {
        output.display_comment("==== RESULTS (top re-tested) ====")?;

        for (index, result) in results.iter().enumerate() {
            output.display_comment(format!("{:>3}. {}", index + 1, result))?;
        }

        output.display_comment(format!("FINISHED AT: {}", Local::now()))?;

        let it: Box<dyn Iterator<Item = SpeedTestResult>> = match max_mirrors_to_output {
            Some(n) => Box::new(results.into_iter().take(n)),
            None => Box::new(results.into_iter()),
        };

        for result in it {
            output.display_mirror(&result.item)?;
        }
    }

    if output.mirror_count == 0 {
        return Err(AppError::BlankOutput);
    }
    output.save_to_file()?;
    Ok(())
}
