use std::borrow::Cow;
use std::collections::HashMap;
use std::fs;
use std::io::ErrorKind;
use std::net::SocketAddr;
use std::sync::Arc;
use std::time::Duration;

use axum::Router;
use axum::body::Body;
use axum::body::to_bytes;
use axum::extract::Json;
use axum::extract::State;
use axum::http::HeaderMap;
use axum::http::HeaderValue;
use axum::http::Method;
use axum::http::Request;
use axum::http::StatusCode;
use axum::http::header::AUTHORIZATION;
use axum::http::header::CONTENT_TYPE;
use axum::http::header::HOST;
use axum::http::header::WWW_AUTHENTICATE;
use axum::middleware;
use axum::middleware::Next;
use axum::response::Response;
use axum::routing::get;
use axum::routing::post;
use rmcp::ErrorData as McpError;
use rmcp::handler::server::ServerHandler;
use rmcp::model::CallToolRequestParams;
use rmcp::model::CallToolResult;
use rmcp::model::JsonObject;
use rmcp::model::ListResourceTemplatesResult;
use rmcp::model::ListResourcesResult;
use rmcp::model::ListToolsResult;
use rmcp::model::PaginatedRequestParams;
use rmcp::model::RawResource;
use rmcp::model::RawResourceTemplate;
use rmcp::model::ReadResourceRequestParams;
use rmcp::model::ReadResourceResult;
use rmcp::model::Resource;
use rmcp::model::ResourceContents;
use rmcp::model::ResourceTemplate;
use rmcp::model::ServerCapabilities;
use rmcp::model::ServerInfo;
use rmcp::model::Tool;
use rmcp::model::ToolAnnotations;
use rmcp::transport::StreamableHttpServerConfig;
use rmcp::transport::StreamableHttpService;
use rmcp::transport::streamable_http_server::session::local::LocalSessionManager;
use serde::Deserialize;
use serde_json::Value;
use serde_json::json;
use tokio::sync::Mutex;
use tokio::task;
use tokio::time::sleep;

#[derive(Clone)]
struct TestToolServer {
    tools: Arc<Vec<Tool>>,
    resources: Arc<Vec<Resource>>,
    resource_templates: Arc<Vec<ResourceTemplate>>,
}

const MEMO_URI: &str = "memo://codex/example-note";
const MEMO_CONTENT: &str = "This is a sample MCP resource served by the rmcp test server.";
const MCP_SESSION_ID_HEADER: &str = "mcp-session-id";
const SESSION_POST_FAILURE_CONTROL_PATH: &str = "/test/control/session-post-failure";
const INITIALIZE_POST_FAILURE_CONTROL_PATH: &str = "/test/control/initialize-post-failure";
const INITIALIZED_NOTIFICATION_POST_FAILURE_CONTROL_PATH: &str =
    "/test/control/initialized-notification-post-failure";
const MAX_MCP_POST_BODY_BYTES: usize = 1024 * 1024;

#[derive(Clone, Default)]
struct PostFailureState {
    armed_failure: Arc<Mutex<Option<ArmedFailure>>>,
}

#[derive(Clone, Copy, Debug)]
enum ArmedFailureTarget {
    Initialize,
    InitializedNotification,
    Session,
}

#[derive(Clone, Debug)]
struct ArmedFailure {
    target: ArmedFailureTarget,
    status: StatusCode,
    remaining: usize,
    /// Raw `WWW-Authenticate` challenge header field values returned with the failure.
    www_authenticate_headers: Vec<HeaderValue>,
    content_type: Option<HeaderValue>,
    body: Option<String>,
}

#[derive(Debug, Deserialize)]
struct ArmSessionPostFailureRequest {
    status: u16,
    remaining: usize,
    /// Raw `WWW-Authenticate` challenge header field values to add to the failure.
    #[serde(default)]
    www_authenticate_headers: Vec<String>,
    content_type: Option<String>,
    body: Option<String>,
}

