use rmcp::{
    ServerHandler,
    handler::server::{router::tool::ToolRouter, tool::Parameters},
    model::*,
    tool, tool_handler, tool_router,
    transport::streamable_http_server::{
        StreamableHttpService, session::local::LocalSessionManager,
    },
};
use std::sync::{Arc, RwLock};

#[derive(Debug, serde::Deserialize, schemars::JsonSchema)]
pub struct ListEntitiesRequest {}

#[derive(Debug, serde::Deserialize, schemars::JsonSchema)]
pub struct QueryEntityRequest {
    #[schemars(description = "Name of the entity to query")]
    pub name: String,
}

#[derive(Debug, serde::Deserialize, schemars::JsonSchema)]
pub struct SpawnEntityRequest {
    #[schemars(description = "Unique name for the entity")]
    pub name: String,
    #[schemars(description = "Mesh type: Cube, Sphere, Cylinder, Cone, Plane, Torus")]
    pub mesh: String,
    #[schemars(description = "Position as [x, y, z]")]
    pub position: [f32; 3],
    #[schemars(description = "Scale as [x, y, z], defaults to [1, 1, 1]")]
    pub scale: Option<[f32; 3]>,
}

#[derive(Debug, serde::Deserialize, schemars::JsonSchema)]
pub struct DespawnEntityRequest {
    #[schemars(description = "Name of the entity to despawn")]
    pub name: String,
}

#[derive(Debug, serde::Deserialize, schemars::JsonSchema)]
pub struct SetPositionRequest {
    #[schemars(description = "Name of the entity")]
    pub name: String,
    #[schemars(description = "New position as [x, y, z]")]
    pub position: [f32; 3],
}

#[derive(Debug, serde::Deserialize, schemars::JsonSchema)]
pub struct SetRotationRequest {
    #[schemars(description = "Name of the entity")]
    pub name: String,
    #[schemars(description = "Rotation as euler angles [pitch, yaw, roll] in radians")]
    pub rotation: [f32; 3],
}

#[derive(Debug, serde::Deserialize, schemars::JsonSchema)]
pub struct SetScaleRequest {
    #[schemars(description = "Name of the entity")]
    pub name: String,
    #[schemars(description = "New scale as [x, y, z]")]
    pub scale: [f32; 3],
}

#[derive(Debug, serde::Deserialize, schemars::JsonSchema)]
pub struct SetMaterialColorRequest {
    #[schemars(description = "Name of the entity")]
    pub name: String,
    #[schemars(description = "Base color as [r, g, b, a] where each component is 0.0-1.0")]
    pub color: [f32; 4],
}

#[derive(Debug, serde::Deserialize, schemars::JsonSchema)]
pub struct SetBehaviorRequest {
    #[schemars(description = "Name of the entity")]
    pub name: String,
    #[schemars(
        description = "Behavior flags as a u64 bitmask. Engine behaviors: ROTATE=1, BOB=2, ORBIT=4, SCALE_PULSE=8, SINE_WAVE=16"
    )]
    pub flags: u64,
}

#[derive(Debug, serde::Deserialize, schemars::JsonSchema)]
pub struct AddBehaviorRequest {
    #[schemars(description = "Name of the entity")]
    pub name: String,
    #[schemars(
        description = "Behavior to add: 'rotate', 'bob', 'orbit', 'scale_pulse', 'sine_wave'"
    )]
    pub behavior: String,
}

#[derive(Debug, serde::Deserialize, schemars::JsonSchema)]
pub struct RemoveBehaviorRequest {
    #[schemars(description = "Name of the entity")]
    pub name: String,
    #[schemars(
        description = "Behavior to remove: 'rotate', 'bob', 'orbit', 'scale_pulse', 'sine_wave'"
    )]
    pub behavior: String,
}

#[derive(Debug, serde::Deserialize, schemars::JsonSchema)]
pub struct GetBehaviorRequest {
    #[schemars(description = "Name of the entity")]
    pub name: String,
}

#[derive(Debug, serde::Deserialize, schemars::JsonSchema)]
pub struct LoadAssetRequest {
    #[schemars(description = "Absolute file path to .glb, .gltf, or .fbx file")]
    pub path: String,
    #[schemars(description = "Unique name to reference this asset")]
    pub asset_name: String,
}

