use axum::{
Router,
body::Body,
extract::State,
http::{StatusCode, header},
response::{IntoResponse, Response, Sse, sse::Event},
routing::get,
};
use rheo_core::{Result, constants::HTML_EXT};
use std::convert::Infallible;
use std::net::SocketAddr;
use std::path::{Path, PathBuf};
use tokio::sync::broadcast;
use tokio_stream::{StreamExt, wrappers::BroadcastStream};
use tracing::{info, warn};
#[derive(Clone)]
pub struct ServerState {
pub reload_tx: broadcast::Sender<()>,
pub html_dir: PathBuf,
}
pub async fn start_server(
html_dir: PathBuf,
port: u16,
) -> Result<(tokio::task::JoinHandle<()>, broadcast::Sender<()>, String)> {
let (reload_tx, _) = broadcast::channel(100);
let state = ServerState {
reload_tx: reload_tx.clone(),
html_dir: html_dir.clone(),
};
let app = Router::new()
.route("/events", get(sse_handler))
.fallback(get(static_handler))
.with_state(state);
let addr = SocketAddr::from(([127, 0, 0, 1], port));
let listener = tokio::net::TcpListener::bind(addr)
.await
.map_err(|e| crate::RheoError::io(e, format!("binding to {}", addr)))?;
let server_url = format!("http://localhost:{}", port);
info!(url = %server_url, "web server started");
let server_handle = tokio::spawn(async move {
if let Err(e) = axum::serve(listener, app).await {
warn!(error = %e, "web server stopped with error");
}
});
Ok((server_handle, reload_tx, server_url))
}
async fn sse_handler(
State(state): State<ServerState>,
) -> Sse<impl tokio_stream::Stream<Item = std::result::Result<Event, Infallible>>> {
let rx = state.reload_tx.subscribe();
let stream = BroadcastStream::new(rx).filter_map(|result| match result {
Ok(_) => Some(Ok(Event::default().event("reload").data("refresh"))),
Err(_) => None, });
Sse::new(stream).keep_alive(
axum::response::sse::KeepAlive::new()
.interval(std::time::Duration::from_secs(30))
.text("ping"),
)
}
async fn static_handler(State(state): State<ServerState>, uri: axum::http::Uri) -> Response {
let path = uri.path().trim_start_matches('/');
let file_path = if path.is_empty() || path.ends_with('/') {
let index_path = state.html_dir.join(path).join("index.html");
if index_path.exists() {
index_path
} else {
return directory_listing(&state.html_dir, path).into_response();
}
} else {
state.html_dir.join(path)
};
if !file_path.exists() {
return (StatusCode::NOT_FOUND, "404 Not Found").into_response();
}
let content = match tokio::fs::read(&file_path).await {
Ok(content) => content,
Err(_) => {
return (StatusCode::INTERNAL_SERVER_ERROR, "Failed to read file").into_response();
}
};
if path.ends_with(HTML_EXT) {
match inject_live_reload_script(&content) {
Ok(modified_content) => {
return Response::builder()
.status(StatusCode::OK)
.header(header::CONTENT_TYPE, "text/html; charset=utf-8")
.body(Body::from(modified_content))
.unwrap();
}
Err(_) => {
warn!(file = ?file_path, "failed to inject live reload script");
}
}
}
let content_type = mime_guess::from_path(&file_path)
.first_or_octet_stream()
.to_string();
Response::builder()
.status(StatusCode::OK)
.header(header::CONTENT_TYPE, content_type)
.body(Body::from(content))
.unwrap()
}
fn inject_live_reload_script(html: &[u8]) -> std::io::Result<Vec<u8>> {
let html_str = String::from_utf8_lossy(html);
const SCRIPT: &str = r#"
<script>
const eventSource = new EventSource('/events');
eventSource.addEventListener('reload', function(e) {
console.log('Reloading page...');
location.reload();
});
eventSource.onerror = function(e) {
console.log('SSE connection error, will retry automatically');
};
</script>
"#;
let modified = if let Some(pos) = html_str.rfind("</body>") {
let mut result = String::with_capacity(html_str.len() + SCRIPT.len());
result.push_str(&html_str[..pos]);
result.push_str(SCRIPT);
result.push_str(&html_str[pos..]);
result
} else {
let mut result = String::with_capacity(html_str.len() + SCRIPT.len());
result.push_str(&html_str);
result.push_str(SCRIPT);
result
};
Ok(modified.into_bytes())
}
fn directory_listing(html_dir: &Path, path: &str) -> Response {
let dir_path = html_dir.join(path);
let entries = match std::fs::read_dir(&dir_path) {
Ok(entries) => entries,
Err(_) => return (StatusCode::NOT_FOUND, "Directory not found").into_response(),
};
let mut html_files: Vec<String> = entries
.filter_map(|entry| {
let entry = entry.ok()?;
let path = entry.path();
if path.extension()? == "html" {
Some(path.file_name()?.to_string_lossy().to_string())
} else {
None
}
})
.collect();
html_files.sort();
let mut html = String::from(
r#"<!DOCTYPE html>
<html>
<head>
<meta charset="utf-8">
<title>rheo - HTML Output</title>
<style>
body { font-family: system-ui, sans-serif; max-width: 800px; margin: 40px auto; padding: 0 20px; }
h1 { color: #333; }
ul { list-style: none; padding: 0; }
li { margin: 10px 0; }
a { color: #0066cc; text-decoration: none; padding: 8px 12px; display: inline-block; border-radius: 4px; }
a:hover { background: #f0f0f0; }
</style>
</head>
<body>
<h1>HTML Output Files</h1>
<ul>
"#,
);
if html_files.is_empty() {
html.push_str("<li>No HTML files found</li>");
} else {
for file in html_files {
html.push_str(&format!(
r#" <li><a href="{}">{}</a></li>
"#,
file, file
));
}
}
html.push_str(
r#" </ul>
</body>
</html>"#,
);
Response::builder()
.status(StatusCode::OK)
.header(header::CONTENT_TYPE, "text/html; charset=utf-8")
.body(Body::from(html))
.unwrap()
}
pub fn open_browser(url: &str) -> Result<()> {
info!(url = %url, "opening browser");
webbrowser::open(url)
.map_err(|e| crate::RheoError::project_config(format!("failed to open browser: {}", e)))
}