import re

with open("server/src/main.rs", "r", encoding="utf-8") as f:
    content = f.read()

cli_gate_block = re.search(r'    if let Some\(command\) = cli\.command \{.*?(?=    run_server\(state\))', content, re.DOTALL)
if cli_gate_block:
    content = content.replace(cli_gate_block.group(0), "")

new_cli_logic = """
    if let Some(command) = cli.command {
        let rt = tokio::runtime::Runtime::new().unwrap();
        rt.block_on(async {
            let client = reqwest::Client::new();
            match command {
                Commands::Gate { subcmd } => match subcmd {
                    GateCommands::Set { action, target, namespace, params, authorize, block, reason } => {
                        let mut param_map = HashMap::new();
                        for p in params {
                            if let Some((k, v)) = p.split_once('=') {
                                param_map.insert(k.to_string(), v.to_string());
                            }
                        }
                        let body = serde_json::json!({
                            "action": action,
                            "target": target,
                            "namespace": namespace,
                            "params": param_map,
                            "authorize": authorize,
                            "block": block,
                            "reason": reason
                        });
                        match client.post("http://127.0.0.1:3000/gate/set").json(&body).send().await {
                            Ok(res) if res.status().is_success() => {
                                println!("Gate state updated via daemon.");
                                std::process::exit(0);
                            }
                            Ok(res) => {
                                eprintln!("Failed to update gate: {}", res.status());
                                std::process::exit(1);
                            }
                            Err(e) => {
                                eprintln!("Failed to connect to daemon: {}", e);
                                std::process::exit(1);
                            }
                        }
                    }
                    GateCommands::Verify { action, target, namespace, params, consume } => {
                        let mut query = vec![
                            ("action".to_string(), action),
                            ("target".to_string(), target),
                            ("consume".to_string(), consume.to_string()),
                        ];
                        if let Some(ns) = namespace {
                            query.push(("namespace".to_string(), ns));
                        }
                        // Reqwest will serialize params correctly if we pass the right struct.
                        // Actually, axum's Query extractor for HashMap requires flat keys or standard serialization.
                        for p in params {
                            if let Some((k, v)) = p.split_once('=') {
                                query.push((k.to_string(), v.to_string()));
                            }
                        }
                        match client.get("http://127.0.0.1:3000/gate/verify").query(&query).send().await {
                            Ok(res) => {
                                let status = res.status();
                                let text = res.text().await.unwrap_or_default();
                                if status.is_success() {
                                    std::process::exit(0);
                                } else if status == reqwest::StatusCode::FORBIDDEN {
                                    eprintln!("❌ {}", text);
                                    std::process::exit(1);
                                } else {
                                    eprintln!("❌ {}", text);
                                    std::process::exit(2);
                                }
                            }
                            Err(e) => {
                                eprintln!("Failed to connect to daemon: {}", e);
                                std::process::exit(2);
                            }
                        }
                    }
                }
            }
        });
        return Ok(());
    }
"""

restart_block = re.search(r'    if cli\.restart \{.*?return Ok\(\(\);\n    \}', content, re.DOTALL)
if restart_block:
    content = content.replace(restart_block.group(0), restart_block.group(0) + "\n" + new_cli_logic)

with open("server/src/main.rs", "w", encoding="utf-8") as f:
    f.write(content)