lobe-core 0.1.1

use std::sync::Arc;
use std::time::Instant;

use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
use tokio::net::{lookup_host, TcpStream};
use tokio_rustls::rustls::pki_types::ServerName;
use tokio_rustls::rustls::{ClientConfig, RootCertStore};
use tokio_rustls::TlsConnector;
use url::Url;
use webpki_roots::TLS_SERVER_ROOTS;

use crate::engine::timing::{build_report_from_phases, TimingPhases};
use crate::engine::tls::{report_for_http, report_for_https_failure, report_for_https_success};
use crate::error::{Result, TloxError};
use crate::models::{CertificateInfo, NewProbeResult, ProbeResult, ProbeStatus, TimingReport, TlsReport};
use crate::storage::SqliteStore;

#[derive(Debug, Clone, Default)]
pub struct ProxyMeasurement {
    phases: TimingPhases,
}

impl ProxyMeasurement {
    pub fn new() -> Self {
        Self::default()
    }

    pub fn record_dns_ms(&mut self, dns_ms: u64) {
        self.phases.dns_ms = dns_ms;
    }

    pub fn record_tcp_ms(&mut self, tcp_ms: u64) {
        self.phases.tcp_ms = tcp_ms;
    }

    pub fn record_tls_ms(&mut self, tls_ms: u64) {
        self.phases.tls_ms = tls_ms;
    }

    pub fn record_ttfb_ms(&mut self, ttfb_ms: u64) {
        self.phases.ttfb_ms = ttfb_ms;
    }

    pub fn record_download_ms(&mut self, download_ms: u64) {
        self.phases.download_ms = download_ms;
    }

    pub fn finish_report(&self) -> TimingReport {
        build_report_from_phases(self.phases)
    }
}

#[derive(Debug, Clone, Default)]
pub struct SimpleProbe;

#[derive(Debug, Clone, PartialEq, Eq)]
struct HttpResponseMetadata {
    total_ms: u64,
    status_code: Option<u16>,
    status_text: Option<String>,
    bytes: u64,
}

impl SimpleProbe {
    pub fn new() -> Self {
        Self
    }

    pub async fn probe(&self, target: &str) -> Result<NewProbeResult> {
        let url = Url::parse(target).map_err(|_| TloxError::InvalidTarget(target.to_string()))?;
        let is_https = url.scheme() == "https";
        let host = url
            .host_str()
            .ok_or_else(|| TloxError::InvalidTarget(target.to_string()))?
            .to_string();
        let port = url
            .port_or_known_default()
            .ok_or_else(|| TloxError::InvalidTarget(target.to_string()))?;
        let request_target = build_request_target(&url);

        let mut measurement = ProxyMeasurement::new();
        let total_start = Instant::now();

        let dns_start = Instant::now();
        let mut addresses = match lookup_host((host.as_str(), port)).await {
            Ok(addresses) => addresses,
            Err(error) => {
                return Ok(failed_probe_result(
                    target,
                    &request_target,
                    finish_report(&measurement, total_start.elapsed().as_millis() as u64),
                    tls_report_for_failure(is_https, error.to_string()),
                    error.to_string(),
                ));
            }
        };
        measurement.record_dns_ms(dns_start.elapsed().as_millis() as u64);

        let Some(address) = addresses.next() else {
            let error_message = "no addresses resolved".to_string();

            return Ok(failed_probe_result(
                target,
                &request_target,
                finish_report(&measurement, total_start.elapsed().as_millis() as u64),
                tls_report_for_failure(is_https, error_message.clone()),
                error_message,
            ));
        };

        let tcp_start = Instant::now();
        let stream = match TcpStream::connect(address).await {
            Ok(stream) => stream,
            Err(error) => {
                return Ok(failed_probe_result(
                    target,
                    &request_target,
                    finish_report(&measurement, total_start.elapsed().as_millis() as u64),
                    tls_report_for_failure(is_https, error.to_string()),
                    error.to_string(),
                ));
            }
        };
        measurement.record_tcp_ms(tcp_start.elapsed().as_millis() as u64);

        if is_https {
            self.probe_https(target, &host, &request_target, stream, measurement, total_start)
                .await
        } else {
            self.probe_http(target, &host, &request_target, stream, measurement, total_start)
                .await
        }
    }

    pub async fn probe_and_store(&self, store: &SqliteStore, target: &str) -> Result<ProbeResult> {
        let result = self.probe(target).await?;

        store.save_result(&result)
    }

    async fn probe_http(
        &self,
        target: &str,
        host: &str,
        request_target: &str,
        mut stream: TcpStream,
        mut measurement: ProxyMeasurement,
        total_start: Instant,
    ) -> Result<NewProbeResult> {
        match execute_request(
            &mut stream,
            host,
            request_target,
            &mut measurement,
            total_start,
        )
        .await
        {
            Ok(response) => Ok(NewProbeResult {
                target: target.to_string(),
                request_path: request_target.to_string(),
                report: finish_report(&measurement, response.total_ms),
                status: ProbeStatus::Succeeded,
                response_status_code: response.status_code,
                response_status_text: response.status_text,
                response_bytes: Some(response.bytes),
                error_message: None,
                tls: Some(report_for_http()),
            }),
            Err(error) => Ok(failed_probe_result(
                target,
                request_target,
                finish_report(&measurement, total_start.elapsed().as_millis() as u64),
                Some(report_for_http()),
                error.to_string(),
            )),
        }
    }