#[derive(Debug, serde::Deserialize, schemars::JsonSchema)]
pub struct SpawnPrefabRequest {
    #[schemars(description = "Name of a previously loaded asset")]
    pub asset_name: String,
    #[schemars(description = "Unique name for the spawned entity")]
    pub entity_name: String,
    #[schemars(description = "Position as [x, y, z]")]
    pub position: [f32; 3],
    #[schemars(description = "Scale as [x, y, z], defaults to [1, 1, 1]")]
    pub scale: Option<[f32; 3]>,
}

#[derive(Debug, serde::Deserialize, schemars::JsonSchema)]
pub struct ListLoadedAssetsRequest {}

#[derive(Debug, serde::Deserialize, schemars::JsonSchema)]
pub struct ClearSceneRequest {}

#[derive(Debug, Clone, serde::Serialize, serde::Deserialize, schemars::JsonSchema)]
pub struct BatchOperation {
    #[schemars(
        description = "Operation type: spawn_entity, despawn_entity, set_position, set_rotation, set_scale, set_material_color, add_behavior, remove_behavior, set_behavior, get_behavior, query_entity, list_entities, clear_scene"
    )]
    pub operation: String,
    #[schemars(description = "Entity name for the operation")]
    pub name: Option<String>,
    #[schemars(
        description = "Mesh type for spawn_entity: Cube, Sphere, Cylinder, Cone, Plane, Torus"
    )]
    pub mesh: Option<String>,
    #[schemars(description = "Position as [x, y, z]")]
    pub position: Option<[f32; 3]>,
    #[schemars(description = "Scale as [x, y, z]")]
    pub scale: Option<[f32; 3]>,
    #[schemars(description = "Rotation as euler angles [pitch, yaw, roll] in radians")]
    pub rotation: Option<[f32; 3]>,
    #[schemars(description = "Color as [r, g, b, a] where each component is 0.0-1.0")]
    pub color: Option<[f32; 4]>,
    #[schemars(description = "Behavior name: rotate, bob, orbit, scale_pulse, sine_wave")]
    pub behavior: Option<String>,
    #[schemars(description = "Behavior flags as u64 bitmask")]
    pub flags: Option<u64>,
}

#[derive(Debug, serde::Deserialize, schemars::JsonSchema)]
pub struct BatchRequest {
    #[schemars(description = "Array of operations to execute atomically in a single frame")]
    pub operations: Vec<BatchOperation>,
}

#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub enum McpCommand {
    ListEntities,
    QueryEntity {
        name: String,
    },
    SpawnEntity {
        name: String,
        mesh: String,
        position: [f32; 3],
        scale: [f32; 3],
    },
    DespawnEntity {
        name: String,
    },
    SetPosition {
        name: String,
        position: [f32; 3],
    },
    SetRotation {
        name: String,
        rotation: [f32; 3],
    },
    SetScale {
        name: String,
        scale: [f32; 3],
    },
    SetMaterialColor {
        name: String,
        color: [f32; 4],
    },
    SetBehavior {
        name: String,
        flags: u64,
    },
    AddBehavior {
        name: String,
        behavior: String,
    },
    RemoveBehavior {
        name: String,
        behavior: String,
    },
    GetBehavior {
        name: String,
    },
    LoadAsset {
        path: String,
        asset_name: String,
    },
    SpawnPrefab {
        asset_name: String,
        entity_name: String,
        position: [f32; 3],
        scale: [f32; 3],
    },
    ListLoadedAssets,
    ClearScene,
    Batch(Vec<BatchOperation>),
}

#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub enum McpResponse {
    EntityList(Vec<String>),
    EntityInfo {
        name: String,
        position: [f32; 3],
        rotation: [f32; 4],
        scale: [f32; 3],
    },
    BehaviorInfo {
        name: String,
        flags: u64,
        active_behaviors: Vec<String>,
    },
    AssetList(Vec<String>),
    Success(String),
    Error(String),
    BatchResult(Vec<McpResponse>),
}