#[derive(Deserialize)]
struct EchoArgs {
    message: String,
    #[allow(dead_code)]
    env_var: Option<String>,
}

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    let bind_addr = parse_bind_addr()?;
    let post_failure_state = PostFailureState::default();
    const MAX_BIND_RETRIES: u32 = 20;
    const BIND_RETRY_DELAY: Duration = Duration::from_millis(50);

    let mut bind_retries = 0;
    let listener = loop {
        match tokio::net::TcpListener::bind(&bind_addr).await {
            Ok(listener) => break listener,
            Err(err) if err.kind() == ErrorKind::PermissionDenied => {
                eprintln!(
                    "failed to bind to {bind_addr}: {err}. make sure the process has network access"
                );
                return Ok(());
            }
            Err(err) if err.kind() == ErrorKind::AddrInUse && bind_retries < MAX_BIND_RETRIES => {
                bind_retries += 1;
                sleep(BIND_RETRY_DELAY).await;
            }
            Err(err) => return Err(err.into()),
        }
    };
    let actual_bind_addr = listener.local_addr()?;
    if let Ok(bound_addr_file) = std::env::var("MCP_STREAMABLE_HTTP_BOUND_ADDR_FILE") {
        fs::write(bound_addr_file, actual_bind_addr.to_string())?;
    }
    eprintln!("starting rmcp streamable http test server on http://{actual_bind_addr}/mcp");

    let router = Router::new()
        .route(
            SESSION_POST_FAILURE_CONTROL_PATH,
            post(arm_session_post_failure),
        )
        .route(
            INITIALIZE_POST_FAILURE_CONTROL_PATH,
            post(arm_initialize_post_failure),
        )
        .route(
            INITIALIZED_NOTIFICATION_POST_FAILURE_CONTROL_PATH,
            post(arm_initialized_notification_post_failure),
        )
        .route(
            "/.well-known/oauth-authorization-server/mcp",
            get({
                move |headers: HeaderMap| async move {
                    let metadata_base = headers
                        .get(HOST)
                        .and_then(|value| value.to_str().ok())
                        .map(|host| format!("http://{host}"))
                        .unwrap_or_else(|| format!("http://{actual_bind_addr}"));
                    #[expect(clippy::expect_used)]
                    Response::builder()
                        .status(StatusCode::OK)
                        .header(CONTENT_TYPE, "application/json")
                        .body(Body::from(
                            serde_json::to_vec(&json!({
                                "authorization_endpoint": format!("{metadata_base}/oauth/authorize"),
                                "token_endpoint": format!("{metadata_base}/oauth/token"),
                                "scopes_supported": [""],
                            })).expect("failed to serialize metadata"),
                        ))
                        .expect("valid metadata response")
                }
            }),
        )
        .nest_service(
            "/mcp",
            StreamableHttpService::new(
                || Ok(TestToolServer::new()),
                Arc::new(LocalSessionManager::default()),
                StreamableHttpServerConfig::default(),
            ),
        )
        .layer(middleware::from_fn_with_state(
            post_failure_state.clone(),
            fail_mcp_post_when_armed,
        ))
        .with_state(post_failure_state);

    let router = if let Ok(token) = std::env::var("MCP_EXPECT_BEARER") {
        let expected = Arc::new(format!("Bearer {token}"));
        router.layer(middleware::from_fn_with_state(expected, require_bearer))
    } else {
        router
    };

    axum::serve(listener, router).await?;
    task::yield_now().await;
    Ok(())
}

impl ServerHandler for TestToolServer {
    fn get_info(&self) -> ServerInfo {
        ServerInfo::new(
            ServerCapabilities::builder()
                .enable_tools()
                .enable_tool_list_changed()
                .enable_resources()
                .build(),
        )
    }