    async fn probe_https(
        &self,
        target: &str,
        host: &str,
        request_target: &str,
        stream: TcpStream,
        mut measurement: ProxyMeasurement,
        total_start: Instant,
    ) -> Result<NewProbeResult> {
        let tls_start = Instant::now();
        let connector = build_tls_connector();
        let server_name = ServerName::try_from(host.to_string())
            .map_err(|_| TloxError::InvalidTarget(target.to_string()))?;

        let mut tls_stream = match connector.connect(server_name, stream).await {
            Ok(stream) => stream,
            Err(error) => {
                return Ok(failed_probe_result(
                    target,
                    request_target,
                    finish_report(&measurement, total_start.elapsed().as_millis() as u64),
                    Some(report_for_https_failure(error.to_string())),
                    error.to_string(),
                ));
            }
        };
        measurement.record_tls_ms(tls_start.elapsed().as_millis() as u64);

        match execute_request(
            &mut tls_stream,
            host,
            request_target,
            &mut measurement,
            total_start,
        )
        .await
        {
            Ok(response) => {
                let (_, connection) = tls_stream.get_ref();
                let tls = Some(report_for_https_success(
                    connection.protocol_version().map(|value| format!("{value:?}")),
                    connection
                        .negotiated_cipher_suite()
                        .map(|value| format!("{:?}", value.suite())),
                    extract_certificate_info(connection.peer_certificates()),
                ));

                Ok(NewProbeResult {
                    target: target.to_string(),
                    request_path: request_target.to_string(),
                    report: finish_report(&measurement, response.total_ms),
                    status: ProbeStatus::Succeeded,
                    response_status_code: response.status_code,
                    response_status_text: response.status_text,
                    response_bytes: Some(response.bytes),
                    error_message: None,
                    tls,
                })
            }
            Err(error) => Ok(failed_probe_result(
                target,
                request_target,
                finish_report(&measurement, total_start.elapsed().as_millis() as u64),
                Some(report_for_https_failure(error.to_string())),
                error.to_string(),
            )),
        }
    }
}

fn failed_probe_result(
    target: &str,
    request_path: &str,
    report: TimingReport,
    tls: Option<TlsReport>,
    error_message: String,
) -> NewProbeResult {
    NewProbeResult {
        target: target.to_string(),
        request_path: request_path.to_string(),
        report,
        status: ProbeStatus::Failed,
        response_status_code: None,
        response_status_text: None,
        response_bytes: None,
        error_message: Some(error_message),
        tls,
    }
}

fn finish_report(measurement: &ProxyMeasurement, total_ms: u64) -> TimingReport {
    let mut report = measurement.finish_report();
    report.total_ms = total_ms;
    report
}

fn tls_report_for_failure(is_https: bool, error_message: String) -> Option<TlsReport> {
    if is_https {
        Some(report_for_https_failure(error_message))
    } else {
        Some(report_for_http())
    }
}

fn build_request_target(url: &Url) -> String {
    let mut target = url.path().to_string();
    if target.is_empty() {
        target.push('/');
    }

    if let Some(query) = url.query() {
        target.push('?');
        target.push_str(query);
    }

    target
}

fn build_tls_connector() -> TlsConnector {
    let roots = RootCertStore::from_iter(TLS_SERVER_ROOTS.iter().cloned());
    let config = ClientConfig::builder()
        .with_root_certificates(roots)
        .with_no_client_auth();

    TlsConnector::from(Arc::new(config))
}

async fn execute_request<S>(
    stream: &mut S,
    host: &str,
    request_target: &str,
    measurement: &mut ProxyMeasurement,
    total_start: Instant,
) -> std::io::Result<HttpResponseMetadata>
where
    S: AsyncRead + AsyncWrite + Unpin,
{
    let request = format!(
        "GET {request_target} HTTP/1.1\r\nHost: {host}\r\nUser-Agent: lobe/0.1\r\nAccept: */*\r\nConnection: close\r\n\r\n"
    );

    let ttfb_start = Instant::now();
    stream.write_all(request.as_bytes()).await?;
    stream.flush().await?;

    let mut buffer = Vec::new();
    let mut chunk = [0_u8; 1024];
    let mut download_start: Option<Instant> = None;

    loop {
        let bytes_read = stream.read(&mut chunk).await?;
        if bytes_read == 0 {
            break;
        }

        if download_start.is_none() {
            measurement.record_ttfb_ms(ttfb_start.elapsed().as_millis() as u64);
            download_start = Some(Instant::now());
        }

        buffer.extend_from_slice(&chunk[..bytes_read]);
    }

    if let Some(start) = download_start {
        measurement.record_download_ms(start.elapsed().as_millis() as u64);
    }

    let metadata = parse_http_response(&buffer);

    Ok(HttpResponseMetadata {
        total_ms: total_start.elapsed().as_millis() as u64,
        status_code: metadata.status_code,
        status_text: metadata.status_text,
        bytes: metadata.bytes,
    })
}