pub type CommandQueue = Arc<RwLock<Vec<McpCommand>>>;
pub type ResponseQueue = Arc<RwLock<Option<McpResponse>>>;

#[derive(Clone)]
pub struct McpServer {
    tool_router: ToolRouter<Self>,
    command_queue: CommandQueue,
    response_queue: ResponseQueue,
}

#[tool_router]
impl McpServer {
    pub fn new(command_queue: CommandQueue, response_queue: ResponseQueue) -> Self {
        Self {
            tool_router: Self::tool_router(),
            command_queue,
            response_queue,
        }
    }

    fn send_command_and_wait(&self, cmd: McpCommand) -> String {
        {
            let mut queue = self.command_queue.write().unwrap();
            queue.push(cmd);
        }

        for _ in 0..100 {
            std::thread::sleep(std::time::Duration::from_millis(50));
            let mut response = self.response_queue.write().unwrap();
            if let Some(resp) = response.take() {
                return serde_json::to_string_pretty(&resp)
                    .unwrap_or_else(|_| format!("{:?}", resp));
            }
        }

        "Timeout waiting for response".to_string()
    }

    #[tool(description = "List all named entities in the scene")]
    async fn list_entities(&self, Parameters(_): Parameters<ListEntitiesRequest>) -> String {
        self.send_command_and_wait(McpCommand::ListEntities)
    }

    #[tool(description = "Query information about a specific entity by name")]
    async fn query_entity(&self, Parameters(req): Parameters<QueryEntityRequest>) -> String {
        self.send_command_and_wait(McpCommand::QueryEntity { name: req.name })
    }

    #[tool(description = "Spawn a new entity with a mesh at a position")]
    async fn spawn_entity(&self, Parameters(req): Parameters<SpawnEntityRequest>) -> String {
        self.send_command_and_wait(McpCommand::SpawnEntity {
            name: req.name,
            mesh: req.mesh,
            position: req.position,
            scale: req.scale.unwrap_or([1.0, 1.0, 1.0]),
        })
    }

    #[tool(description = "Despawn an entity by name")]
    async fn despawn_entity(&self, Parameters(req): Parameters<DespawnEntityRequest>) -> String {
        self.send_command_and_wait(McpCommand::DespawnEntity { name: req.name })
    }

    #[tool(description = "Set the position of an entity")]
    async fn set_position(&self, Parameters(req): Parameters<SetPositionRequest>) -> String {
        self.send_command_and_wait(McpCommand::SetPosition {
            name: req.name,
            position: req.position,
        })
    }

    #[tool(description = "Set the rotation of an entity using euler angles in radians")]
    async fn set_rotation(&self, Parameters(req): Parameters<SetRotationRequest>) -> String {
        self.send_command_and_wait(McpCommand::SetRotation {
            name: req.name,
            rotation: req.rotation,
        })
    }

    #[tool(description = "Set the scale of an entity")]
    async fn set_scale(&self, Parameters(req): Parameters<SetScaleRequest>) -> String {
        self.send_command_and_wait(McpCommand::SetScale {
            name: req.name,
            scale: req.scale,
        })
    }

    #[tool(description = "Set the material base color of an entity")]
    async fn set_material_color(
        &self,
        Parameters(req): Parameters<SetMaterialColorRequest>,
    ) -> String {
        self.send_command_and_wait(McpCommand::SetMaterialColor {
            name: req.name,
            color: req.color,
        })
    }