    fn list_tools(
        &self,
        _request: Option<PaginatedRequestParams>,
        _context: rmcp::service::RequestContext<rmcp::service::RoleServer>,
    ) -> impl std::future::Future<Output = Result<ListToolsResult, McpError>> + Send + '_ {
        let tools = self.tools.clone();
        async move {
            Ok(ListToolsResult {
                tools: (*tools).clone(),
                next_cursor: None,
                meta: None,
            })
        }
    }

    fn list_resources(
        &self,
        _request: Option<PaginatedRequestParams>,
        _context: rmcp::service::RequestContext<rmcp::service::RoleServer>,
    ) -> impl std::future::Future<Output = Result<ListResourcesResult, McpError>> + Send + '_ {
        let resources = self.resources.clone();
        async move {
            Ok(ListResourcesResult {
                resources: (*resources).clone(),
                next_cursor: None,
                meta: None,
            })
        }
    }

    async fn list_resource_templates(
        &self,
        _request: Option<PaginatedRequestParams>,
        _context: rmcp::service::RequestContext<rmcp::service::RoleServer>,
    ) -> Result<ListResourceTemplatesResult, McpError> {
        Ok(ListResourceTemplatesResult {
            resource_templates: (*self.resource_templates).clone(),
            next_cursor: None,
            meta: None,
        })
    }

    async fn read_resource(
        &self,
        ReadResourceRequestParams { uri, .. }: ReadResourceRequestParams,
        _context: rmcp::service::RequestContext<rmcp::service::RoleServer>,
    ) -> Result<ReadResourceResult, McpError> {
        if uri == MEMO_URI {
            Ok(ReadResourceResult::new(vec![
                ResourceContents::TextResourceContents {
                    uri,
                    mime_type: Some("text/plain".to_string()),
                    text: Self::memo_text().to_string(),
                    meta: None,
                },
            ]))
        } else {
            Err(McpError::resource_not_found(
                "resource_not_found",
                Some(json!({ "uri": uri })),
            ))
        }
    }

    async fn call_tool(
        &self,
        request: CallToolRequestParams,
        _context: rmcp::service::RequestContext<rmcp::service::RoleServer>,
    ) -> Result<CallToolResult, McpError> {
        match request.name.as_ref() {
            "echo" => {
                let args: EchoArgs = match request.arguments {
                    Some(arguments) => serde_json::from_value(serde_json::Value::Object(
                        arguments.into_iter().collect(),
                    ))
                    .map_err(|err| McpError::invalid_params(err.to_string(), None))?,
                    None => {
                        return Err(McpError::invalid_params(
                            "missing arguments for echo tool",
                            None,
                        ));
                    }
                };

                let env_snapshot: HashMap<String, String> = std::env::vars().collect();
                let structured_content = json!({
                    "echo": format!("ECHOING: {}", args.message),
                    "env": env_snapshot.get("MCP_TEST_VALUE"),
                });

                let mut result = CallToolResult::success(Vec::new());
                result.structured_content = Some(structured_content);
                Ok(result)
            }
            other => Err(McpError::invalid_params(
                format!("unknown tool: {other}"),
                None,
            )),
        }
    }
}

impl TestToolServer {
    fn new() -> Self {
        let tools = vec![Self::echo_tool()];
        let resources = vec![Self::memo_resource()];
        let resource_templates = vec![Self::memo_template()];
        Self {
            tools: Arc::new(tools),
            resources: Arc::new(resources),
            resource_templates: Arc::new(resource_templates),
        }
    }

    fn echo_tool() -> Tool {
        #[expect(clippy::expect_used)]
        let schema: JsonObject = serde_json::from_value(json!({
            "type": "object",
            "properties": {
                "message": { "type": "string" },
                "env_var": { "type": "string" }
            },
            "required": ["message"],
            "additionalProperties": false
        }))
        .expect("echo tool schema should deserialize");

        let mut tool = Tool::new(
            Cow::Borrowed("echo"),
            Cow::Borrowed("Echo back the provided message and include environment data."),
            Arc::new(schema),
        );
        #[expect(clippy::expect_used)]
        let output_schema: JsonObject = serde_json::from_value(json!({
            "type": "object",
            "properties": {
                "echo": { "type": "string" },
                "env": {
                    "anyOf": [
                        { "type": "string" },
                        { "type": "null" }
                    ]
                }
            },
            "required": ["echo", "env"],
            "additionalProperties": false
        }))
        .expect("echo tool output schema should deserialize");
        tool.output_schema = Some(Arc::new(output_schema));
        tool.annotations = Some(ToolAnnotations::new().read_only(true));
        tool
    }

    fn memo_resource() -> Resource {
        let raw = RawResource {
            uri: MEMO_URI.to_string(),
            name: "example-note".to_string(),
            title: Some("Example Note".to_string()),
            description: Some("A sample MCP resource exposed for integration tests.".to_string()),
            mime_type: Some("text/plain".to_string()),
            size: None,
            icons: None,
            meta: None,
        };
        Resource::new(raw, None)
    }

    fn memo_template() -> ResourceTemplate {
        let raw = RawResourceTemplate {
            uri_template: "memo://codex/{slug}".to_string(),
            name: "codex-memo".to_string(),
            title: Some("Codex Memo".to_string()),
            description: Some(
                "Template for memo://codex/{slug} resources used in tests.".to_string(),
            ),
            mime_type: Some("text/plain".to_string()),
            icons: None,
        };
        ResourceTemplate::new(raw, None)
    }

    fn memo_text() -> &'static str {
        MEMO_CONTENT
    }
}

fn parse_bind_addr() -> Result<SocketAddr, Box<dyn std::error::Error>> {
    let default_addr = "127.0.0.1:3920";
    let bind_addr = std::env::var("MCP_STREAMABLE_HTTP_BIND_ADDR")
        .or_else(|_| std::env::var("BIND_ADDR"))
        .unwrap_or_else(|_| default_addr.to_string());
    Ok(bind_addr.parse()?)
}

