flq 1.0.0

use rmcp::{
	handler::server::{router::tool::ToolRouter, wrapper::Parameters},
	model::*,
	schemars::{self, JsonSchema},
	tool, tool_handler, tool_router,
	transport::stdio,
	ErrorData as McpError, ServerHandler, ServiceExt,
};
use serde::Deserialize;

use crate::{crawl, extract::Filter, process, Opts};

#[derive(Clone)]
pub struct FlqServer {
	// Read by the #[tool_router] / #[tool_handler] generated code.
	#[allow(dead_code)]
	tool_router: ToolRouter<Self>,
}

#[derive(Debug, Deserialize, JsonSchema)]
pub struct FetchPageParams {
	#[schemars(description = "target URL to fetch")]
	pub url: String,
	#[schemars(
		description = "block filter (comma-separated): headings,code,tables,links,img,text"
	)]
	pub filter: Option<String>,
	#[schemars(description = "CSS selector to scope extraction")]
	pub select: Option<String>,
	#[schemars(description = "approximate token budget for output")]
	pub max_tokens: Option<usize>,
	#[schemars(description = "grep body lines by pattern (case-insensitive)")]
	pub grep: Option<String>,
}

#[derive(Debug, Deserialize, JsonSchema)]
pub struct FetchPagesParams {
	#[schemars(description = "list of URLs to fetch")]
	pub urls: Vec<String>,
	#[schemars(description = "block filter")]
	pub filter: Option<String>,
	#[schemars(description = "approximate token budget per page")]
	pub max_tokens: Option<usize>,
}

#[derive(Debug, Deserialize, JsonSchema)]
pub struct CrawlParams {
	#[schemars(description = "seed URL")]
	pub url: String,
	#[schemars(description = "crawl depth (default 1)")]
	pub depth: Option<u8>,
	#[schemars(description = "max pages to fetch (default 10)")]
	pub max_pages: Option<usize>,
	#[schemars(description = "block filter")]
	pub filter: Option<String>,
}

#[derive(Debug, Deserialize, JsonSchema)]
pub struct HeadParams {
	#[schemars(description = "URL to check")]
	pub url: String,
}

fn parse_filters(s: &Option<String>) -> Vec<Filter> {
	match s {
		Some(s) => Filter::parse(s).unwrap_or_default(),
		None => Vec::new(),
	}
}

fn make_opts(
	filter: &Option<String>,
	select: Option<String>,
	max_tokens: Option<usize>,
	grep: Option<String>,
) -> Opts {
	let filters = parse_filters(filter);
	Opts {
		agent: "Mozilla/5.0 (compatible; flq/0.2)".into(),
		max_time: 30,
		filters,
		select,
		max_tokens,
		grep,
		no_meta: filter.is_some(),
		..Default::default()
	}
}

// ureq is blocking; run each tool body off the async executor.
async fn blocking<T: Send + 'static>(
	f: impl FnOnce() -> Result<T, McpError> + Send + 'static,
) -> Result<T, McpError> {
	tokio::task::spawn_blocking(f)
		.await
		.map_err(|e| McpError::internal_error(e.to_string(), None))?
}

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

#[tool_router]
impl FlqServer {
	pub fn new() -> Self {
		Self {
			tool_router: Self::tool_router(),
		}
	}

	#[tool(description = "Fetch a web page as lean, token-efficient text. \
		Strips HTML to headings, paragraphs, lists, tables, code blocks. \
		Use filter for specific block types: headings,code,tables,links,img,text")]
	async fn fetch_page(
		&self,
		Parameters(p): Parameters<FetchPageParams>,
	) -> Result<String, McpError> {
		blocking(move || {
			let opts = make_opts(&p.filter, p.select, p.max_tokens, p.grep);
			match process(&p.url, &opts) {
				Ok(page) => Ok(if opts.no_meta {
					page.body
				} else {
					format_page(&page)
				}),
				Err(e) => Err(McpError::internal_error(e, None)),
			}
		})
		.await
	}

	#[tool(description = "Fetch multiple web pages in batch. Returns pages separated by ===")]
	async fn fetch_pages(
		&self,
		Parameters(p): Parameters<FetchPagesParams>,
	) -> Result<String, McpError> {
		blocking(move || {
			let opts = make_opts(&p.filter, None, p.max_tokens, None);
			let out = p
				.urls
				.iter()
				.map(|u| match process(u, &opts) {
					Ok(page) => format_page(&page),
					Err(e) => format!("url: {u}\n---\nerror: {e}\n"),
				})
				.collect::<Vec<_>>()
				.join("\n===\n");
			Ok(out)
		})
		.await
	}

	#[tool(
		description = "Crawl a website from a seed URL. Follows same-host links to specified depth"
	)]
	async fn crawl_site(&self, Parameters(p): Parameters<CrawlParams>) -> Result<String, McpError> {
		blocking(move || {
			let opts = make_opts(&p.filter, None, None, None);
			let results = crawl(
				&p.url,
				p.depth.unwrap_or(1),
				p.max_pages.unwrap_or(10),
				&opts,
			);
			let out = results
				.iter()
				.map(|r| match r {
					Ok(page) => format_page(page),
					Err(e) => format!("error: {e}\n"),
				})
				.collect::<Vec<_>>()
				.join("\n===\n");
			Ok(out)
		})
		.await
	}

	#[tool(description = "Get HTTP response headers for a URL (like curl -I)")]
	async fn head(&self, Parameters(p): Parameters<HeadParams>) -> Result<String, McpError> {
		blocking(move || {
			let opts = Opts {
				agent: "Mozilla/5.0 (compatible; flq/0.2)".into(),
				max_time: 30,
				..Default::default()
			};
			match crate::fetch::get(&p.url, &opts) {
				Ok(res) => {
					let mut out = format!("status: {}\n", res.status);
					for (k, v) in &res.headers {
						out.push_str(&format!("{k}: {v}\n"));
					}
					Ok(out)
				}
				Err(e) => Err(McpError::internal_error(e, None)),
			}
		})
		.await
	}
}

#[tool_handler]
impl ServerHandler for FlqServer {
	fn get_info(&self) -> ServerInfo {
		ServerInfo::new(ServerCapabilities::builder().enable_tools().build())
			.with_instructions("Fetch web pages as lean, token-efficient text for LLM consumption")
	}
}

fn format_page(page: &crate::Page) -> String {
	let mut s = String::new();
	if !page.title.is_empty() {
		s.push_str(&format!("title: {}\n", page.title));
	}
	s.push_str(&format!("url: {}\n", page.url));
	if !page.lang.is_empty() {
		s.push_str(&format!("lang: {}\n", page.lang));
	}
	if !page.desc.is_empty() {
		s.push_str(&format!("desc: {}\n", page.desc));
	}
	s.push_str("---\n");
	s.push_str(&page.body);
	if page.truncated {
		s.push_str("\n...(truncated)\n");
	}
	s
}

pub async fn run() -> Result<(), Box<dyn std::error::Error>> {
	let service = FlqServer::new().serve(stdio()).await?;
	service.waiting().await?;
	Ok(())
}