fn parse_http_response(buffer: &[u8]) -> HttpResponseMetadata {
    let Some(header_end) = find_header_end(buffer) else {
        return HttpResponseMetadata {
            total_ms: 0,
            status_code: None,
            status_text: None,
            bytes: buffer.len() as u64,
        };
    };

    let header_bytes = &buffer[..header_end];
    let body_bytes = &buffer[header_end + 4..];
    let status_line = header_bytes
        .split(|byte| *byte == b'\n')
        .next()
        .and_then(|line| std::str::from_utf8(line).ok())
        .map(str::trim)
        .unwrap_or_default();

    let mut parts = status_line.splitn(3, ' ');
    let _http_version = parts.next();
    let status_code = parts.next().and_then(|part| part.parse::<u16>().ok());
    let status_text = parts.next().map(|part| part.trim().to_string());

    HttpResponseMetadata {
        total_ms: 0,
        status_code,
        status_text,
        bytes: body_bytes.len() as u64,
    }
}

fn find_header_end(buffer: &[u8]) -> Option<usize> {
    buffer.windows(4).position(|window| window == b"\r\n\r\n")
}

fn extract_certificate_info(
    certificates: Option<&[tokio_rustls::rustls::pki_types::CertificateDer<'_>]>,
) -> Option<CertificateInfo> {
    let certificate = certificates?.first()?;
    let (_, parsed) = x509_parser::parse_x509_certificate(certificate.as_ref()).ok()?;

    Some(CertificateInfo {
        subject: parsed.subject().to_string(),
        issuer: parsed.issuer().to_string(),
        not_before: Some(parsed.validity().not_before.to_string()),
        not_after: Some(parsed.validity().not_after.to_string()),
    })
}

#[cfg(test)]
mod tests {
    use url::Url;

    use crate::models::ProbeStatus;
    use crate::models::TlsStatus;
    use crate::storage::SqliteStore;

    use super::{build_request_target, ProxyMeasurement, SimpleProbe};

    #[test]
    fn finish_report_builds_report_from_recorded_phases() {
        let mut measurement = ProxyMeasurement::new();
        measurement.record_dns_ms(12);
        measurement.record_tcp_ms(18);
        measurement.record_tls_ms(24);
        measurement.record_ttfb_ms(30);
        measurement.record_download_ms(6);

        let report = measurement.finish_report();

        assert_eq!(report.dns_ms, 12);
        assert_eq!(report.tcp_ms, 18);
        assert_eq!(report.tls_ms, 24);
        assert_eq!(report.ttfb_ms, 30);
        assert_eq!(report.download_ms, 6);
        assert_eq!(report.total_ms, 90);
    }

    #[test]
    fn build_request_target_includes_query_string() {
        let url = Url::parse("https://example.com/path?hello=world").expect("url should parse");

        let target = build_request_target(&url);

        assert_eq!(target, "/path?hello=world");
    }

    #[tokio::test]
    async fn simple_probe_returns_failed_result_for_unreachable_target() {
        let probe = SimpleProbe::new();
        let result = probe
            .probe("http://127.0.0.1:1")
            .await
            .expect("probe should return a result");

        assert_eq!(result.status, ProbeStatus::Failed);
        assert_eq!(result.report.tls_ms, 0);
        assert_eq!(result.report.ttfb_ms, 0);
        assert!(result.report.dns_ms <= result.report.total_ms);
        assert!(result.error_message.is_some());
        assert_eq!(
            result.tls.expect("tls report should exist").status,
            TlsStatus::NotUsed
        );
    }

    #[tokio::test]
    async fn simple_probe_rejects_invalid_targets() {
        let probe = SimpleProbe::new();
        let result = probe.probe("not-a-url").await;

        assert!(matches!(
            result,
            Err(crate::error::TloxError::InvalidTarget(_))
        ));
    }

    #[tokio::test]
    async fn probe_and_store_persists_probe_result() {
        let probe = SimpleProbe::new();
        let store = SqliteStore::in_memory().expect("store should initialize");

        let saved = probe
            .probe_and_store(&store, "http://127.0.0.1:1")
            .await
            .expect("probe result should save");

        let loaded = store
            .latest_result_for_target("http://127.0.0.1:1")
            .expect("latest result should load")
            .expect("target should have a saved result");

        assert_eq!(loaded.id, saved.id);
        assert_eq!(loaded.target, "http://127.0.0.1:1");
        assert_eq!(loaded.status, ProbeStatus::Failed);
        assert!(loaded.error_message.is_some());
    }
}