async fn require_bearer(
    State(expected): State<Arc<String>>,
    request: Request<Body>,
    next: Next,
) -> Result<Response, StatusCode> {
    if request.uri().path().contains("/.well-known/") {
        return Ok(next.run(request).await);
    }
    if request
        .headers()
        .get(AUTHORIZATION)
        .is_some_and(|value| value.as_bytes() == expected.as_bytes())
    {
        Ok(next.run(request).await)
    } else {
        Err(StatusCode::UNAUTHORIZED)
    }
}

async fn arm_session_post_failure(
    State(state): State<PostFailureState>,
    Json(request): Json<ArmSessionPostFailureRequest>,
) -> Result<StatusCode, StatusCode> {
    arm_post_failure(state, request, ArmedFailureTarget::Session).await
}

async fn arm_initialize_post_failure(
    State(state): State<PostFailureState>,
    Json(request): Json<ArmSessionPostFailureRequest>,
) -> Result<StatusCode, StatusCode> {
    arm_post_failure(state, request, ArmedFailureTarget::Initialize).await
}

async fn arm_initialized_notification_post_failure(
    State(state): State<PostFailureState>,
    Json(request): Json<ArmSessionPostFailureRequest>,
) -> Result<StatusCode, StatusCode> {
    arm_post_failure(state, request, ArmedFailureTarget::InitializedNotification).await
}

async fn arm_post_failure(
    state: PostFailureState,
    request: ArmSessionPostFailureRequest,
    target: ArmedFailureTarget,
) -> Result<StatusCode, StatusCode> {
    let status = StatusCode::from_u16(request.status).map_err(|_| StatusCode::BAD_REQUEST)?;
    let www_authenticate_headers = request
        .www_authenticate_headers
        .into_iter()
        .map(|value| HeaderValue::from_str(&value).map_err(|_| StatusCode::BAD_REQUEST))
        .collect::<Result<Vec<_>, _>>()?;
    let content_type = request
        .content_type
        .map(|value| HeaderValue::from_str(&value).map_err(|_| StatusCode::BAD_REQUEST))
        .transpose()?;
    let armed_failure = if request.remaining == 0 {
        None
    } else {
        Some(ArmedFailure {
            target,
            status,
            remaining: request.remaining,
            www_authenticate_headers,
            content_type,
            body: request.body,
        })
    };
    *state.armed_failure.lock().await = armed_failure;
    Ok(StatusCode::NO_CONTENT)
}

async fn fail_mcp_post_when_armed(
    State(state): State<PostFailureState>,
    request: Request<Body>,
    next: Next,
) -> Response {
    if request.uri().path() != "/mcp" || request.method() != Method::POST {
        return next.run(request).await;
    }
    let (parts, body) = request.into_parts();
    let body_bytes = match to_bytes(body, MAX_MCP_POST_BODY_BYTES).await {
        Ok(body_bytes) => body_bytes,
        Err(_) => {
            let mut response = Response::new(Body::from("failed to read request body"));
            *response.status_mut() = StatusCode::BAD_REQUEST;
            return response;
        }
    };
    let has_session_id = parts.headers.contains_key(MCP_SESSION_ID_HEADER);
    let mcp_method = request_mcp_method(&body_bytes);

    {
        let mut armed_failure = state.armed_failure.lock().await;
        if let Some(failure) = armed_failure.as_mut()
            && failure.remaining > 0
            && match failure.target {
                ArmedFailureTarget::Initialize => !has_session_id,
                ArmedFailureTarget::InitializedNotification => {
                    has_session_id && mcp_method.as_deref() == Some("notifications/initialized")
                }
                ArmedFailureTarget::Session => {
                    has_session_id && mcp_method.as_deref() != Some("notifications/initialized")
                }
            }
        {
            failure.remaining -= 1;
            let status = failure.status;
            let www_authenticate_headers = failure.www_authenticate_headers.clone();
            let content_type = failure.content_type.clone();
            let body = failure
                .body
                .clone()
                .unwrap_or_else(|| format!("forced session failure with status {status}"));
            if failure.remaining == 0 {
                *armed_failure = None;
            }
            let mut response = Response::new(Body::from(body));
            *response.status_mut() = status;
            if let Some(content_type) = content_type {
                response.headers_mut().insert(CONTENT_TYPE, content_type);
            }
            for www_authenticate_header in www_authenticate_headers {
                response
                    .headers_mut()
                    .append(WWW_AUTHENTICATE, www_authenticate_header);
            }
            return response;
        }
    }

    next.run(Request::from_parts(parts, Body::from(body_bytes)))
        .await
}

fn request_mcp_method(body: &[u8]) -> Option<String> {
    serde_json::from_slice::<Value>(body)
        .ok()?
        .get("method")?
        .as_str()
        .map(ToString::to_string)
}