    #[tool(
        description = "Set behavior flags on an entity (u64 bitmask). ROTATE=1, BOB=2, ORBIT=4, SCALE_PULSE=8, SINE_WAVE=16"
    )]
    async fn set_behavior(&self, Parameters(req): Parameters<SetBehaviorRequest>) -> String {
        self.send_command_and_wait(McpCommand::SetBehavior {
            name: req.name,
            flags: req.flags,
        })
    }

    #[tool(
        description = "Add a behavior to an entity: 'rotate', 'bob', 'orbit', 'scale_pulse', 'sine_wave'"
    )]
    async fn add_behavior(&self, Parameters(req): Parameters<AddBehaviorRequest>) -> String {
        self.send_command_and_wait(McpCommand::AddBehavior {
            name: req.name,
            behavior: req.behavior,
        })
    }

    #[tool(
        description = "Remove a behavior from an entity: 'rotate', 'bob', 'orbit', 'scale_pulse', 'sine_wave'"
    )]
    async fn remove_behavior(&self, Parameters(req): Parameters<RemoveBehaviorRequest>) -> String {
        self.send_command_and_wait(McpCommand::RemoveBehavior {
            name: req.name,
            behavior: req.behavior,
        })
    }

    #[tool(description = "Get the current behavior flags for an entity")]
    async fn get_behavior(&self, Parameters(req): Parameters<GetBehaviorRequest>) -> String {
        self.send_command_and_wait(McpCommand::GetBehavior { name: req.name })
    }

    #[tool(description = "Load a 3D asset from file (.glb, .gltf, or .fbx)")]
    async fn load_asset(&self, Parameters(req): Parameters<LoadAssetRequest>) -> String {
        self.send_command_and_wait(McpCommand::LoadAsset {
            path: req.path,
            asset_name: req.asset_name,
        })
    }

    #[tool(description = "Spawn a prefab from a previously loaded asset")]
    async fn spawn_prefab(&self, Parameters(req): Parameters<SpawnPrefabRequest>) -> String {
        self.send_command_and_wait(McpCommand::SpawnPrefab {
            asset_name: req.asset_name,
            entity_name: req.entity_name,
            position: req.position,
            scale: req.scale.unwrap_or([1.0, 1.0, 1.0]),
        })
    }

    #[tool(description = "List all loaded assets available for spawning")]
    async fn list_loaded_assets(
        &self,
        Parameters(_): Parameters<ListLoadedAssetsRequest>,
    ) -> String {
        self.send_command_and_wait(McpCommand::ListLoadedAssets)
    }

    #[tool(description = "Clear all named entities from the scene")]
    async fn clear_scene(&self, Parameters(_): Parameters<ClearSceneRequest>) -> String {
        self.send_command_and_wait(McpCommand::ClearScene)
    }

    #[tool(
        description = "Execute multiple operations atomically in a single frame. Supported operations: spawn_entity, despawn_entity, set_position, set_rotation, set_scale, set_material_color, add_behavior, remove_behavior, set_behavior, get_behavior, query_entity, list_entities, clear_scene. Each operation requires 'name' (except list_entities and clear_scene) and operation-specific fields."
    )]
    async fn batch(&self, Parameters(req): Parameters<BatchRequest>) -> String {
        self.send_command_and_wait(McpCommand::Batch(req.operations))
    }
}

#[tool_handler]
impl ServerHandler for McpServer {
    fn get_info(&self) -> ServerInfo {
        ServerInfo {
            instructions: Some(
                "Nightshade Engine MCP Server - Control entities and scene elements".into(),
            ),
            capabilities: ServerCapabilities::builder().enable_tools().build(),
            ..Default::default()
        }
    }
}

pub fn start_mcp_server(command_queue: CommandQueue, response_queue: ResponseQueue) {
    std::thread::spawn(move || {
        let runtime = tokio::runtime::Runtime::new().unwrap();
        runtime.block_on(async {
            let command_queue_clone = command_queue.clone();
            let response_queue_clone = response_queue.clone();

            let service = StreamableHttpService::new(
                move || Ok(McpServer::new(command_queue_clone.clone(), response_queue_clone.clone())),
                LocalSessionManager::default().into(),
                Default::default(),
            );

            let router = axum::Router::new().nest_service("/mcp", service);
            let tcp_listener = tokio::net::TcpListener::bind("127.0.0.1:3333").await.unwrap();

            tracing::info!("MCP server listening on http://127.0.0.1:3333/mcp");
            tracing::info!("Add to Claude Code: claude mcp add --transport http nightshade http://127.0.0.1:3333/mcp");

            axum::serve(tcp_listener, router)
                .with_graceful_shutdown(async {
                    tokio::signal::ctrl_c().await.ok();
                })
                .await
                .ok();
        });
    });
}

pub fn create_mcp_queues() -> (CommandQueue, ResponseQueue) {
    (
        Arc::new(RwLock::new(Vec::new())),
        Arc::new(RwLock::new(None)),
    )
}
