/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

//! Tests cases that are more relavant for server use-cases.
//!
//! The server use-cases are different from the client in a few ways:
//! - The server has a non-lazy `IdMap`.
//! - The server might store `IdMap` in SQL without using transactions (because
//!   of performance concerns and size limits). That could cause `IdMap` having
//!   more entries than `IdDag`. The client uses `MultiLog` in indexedlog to
//!   ensure `IdMap` and `IdDag` are in sync.

use crate::idmap::IdMapWrite;
use crate::ops::DagAddHeads;
use crate::tests::DrawDag;
use crate::tests::TestDag;
use crate::Group;
use crate::Id;
use crate::VertexListWithOptions;

#[tokio::test]
async fn test_idmap_more_entries_than_iddag() {
    let mut dag = TestDag::draw("A--B  # master: B");

    // Make IdMap contain more entries than IdDag.
    let map = &mut dag.dag.map;
    IdMapWrite::insert(map, Id(5), b"C").await.unwrap();
    IdMapWrite::insert(map, Id(6), b"D").await.unwrap();
    IdMapWrite::insert(map, Id(8), b"E").await.unwrap();
    IdMapWrite::insert(map, Id(20), b"H").await.unwrap();
    IdMapWrite::insert(map, Id(21), b"I").await.unwrap();

    let draw = DrawDag::from(
        r#"
           B--C--D--F--G  H--I--J--K
                  /
                 E"#,
    );
    let heads = VertexListWithOptions::from(&["G".into(), "K".into()][..])
        .with_highest_group(Group::MASTER);
    dag.dag.add_heads(&draw, &heads).await.unwrap();

    // C, D, E, H, I are inserted to the IdDag respecting their IdMap values.
    assert_eq!(
        dag.debug_segments(0, Group::MASTER),
        r#"
        H+20 : K+23 [] Root
        F+9 : G+10 [D+6, E+8]
        E+8 : E+8 [] Root
        C+5 : D+6 [B+1]
        A+0 : B+1 [] Root OnlyHead"#
    );
}

#[tokio::test]
async fn test_idmap_more_entries_conflict_with_assign_head() {
    let mut dag = TestDag::draw("A--B  # master: B");

    // Make IdMap contain entry that conflicts with new Id assignment.
    let map = &mut dag.dag.map;
    IdMapWrite::insert(map, Id(2), b"C").await.unwrap();

    let draw = DrawDag::from("B--D");
    let heads = VertexListWithOptions::from(&["D".into()][..]).with_highest_group(Group::MASTER);
    dag.dag.add_heads(&draw, &heads).await.unwrap();

    // D takes id 3, skips the taken id 2.
    assert_eq!(
        dag.debug_segments(0, Group::MASTER),
        "\n        D+3 : D+3 [B+1]\n        A+0 : B+1 [] Root OnlyHead"
    );
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

use std::sync::Arc;

use futures::TryStreamExt;

use super::ProtocolMonitor;
use super::TestDag;
use crate::ops::DagAddHeads;
use crate::ops::DagAlgorithm;
use crate::ops::DagExportPullData;
use crate::ops::DagImportPullData;
use crate::ops::DagPersistent;
use crate::ops::IdConvert;
use crate::Group;
use crate::Id;
use crate::VertexListWithOptions;
use crate::VertexName;

#[tokio::test]
async fn test_sparse_dag() {
    // In this test, we have 3 dags:
    // - server1: A complete dag. Emulates server-side.
    // - server2: An isomorphism of server1. But uses different ids.
    // - client: A sparse dag. Emulates client-side. Cloned from server2.
    //   Speaks remote protocol with server1 or server2.
    let mut server1 = TestDag::new();
    server1.drawdag(
        r#"
        A----B----C----D----E----M----X
         \       /         /
          F-----G--H--I---J
                         /
                     K--L"#,
        &["M"],
    );

    // server2 is an isomorphism of server1 with different ids.
    let mut server2 = TestDag::new();
    server2.drawdag("A-F-G-H-I", &["I"]);
    server2.drawdag("K-L-J I-J", &["J"]);
    server2.drawdag("A-B-C G-C", &["C"]);
    server2.drawdag("C-D-E-M-X J-E", &["M"]);

    for opt_remote_dag in vec![Some(server1.dag), None] {
        let mut client = server2.client_cloned_data().await;

        // Note: some ids (ex. 11) does not have matching name in its IdMap.
        // The server-side non-master (X) is not cloned.
        assert_eq!(
            format!("{:?}", &client.dag),
            r#"Max Level: 0
 Level 0
  Group Master:
   Segments: 6
    11 : M+12 [D+10, J+7] OnlyHead
    9 : D+10 [B+8, G+2]
    B+8 : B+8 [0]
    J+7 : J+7 [I+4, L+6] OnlyHead
    5 : L+6 [] Root
    0 : I+4 [] Root OnlyHead
  Group Non-Master:
   Segments: 0
"#
        );

        // The remote protocol could be backed by either server1 or server2.
        if let Some(remote_dag) = opt_remote_dag {
            let protocol = ProtocolMonitor {
                inner: Box::new(remote_dag),
                output: client.output.clone(),
            };
            client.dag.set_remote_protocol(Arc::new(protocol));
        }

        // With remote protocol. Be able to resolve id <-> names.
        assert_eq!(client.dag.vertex_name(Id(9)).await.unwrap(), "C".into());
        assert_eq!(client.dag.vertex_id("E".into()).await.unwrap(), Id(11));

        // NameSet iteration works too, and resolve Ids in batch.
        let all: Vec<VertexName> = {
            let all = client.dag.all().await.unwrap();
            let iter = all.iter().await.unwrap();
            iter.try_collect().await.unwrap()
        };
        assert_eq!(
            format!("{:?}", all),
            "[M, E, D, C, B, J, L, K, I, H, G, F, A]"
        );

        assert_eq!(
            client.output(),
            [
                "resolve paths: [D~1]",
                "resolve names: [E], heads: [M]",
                "resolve paths: [L~1, I~1, G~1(+2)]"
            ]
        );
    }
}

#[tokio::test]
async fn test_lazy_hash_on_non_master_group() {
    // In this test, we have 3 dags:
    // - server1: A dag.
    // - server2: A dag similar to server1 but in non-master group.
    // - client: A sparse dag. Emulates client-side. Cloned from server1.
    //   Speaks remote protocol with server2.
    let ascii_graph = r#"
        A--B--C--D--I
                  /
        E--F--G--H
        "#;
    let mut server1 = TestDag::new();
    server1.drawdag(ascii_graph, &["I"]);

    // server2 has a same graph but in the non-master group.
    let server2 = TestDag::draw(ascii_graph);

    // server2 can answer name->location queries.
    let client = server1.client_cloned_data().await.with_remote(&server2);
    assert_eq!(client.dag.vertex_id("C".into()).await.unwrap(), Id(2));
    assert_eq!(client.dag.vertex_id("F".into()).await.unwrap(), Id(5));
}

#[tokio::test]
async fn test_negative_cache() {
    let server = TestDag::draw("A-B  # master: B");

    let mut client = server.client_cloned_data().await;

    // Lookup "C" - not found.
    assert!(client.dag.vertex_id("C".into()).await.is_err());
    assert_eq!(client.output(), ["resolve names: [C], heads: [B]"]);

    // Lookup again - no need to resolve again.
    assert!(client.dag.vertex_id("C".into()).await.is_err());
    assert_eq!(client.output(), Vec::<String>::new());

    // The negative cache does not affect inserting the name.
    client.drawdag("B-C-D", &[]);
    assert!(client.dag.vertex_id("C".into()).await.is_ok());
}

#[tokio::test]
async fn test_add_heads() {
    let server = TestDag::draw("A-B  # master: B");
    let mut client = server.client_cloned_data().await;

    let pending = TestDag::draw("A-C B-C-D-E-F-G E-H-K I-J-K");
    let parents = pending.dag.dag_snapshot().unwrap();
    client
        .dag
        .add_heads(
            &parents,
            &vec![VertexName::from("G"), VertexName::from("K")].into(),
        )
        .await
        .unwrap();
    assert_eq!(
        client.output(),
        [
            "resolve names: [I, A], heads: [B]",
            "resolve names: [C], heads: [B]"
        ]
    );

    client.flush("G").await;
    assert_eq!(client.output(), ["resolve names: [I], heads: [B]"]);

    let mut client = server.client_cloned_data().await;
    let heads = VertexListWithOptions::from(&["K".into()][..])
        .with_highest_group(Group::MASTER)
        .chain(&["G".into()][..]);
    client
        .dag
        .add_heads_and_flush(&parents, &heads)
        .await
        .unwrap();
    assert_eq!(
        client.output(),
        [
            "resolve names: [I, A], heads: [B]",
            "resolve names: [C], heads: [B]"
        ]
    );
}

#[tokio::test]
async fn test_basic_pull() {
    let server = TestDag::draw("A-B-C-D  # master: D");
    let mut client = server.client().await;
    client.drawdag("A-B", &["B"]);

    let missing = server.dag.only("D".into(), "B".into()).await.unwrap();
    let pull_data = server.dag.export_pull_data(&missing).await.unwrap();

    let heads = VertexListWithOptions::from(&["D".into()][..]).with_highest_group(Group::MASTER);
    client
        .dag
        .import_pull_data(pull_data, &heads)
        .await
        .unwrap();

    #[cfg(feature = "render")]
    assert_eq!(server.render_graph(), client.render_graph());
}

#[tokio::test]
async fn test_pull_remap() {
    // In this test client and server going to have different IDs, but isomorphic graphs
    let mut server = TestDag::new();
    server.drawdag(
        r#"
        A-B--C--D
           \   /
            F-G
    "#,
        &["D"],
    );
    server.drawdag("B-E", &["E"]);
    let mut client = server.client().await;
    client.drawdag("A-B-E", &["E"]);

    client.pull_ff_master(&server, "E", "D").await.unwrap();

    assert_eq!(
        client.output(),
        [
            "resolve names: [A], heads: []",
            "resolve names: [C, F], heads: [E]"
        ]
    );

    #[cfg(feature = "render")]
    assert_eq!(
        client.render_graph(),
        "
            D    6
            â”œâ”€â•®
            â”‚ G  5
            â”‚ â”‚
            â”‚ F  4
            â”‚ â”‚
            C â”‚  3
            â”œâ”€â•¯
            â”‚ E  2
            â”œâ”€â•¯
            B  1
            â”‚
            A  0"
    );

    #[cfg(feature = "render")]
    assert_eq!(
        server.render_graph(),
        "
            E  6
            â”‚
            â”‚ D    5
            â”‚ â”œâ”€â•®
            â”‚ â”‚ G  4
            â”‚ â”‚ â”‚
            â”‚ â”‚ F  3
            â”œâ”€â”€â”€â•¯
            â”‚ C  2
            â”œâ”€â•¯
            B  1
            â”‚
            A  0"
    );
}

#[tokio::test]
async fn test_pull_overlap() {
    let mut server = TestDag::new();
    server.drawdag("A-B-C-D-E-F", &["F"]);
    let mut client = server.client().await;
    client.drawdag("A", &["A"]);

    client.pull_ff_master(&server, "A", "D").await.unwrap();

    let e = client.pull_ff_master(&server, "B", "F").await.unwrap_err();
    assert_eq!(e.to_string(), "NeedSlowPath: C exists in local graph");

    #[cfg(feature = "render")]
    assert_eq!(
        client.render_graph(),
        r#"
            D  3
            â”‚
            C  2
            â”‚
            B  1
            â”‚
            A  0"#
    );
}

#[tokio::test]
async fn test_pull_lazy_with_merges() {
    // Test fast-forward pull on a lazy graph with merges.
    let mut server = TestDag::new();

    // Initial state. Both client and server has just one vertex.
    server.drawdag("A", &["A"]);
    let mut client = server.client_cloned_data().await;

    // Take some IDs so IDs are different from client graph.
    server.drawdag("X-Y", &["Y"]);

    // Linear fast-forward. The client has a lazy graph.
    server.drawdag("A-B-C-D-E", &["E"]);
    client.pull_ff_master(&server, "A", "E").await.unwrap();
    assert_eq!(client.output(), ["resolve names: [B], heads: [A]"]);

    // C, D are lazy, but E is not.
    assert!(!client.contains_vertex_locally("C"));
    assert!(!client.contains_vertex_locally("D"));
    assert!(client.contains_vertex_locally("E"));

    // Add merges. Test parents remap.
    server.drawdag(
        r#"
                  D
                   \
        C E-F---G-H-I-J-K
         \     /   /
          L-M-N   M
        "#,
        &["K"],
    );
    assert_eq!(
        server.debug_segments(0, Group::MASTER),
        r#"
        I+13 : K+15 [D+5, H+12, M+8]
        G+11 : H+12 [F+9, N+10]
        N+10 : N+10 [M+8]
        F+9 : F+9 [E+6]
        L+7 : M+8 [C+4]
        B+3 : E+6 [A+0]
        X+1 : Y+2 [] Root
        A+0 : A+0 [] Root OnlyHead"#
    );

    client.pull_ff_master(&server, "E", "K").await.unwrap();
    assert_eq!(
        client.debug_segments(0, Group::MASTER),
        r#"
        I+11 : K+13 [D+3, H+10, M+7] OnlyHead
        G+9 : H+10 [F+5, N+8] OnlyHead
        L+6 : N+8 [C+2]
        A+0 : F+5 [] Root OnlyHead"#
    );
    assert_eq!(client.output(), ["resolve names: [C, D, F, L], heads: [E]"]);
}

#[tokio::test]
async fn test_pull_no_pending_changes() {
    let mut server = TestDag::draw("A # master: A");
    let mut client = server.client_cloned_data().await;
    server.drawdag("A-B-C", &["C"]);
    client.drawdag("A-D", &[]);
    let e = client.pull_ff_master(&server, "A", "C").await.unwrap_err();
    assert_eq!(
        e.to_string(),
        "ProgrammingError: import_pull_data called with pending heads ([D])"
    );
}

#[tokio::test]
async fn test_flush_reassign_master() {
    // Test remote calls when flush() causes id reassignment
    // from non-master to master.

    let mut server = TestDag::draw("A-B-C-D-E # master: E");
    let mut client = server.client_cloned_data().await;

    // Add vertexes in the non-master group.
    // There are 2 parts: ::G will be reassigned to the master group.
    // The rest (H, ::L) will remain in the non-master group.
    client.drawdag("B-X-Y-Z-F D-F-G-H B-I-J-K-L", &[]);
    assert_eq!(
        client.output(),
        [
            "resolve names: [B, D], heads: [E]",
            "resolve names: [I, X], heads: [E]",
            "resolve paths: [E~2, E~4]"
        ]
    );

    client.flush("").await;
    assert!(client.output().is_empty());

    // The server needs to have the new master group vertexes (up to G)
    // for the client to be able to assign them in the master group.
    server.drawdag("B-X-Y-Z-F D-F-G", &["G"]);
    assert!(client.output().is_empty());
    client.set_remote(&server);

    // To avoid reusing caches, reopen the graph from disk.
    client.reopen();

    assert_eq!(
        client.dump_state().await,
        r#"<spans [X:L+N0:N9, A:E+0:4]>
Lv0: RH0-4[], N0-N2[1], N3-N5[3, N2], N6-N9[1]
Lv1: N0-N5[1, 3]
P->C: 1->N0, 1->N6, 3->N3, N2->N3
0->A 1->B 2->C 3->D 4->E N0->X N1->Y N2->Z N3->F N4->G N5->H N6->I N7->J N8->K N9->L"#
    );

    // Force reassign of vertexes in the non-master group.
    client.flush("G").await;
    assert!(client.output().is_empty());

    //       I-J-K-L
    //      /
    //   A-B--C--D-E
    //      \     \
    //       X-Y-Z-F-G-H
    //
    // Check that:
    // - G % E = G+F+Z+Y+X are added to the master group.
    // - I::L remains in the non-master group unchanged.
    // - X::G are removed from the non-master group.
    // - H is re-inserted in the non-master group.
    assert_eq!(
        client.dump_state().await,
        r#"<spans [I:L+N6:N9, H+N0, A:G+0:9]>
Lv0: RH0-4[], 5-7[1], 8-9[3, 7], N0-N0[9], N6-N9[1]
Lv1: R0-4[]
P->C: 1->5, 1->N6, 3->8, 7->8, 9->N0
0->A 1->B 2->C 3->D 4->E 5->X 6->Y 7->Z 8->F 9->G N0->H N6->I N7->J N8->K N9->L"#
    );
}

#[tokio::test]
async fn test_resolve_misleading_merges() {
    // Test when the server graph gets more merges making vertexes
    // previously not a parent of a merge become a parent of a merge.

    // Initially, B is not a parent of a merge. A can be resolved as
    // C~2.
    let server1 = TestDag::draw("A-B-C # master: C");

    // Make B a parent of a merge by adding extra vertexes.
    // So A might be resolved as B~1 or C~2.
    let server2 = TestDag::draw(
        r#"
        A-B-C
           \
          D-E   # master: E C"#,
    );

    let client1 = server1.client_cloned_data().await.with_remote(&server2);

    // A is resolved to C~2, not B~1, since in the ancestors(C) sub-graph,
    // B is not a parent of a merge.
    client1.dag.vertex_id("A".into()).await.unwrap();
    assert_eq!(client1.output(), ["resolve names: [A], heads: [C]"]);
}

#[tokio::test]
async fn test_resolve_pick_path() {
    // Test when a vertex can be resolved in multiple ways, pick the one that is
    // supported by the client.

    // Make A resolve to D~1, and C~2.
    let server2 = TestDag::draw(
        r#"
           A-B-C
            \
             \ E
              \ \
               D-F
                \
                 G-H-I   # master: F C I"#,
    );

    // The client does not have D:: part of the graph.
    let server1 = TestDag::draw("A-B-C # master: C");
    let client1 = server1.client_cloned_data().await.with_remote(&server2);

    // Can resolve A using B~1, not D~1.
    assert_eq!(client1.dag.vertex_id("A".into()).await.unwrap(), Id(0));
    assert_eq!(client1.output(), ["resolve names: [A], heads: [C]"]);

    // The client does not have B::, E:: part of the graph.
    // D is not considered as a parent of a merge in the graph.
    let server2 = TestDag::draw("A-D-G-H # master: H");
    let client2 = server2.client_cloned_data().await.with_remote(&server2);

    // Can resolve A using H~3, not B~1 or D~1 as requested.
    assert_eq!(client2.dag.vertex_id("A".into()).await.unwrap(), Id(0));
    assert_eq!(client2.output(), ["resolve names: [A], heads: [H]"]);
}

#[tokio::test]
async fn test_resolve_mixed_result() {
    // Test that Ok and Err can be both present in vertex_id_batch return value.
    let servers: Vec<_> = ["F J", "D J", "C E D F J"]
        .iter()
        .map(|master| {
            TestDag::draw(&format!(
                r#"
                    A-B-C-D-G-H-I-J
                       \   /
                        E-F   # master: {}"#,
                master
            ))
        })
        .collect();

    assert_eq!(
        servers
            .iter()
            .map(|s| s.debug_segments(0, Group::MASTER))
            .collect::<Vec<_>>()
            .join("\n"),
        r#"
        G+6 : J+9 [D+5, F+3] OnlyHead
        C+4 : D+5 [B+1]
        A+0 : F+3 [] Root OnlyHead

        G+6 : J+9 [D+3, F+5] OnlyHead
        E+4 : F+5 [B+1]
        A+0 : D+3 [] Root OnlyHead

        G+6 : J+9 [D+4, F+5] OnlyHead
        F+5 : F+5 [E+3]
        D+4 : D+4 [C+2]
        E+3 : E+3 [B+1]
        A+0 : C+2 [] Root OnlyHead"#
    );

    let names: Vec<_> = "A B C D E F G H I J X".split(' ').map(Into::into).collect();
    for server in servers {
        let client = TestDag::draw("A-B-C-D-G-H B-E-F-G # master: H")
            .client_cloned_data()
            .await
            .with_remote(&server);
        let ids = client.dag.vertex_id_batch(&names).await;
        assert_eq!(
            format!("{:?}", ids),
            "Ok([Ok(0), Ok(1), Ok(2), Ok(3), Ok(4), Ok(5), Ok(6), Ok(7), Err(VertexNotFound(I)), Err(VertexNotFound(J)), Err(VertexNotFound(X))])",
        );
        assert_eq!(
            client.output(),
            ["resolve names: [A, B, C, E, G, I, J, X], heads: [H]"]
        );
    }
}

#[tokio::test]
async fn test_flush_lazy_vertex() {
    // Test flushing with main vertex set to a lazy vertex.
    let server = TestDag::draw("A-B-C-D # master: D");
    let mut client = server.client_cloned_data().await;
    client.flush("B").await;
}

async fn client_for_local_cache_test() -> TestDag {
    let server = TestDag::draw("A-B-C-D-E-F-G # master: G");
    server.client_cloned_data().await
}

async fn check_local_cache(client: &TestDag, v: VertexName, id: Id) {
    // Try looking up vertex using different APIs.
    assert_eq!(client.dag.vertex_id(v.clone()).await.unwrap(), id);
    assert!(client.output().is_empty());

    assert!(client.dag.contains_vertex_name(&v.clone()).await.unwrap());
    assert!(client.output().is_empty());

    assert_eq!(
        client.dag.vertex_id_optional(&v.clone()).await.unwrap(),
        Some(id)
    );
    assert!(client.output().is_empty());

    assert_eq!(
        client
            .dag
            .vertex_id_with_max_group(&v.clone(), Group::MASTER)
            .await
            .unwrap(),
        Some(id)
    );
    assert!(client.output().is_empty());

    assert!(matches!(
        &client
            .dag
            .contains_vertex_name_locally(&[v.clone()])
            .await
            .unwrap()[..],
        [true]
    ));
    assert!(client.output().is_empty());

    assert!(matches!(
        &client.dag.vertex_id_batch(&[v.clone()]).await.unwrap()[..],
        [Ok(i)] if *i == id
    ));
    assert!(client.output().is_empty());

    // Try looking up Id using different APIs.
    assert_eq!(client.dag.vertex_name(id).await.unwrap(), v.clone());
    assert!(client.output().is_empty());

    assert!(matches!(
        &client.dag.contains_vertex_id_locally(&[id]).await.unwrap()[..],
        [true]
    ));
    assert!(client.output().is_empty());

    assert!(matches!(
        &client.dag.vertex_name_batch(&[id]).await.unwrap()[..],
        [Ok(n)] if n == &v
    ));
    assert!(client.output().is_empty());
}

#[tokio::test]
async fn test_local_cache_existing_vertex_to_id() {
    let client = client_for_local_cache_test().await;

    let v: VertexName = "C".into();
    let id = client.dag.vertex_id(v.clone()).await.unwrap();
    assert_eq!(client.output(), ["resolve names: [C], heads: [G]"]);

    check_local_cache(&client, v, id).await;
}

#[tokio::test]
async fn test_local_cache_existing_id_to_vertex() {
    let client = client_for_local_cache_test().await;

    let id = Id(3);
    let v = client.dag.vertex_name(id).await.unwrap();
    assert_eq!(client.output(), ["resolve paths: [G~3]"]);

    check_local_cache(&client, v, id).await;
}

#[tokio::test]
async fn test_local_cache_missing_vertex_to_id() {
    let client = client_for_local_cache_test().await;

    // Test that the local cache can prevent remote lookups resolving Vertex -> None.
    assert!(client.dag.vertex_id("Z".into()).await.is_err());
    assert_eq!(client.output(), ["resolve names: [Z], heads: [G]"]);

    // Try looking up using different APIs.
    assert!(client.dag.vertex_id("Z".into()).await.is_err());
    assert!(client.output().is_empty());

    assert!(!client.dag.contains_vertex_name(&"Z".into()).await.unwrap());
    assert!(client.output().is_empty());

    assert_eq!(
        client.dag.vertex_id_optional(&"Z".into()).await.unwrap(),
        None,
    );
    assert!(client.output().is_empty());

    assert_eq!(
        client
            .dag
            .vertex_id_with_max_group(&"Z".into(), Group::MASTER)
            .await
            .unwrap(),
        None,
    );
    assert!(client.output().is_empty());

    assert!(matches!(
        &client
            .dag
            .contains_vertex_name_locally(&["Z".into()])
            .await
            .unwrap()[..],
        [false]
    ));
    assert!(client.output().is_empty());

    assert!(matches!(
        &client.dag.vertex_id_batch(&["Z".into()]).await.unwrap()[..],
        [Err(_)]
    ));
    assert!(client.output().is_empty());
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

use super::TestDag;

#[tokio::test]
async fn test_strip_basic() {
    // Test strip C (and its descendants D, E), then H.
    // C, D are in the master group. E, F::I are in non-master.
    let mut dag = TestDag::draw_client(
        r#"
        A--B--C--D--E
         \
          F--G--H--I
        # master: D"#,
    )
    .await;

    // Note "B" is not known locally - it's lazy.
    let state_before_strip = dag.dump_state().await;
    assert_eq!(
        &state_before_strip,
        r#"<spans [F:E+N0:N4, A:D+0:3]>
Lv0: RH0-3[], N0-N3[0], N4-N4[3]
P->C: 0->N0, 3->N4
0->A 3->D N0->F N1->G N2->H N3->I N4->E"#
    );

    // Strip C.
    // Note "B" becomes the new "head" and is currently lazy.
    // So strip needs extra work to ensure "B" is non-lazy.
    dag.strip("C").await;

    let state_after_strip = dag.dump_state().await;
    assert_eq!(
        &state_after_strip,
        r#"<spans [F:I+N0:N3, A:B+0:1]>
Lv0: RH0-1[], N0-N3[0]
P->C: 0->N0
0->A 1->B N0->F N1->G N2->H N3->I"#
    );

    // Check that strip is persisted and can be seen after reopen.
    dag.reopen();
    assert_eq!(&dag.dump_state().await, &state_after_strip);

    // Strip H. The F-I segment is truncated to F-G.
    dag.strip("H").await;
    assert_eq!(
        dag.dump_state().await,
        r#"<spans [F:G+N0:N1, A:B+0:1]>
Lv0: RH0-1[], N0-N1[0]
P->C: 0->N0
0->A 1->B N0->F N1->G"#
    );
}

#[tokio::test]
async fn test_strip_branch_and_parent_remains_lazy() {
    // Test that stripping a (broken) branch does not resolve
    // the branching point. This allows strip to remove "broken"
    // branches (ex. removed by server) without resolving new
    // vertexes remotely, because the remote protocol might not
    // work.
    let ascii = r#"
        A--B--C--D--E
            \
             F--G
        # master: D F"#;

    let dag = TestDag::draw_client(ascii).await;
    assert_eq!(
        dag.dump_state().await,
        r#"<spans [G:E+N0:N1, 0:4]>
Lv0: RH0-3[], 4-4[1], N0-N0[4], N1-N1[3]
P->C: 1->4, 3->N1, 4->N0
3->D 4->F N0->G N1->E"#
    );

    let strip = |s: &'static str| async move {
        let mut dag = TestDag::draw_client(ascii).await;
        dag.strip(s).await;
        dag.dump_state().await
    };

    // Strip C. B is still lazy because it's not a head.
    assert_eq!(
        strip("C").await,
        r#"<spans [G+N0, F+4, 0:1]>
Lv0: RH0-1[], 4-4[1], N0-N0[4]
P->C: 1->4, 4->N0
4->F N0->G"#
    );

    // Strip F. B is still lazy as well.
    assert_eq!(
        strip("F").await,
        "<spans [E+N1, 0:3]>\nLv0: RH0-3[], N1-N1[3]\nP->C: 3->N1\n3->D N1->E"
    );

    // Strip C+F. B is no longer lazy.
    assert_eq!(strip("C F").await, "<spans [0:1]>\nLv0: RH0-1[]\n1->B");
}

#[tokio::test]
async fn test_reinsert_then_create_higher_level() {
    // Test re-inserting a stripped "gap" and then creating higher level segments.
    //
    // Initial state:
    //
    // Lv0: |0|1|2|3|4|5|6|7|8|9|10|
    // Lv1: |-|---|-|---|---|---|
    // Lv2: |-----|-|-------|
    //
    // Strip 3:
    //
    // Lv0: |0|1|2| |4|5|6|7|8|9|10|
    // Lv1: |-|---| |---|---|---|
    // Lv2: |-----| |-------|
    //
    // Reinsert 3 (Note: 3 no longer has Lv1 and Lv2 segments):
    //
    // Lv0: |0|1|2|3|4|5|6|7|8|9|10|
    // Lv1: |-|---| |---|---|---|
    // Lv2: |-----| |-------|
    //
    // Trigger Lv3 segment creation:
    //
    // Lv0: |0|1|2|3|4|5|6|7|8|9|10|11|12|13|14|
    // Lv1: |-|---| |---|---|---|-----|-----|
    // Lv2: |-----| |-------|---------|
    // Lv3: ?       |-----------------|

    let mut dag = TestDag::new_with_segment_size(2);

    // Use "Z" as an extra parent to avoid merging flat segments.
    dag.drawdag("Z", &[]);
    for i in 1..=10 {
        let ascii = match i {
            1 | 3 | 4 => format!("A{i}", i = i),
            _ => format!("Z-A{i} A{p}-A{i}", p = i - 1, i = i),
        };
        dag.drawdag(&ascii, &[]);
    }
    dag.flush("").await;

    assert_eq!(
        dag.dump_segments_ascii(),
        r#"
        Lv0: |N0|N1|N2|N3|N4|N5|N6|N7|N8|N9|N10|
        Lv1: |N0|N1 N2|N3|N4 N5|N6 N7|N8 N9|
        Lv2: |N0 N1 N2|N3|N4 N5 N6 N7|"#
    );

    // Strip 3.
    dag.strip("A3").await;
    assert_eq!(
        dag.dump_segments_ascii(),
        r#"
        Lv0: |N0|N1|N2| |N4|N5|N6|N7|N8|N9|N10|
        Lv1: |N0|N1 N2| |N4 N5|N6 N7|N8 N9|
        Lv2: |N0 N1 N2| |N4 N5 N6 N7|"#
    );

    // Reinsert 3. Note A3 does not have Lv1 or Lv2 segments.
    dag.drawdag("A2-A3 Z-A3", &[]);
    assert_eq!(
        dag.dump_segments_ascii(),
        r#"
        Lv0: |N0|N1|N2|N3|N4|N5|N6|N7|N8|N9|N10|
        Lv1: |N0|N1 N2|  |N4 N5|N6 N7|N8 N9|
        Lv2: |N0 N1 N2|  |N4 N5 N6 N7|"#
    );

    // Try to create Lv3 segments. This does not crash.
    dag.drawdag("A10-A11 A3-A11", &[]);
    for i in 12..=18 {
        let ascii = format!("Z-A{i} A{p}-A{i}", p = i - 1, i = i);
        dag.drawdag(&ascii, &[]);
    }
    assert_eq!(
        dag.dump_segments_ascii(),
        r#"
        Lv0: |N0|N1|N2|N3|N4|N5|N6|N7|N8|N9|N10|N11|N12|N13|N14|N15|N16|N17|N18|
        Lv1: |N0|N1 N2|  |N4 N5|N6 N7|N8 N9|N10 N11|N12 N13|N14 N15|N16 N17|
        Lv2: |N0 N1 N2|  |N4 N5 N6 N7|N8 N9 N10 N11|N12 N13 N14 N15|
        Lv3: |N0 N1 N2|  |N4 N5 N6 N7 N8 N9 N10 N11|"#
    );
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

use super::TestDag;
use crate::ops::DagAlgorithm;
use crate::ops::Parents;
use crate::Set;
use crate::Vertex;

#[tokio::test]
async fn test_to_parents() {
    let dag = TestDag::draw(
        r#"
        A---B---C---D---E
              /  \     /
         F---G    H---I"#,
    );

    assert_eq!(dump_to_parents("A E", &dag).await, "A -> []; E -> [A]");
    assert_eq!(dump_to_parents("A B", &dag).await, "A -> []; B -> [A]");
    assert_eq!(
        dump_to_parents("A B C D E", &dag).await,
        "A -> []; B -> [A]; C -> [B]; D -> [C]; E -> [D]"
    );
    assert_eq!(
        dump_to_parents("A G D E", &dag).await,
        "A -> []; G -> []; D -> [G, A]; E -> [D]"
    );
    assert_eq!(
        dump_to_parents("A G C I E", &dag).await,
        "A -> []; G -> []; C -> [G, A]; I -> [C]; E -> [I]"
    );
    assert_eq!(
        dump_to_parents("C H E", &dag).await,
        "C -> []; H -> [C]; E -> [H]"
    );
    assert_eq!(
        dump_to_parents("C D H E", &dag).await,
        "C -> []; D -> [C]; H -> [C]; E -> [H, D]"
    );
    assert_eq!(
        dump_to_parents("C D I E", &dag).await,
        "C -> []; D -> [C]; I -> [C]; E -> [D, I]"
    );
    assert_eq!(
        dump_to_parents("A B C D E F G H I", &dag).await,
        "A -> []; B -> [A]; C -> [B, G]; D -> [C]; E -> [D, I]; F -> []; G -> [F]; H -> [C]; I -> [H]"
    );
}

/// Dump parents information for given set (separated by space) as a string.
async fn dump_to_parents(names: &'static str, dag: &TestDag) -> String {
    let names: Vec<Vertex> = names.split(' ').map(Into::into).collect();
    let set = Set::from_static_names(names.clone());
    assert!(set.to_parents().await.unwrap().is_none());
    let set = dag.dag.sort(&set).await.unwrap();
    let parents = set.to_parents().await.unwrap().unwrap();
    let mut output = Vec::new();
    for name in names {
        let parents = parents.parent_names(name.clone()).await.unwrap();
        output.push(format!("{:?} -> {:?}", name, parents));
    }
    output.join("; ")
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

use std::collections::HashMap;

use anyhow::Error;
#[cfg(test)]
use chrono::TimeZone;
use chrono::Utc;
use runlog::Entry;
use runlog::Progress;

struct EntryState {
    last_time: chrono::DateTime<chrono::Utc>,
    last_download_bytes: usize,
    last_upload_bytes: usize,
}

pub struct TableGenerator {
    column_titles: Vec<String>,
    row_generator: Vec<fn(&Entry, fn() -> chrono::DateTime<Utc>, Option<&EntryState>) -> String>,
    entry_removal_delay: chrono::Duration,
    entry_states: HashMap<String, EntryState>,
}

impl TableGenerator {
    pub fn new(
        column_titles_str: String,
        entry_removal_delay: chrono::Duration,
    ) -> Result<TableGenerator, Vec<String>> {
        let column_funcs: Vec<(
            &str,
            fn(&Entry, fn() -> chrono::DateTime<Utc>, Option<&EntryState>) -> String,
        )> = vec![
            ("PID", |entry, _, _| entry.pid.to_string()),
            ("STATUS", |entry, _, _| top_status_entry(&entry.exit_code)),
            ("PROGRESS", |entry, _, _| {
                top_progress_entry(&entry.progress)
            }),
            ("TIME SPENT", |entry, current_time, _| {
                let time_spent = if let Some(end_time) = entry.end_time {
                    end_time
                } else {
                    current_time()
                } - entry.start_time;
                top_time_entry(time_spent)
            }),
            ("NET DOWN", |entry, current_time, entry_state| {
                let (duration, bytes_transferred) = match entry_state {
                    Some(entry_state) => (
                        current_time() - entry_state.last_time,
                        entry.download_bytes - entry_state.last_download_bytes,
                    ),
                    None => (chrono::Duration::seconds(0), 0_usize),
                };
                network_entry(duration, bytes_transferred)
            }),
            ("NET UP", |entry, current_time, entry_state| {
                let (duration, bytes_transferred) = match entry_state {
                    Some(entry_state) => (
                        current_time() - entry_state.last_time,
                        entry.upload_bytes - entry_state.last_upload_bytes,
                    ),
                    None => (chrono::Duration::seconds(0), 0_usize),
                };
                network_entry(duration, bytes_transferred)
            }),
            ("CMD", |entry, _, _| entry.command.join(" ")),
        ];
        let columns_map: HashMap<_, _> = column_funcs.iter().copied().collect();
        let column_titles: Vec<_> = if !column_titles_str.is_empty() {
            column_titles_str.split(',').map(|x| x.trim()).collect()
        } else {
            column_funcs.iter().map(|(x, _)| *x).collect()
        };

        let unexpected_columns: Vec<_> = column_titles
            .iter()
            .filter(|&c| !columns_map.contains_key(c))
            .copied()
            .collect();
        if !unexpected_columns.is_empty() {
            return Err(unexpected_columns
                .iter()
                .map(|&x| String::from(x))
                .collect());
        }
        Ok(TableGenerator {
            column_titles: column_titles.iter().map(|&x| String::from(x)).collect(),
            row_generator: column_titles
                .iter()
                .map(|&c| *columns_map.get(c).unwrap())
                .collect(),
            entry_removal_delay,
            entry_states: HashMap::new(),
        })
    }

    pub fn column_titles(&self) -> &Vec<String> {
        &self.column_titles
    }

    pub fn generate_rows<'a>(
        &'a mut self,
        runlog_entries: impl Iterator<Item = Result<(Entry, bool), Error>> + 'a,
        current_time: fn() -> chrono::DateTime<Utc>,
    ) -> impl Iterator<Item = Vec<String>> + 'a {
        runlog_entries.filter_map(move |entry| {
            let (entry, running) = match entry {
                Ok((entry, running)) => (entry, running),
                Err(_) => {
                    return None;
                }
            };
            let end_time_filter_fn =
                |end_time| current_time() - end_time <= self.entry_removal_delay;
            if running || entry.end_time.map_or(false, end_time_filter_fn) {
                let entry_state = self.entry_states.get(&entry.id);
                let row = self
                    .row_generator
                    .iter()
                    .map(|&f| f(&entry, current_time, entry_state))
                    .collect();
                self.entry_states.insert(
                    entry.id,
                    EntryState {
                        last_time: current_time(),
                        last_download_bytes: entry.download_bytes,
                        last_upload_bytes: entry.upload_bytes,
                    },
                );
                Some(row)
            } else {
                None
            }
        })
    }
}

fn top_status_entry(exit_code: &Option<i32>) -> String {
    match exit_code {
        Some(code) => format!("EXITED ({})", code),
        None => "RUNNING".to_string(),
    }
}

fn top_time_entry(time_spent: chrono::Duration) -> String {
    static MILLIS_IN_SECOND: i64 = 1000;
    static MILLIS_IN_MINUTE: i64 = MILLIS_IN_SECOND * 60;
    static MILLIS_IN_HOUR: i64 = MILLIS_IN_MINUTE * 60;
    let milliseconds_spent = time_spent.num_milliseconds();
    let (unit_str, millis_in_unit) = if milliseconds_spent >= MILLIS_IN_HOUR {
        ("h", MILLIS_IN_HOUR)
    } else if milliseconds_spent >= MILLIS_IN_MINUTE {
        ("m", MILLIS_IN_MINUTE)
    } else {
        ("s", MILLIS_IN_SECOND)
    };
    let units_spent = milliseconds_spent as f64 / millis_in_unit as f64;
    if units_spent < 10_f64 {
        format!("{:.1}{}", units_spent, unit_str)
    } else {
        format!("{:.0}{}", units_spent, unit_str)
    }
}

fn top_progress_entry(progress_bars: &[Progress]) -> String {
    for progress_bar in progress_bars.iter() {
        if progress_bar.total > 0 {
            return format!(
                "{:.1}%",
                progress_bar.position as f32 / progress_bar.total as f32 * 100_f32
            );
        }
    }
    String::from("-")
}

fn network_entry(time_spent: chrono::Duration, bytes_transferred: usize) -> String {
    let time_spent = time_spent.num_seconds();
    if time_spent == 0 {
        return String::from("-");
    }
    let rate = ((bytes_transferred * 8) as f32) / (time_spent as f32);
    let (rate, prefix) = if rate >= 1e5 {
        (rate / 1e6, "M")
    } else {
        (rate / 1e3, "k")
    };
    format!("{:.1} {}b/s", rate, prefix)
}

#[cfg(test)]
mod tests {
    use super::*;

    #[derive(Debug)]
    struct ParseError {}

    impl std::fmt::Display for ParseError {
        fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
            write!(f, "Oh no, something bad went down")
        }
    }

    impl std::error::Error for ParseError {
        fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
            None
        }
    }

    #[test]
    fn test_table_generator() {
        let default_removal_delay = chrono::Duration::seconds(1);
        let default_generator = || TableGenerator::new("".to_string(), default_removal_delay);
        let default_time = || Utc.timestamp_opt(3, 100000000).unwrap();

        // Test invalid columns
        assert!(
            matches!(TableGenerator::new("not a valid column, PID, not_valid_either".to_string(), default_removal_delay),
            Err(x) if x == vec![
                "not a valid column".to_string(),
                "not_valid_either".to_string()
            ])
        );
        // Test default column order
        let default_columns = vec![
            "PID".to_string(),
            "STATUS".to_string(),
            "PROGRESS".to_string(),
            "TIME SPENT".to_string(),
            "NET DOWN".to_string(),
            "NET UP".to_string(),
            "CMD".to_string(),
        ];
        assert_eq!(
            *default_generator().unwrap().column_titles(),
            default_columns
        );
        // Test specific columns
        assert_eq!(
            *TableGenerator::new("CMD,   PID".to_string(), default_removal_delay)
                .unwrap()
                .column_titles(),
            vec!["CMD".to_string(), "PID".to_string()]
        );
        // Test single row column format
        let runlog_entries = vec![Ok((
            Entry {
                id: "1".to_string(),
                command: vec!["somecommand".to_string(), "somearg".to_string()],
                pid: 101,
                download_bytes: 0,
                upload_bytes: 0,
                start_time: Utc.timestamp_opt(0, 0).unwrap(),
                end_time: None,
                exit_code: None,
                progress: vec![Progress {
                    topic: "spinning".to_string(),
                    unit: "".to_string(),
                    total: 100,
                    position: 2,
                }],
            },
            true,
        ))];
        let expected_rows = vec![vec![
            "101".to_string(),
            "RUNNING".to_string(),
            "2.0%".to_string(),
            "3.1s".to_string(),
            "-".to_string(),
            "-".to_string(),
            "somecommand somearg".to_string(),
        ]];
        assert_eq!(
            default_generator()
                .unwrap()
                .generate_rows(runlog_entries.into_iter(), default_time)
                .collect::<Vec<_>>(),
            expected_rows
        );
        // Test row filtering
        let err = ParseError {};
        let runlog_entries: Vec<Result<(Entry, bool), Error>> = vec![
            Err(Error::new(err)),
            Ok((
                Entry {
                    id: "0".to_string(),
                    command: vec!["".to_string()],
                    pid: 0,
                    download_bytes: 0,
                    upload_bytes: 0,
                    start_time: Utc.timestamp_opt(0, 0).unwrap(),
                    end_time: Some(Utc.timestamp_opt(2, 0).unwrap()),
                    exit_code: Some(0),
                    progress: vec![],
                },
                false,
            )),
            Ok((
                Entry {
                    id: "1".to_string(),
                    command: vec!["notarealcommand".to_string()],
                    pid: 321,
                    download_bytes: 0,
                    upload_bytes: 0,
                    start_time: Utc.timestamp_opt(0, 0).unwrap(),
                    end_time: Some(Utc.timestamp_opt(3, 0).unwrap()),
                    exit_code: Some(123),
                    progress: vec![],
                },
                false,
            )),
        ];

        let expected_rows = vec![vec![
            "321".to_string(),
            "EXITED (123)".to_string(),
            "-".to_string(),
            "3.0s".to_string(),
            "-".to_string(),
            "-".to_string(),
            "notarealcommand".to_string(),
        ]];
        assert_eq!(
            default_generator()
                .unwrap()
                .generate_rows(runlog_entries.into_iter(), default_time)
                .collect::<Vec<_>>(),
            expected_rows
        );
    }

    #[test]
    fn test_status_entry() {
        assert_eq!(top_status_entry(&Some(0)), "EXITED (0)");
        assert_eq!(top_status_entry(&None), "RUNNING");
    }

    #[test]
    fn test_time_entry() {
        assert_eq!(
            top_time_entry(chrono::Duration::hours(5) + chrono::Duration::minutes(6)),
            "5.1h"
        );
        assert_eq!(
            top_time_entry(chrono::Duration::hours(12) + chrono::Duration::minutes(6)),
            "12h"
        );
        assert_eq!(
            top_time_entry(chrono::Duration::minutes(2) + chrono::Duration::seconds(42)),
            "2.7m"
        );
        assert_eq!(
            top_time_entry(chrono::Duration::minutes(30) + chrono::Duration::seconds(42)),
            "31m"
        );
        assert_eq!(
            top_time_entry(chrono::Duration::seconds(1) + chrono::Duration::milliseconds(357)),
            "1.4s"
        );
        assert_eq!(
            top_time_entry(chrono::Duration::seconds(41) + chrono::Duration::milliseconds(357)),
            "41s"
        );
    }

    #[test]
    fn test_progress_entry() {
        // Test picking only the first progress bar that has a positive total
        assert_eq!(
            top_progress_entry(&[
                Progress {
                    topic: String::from(""),
                    unit: String::from(""),
                    position: 0,
                    total: 0,
                },
                Progress {
                    topic: String::from(""),
                    unit: String::from(""),
                    position: 123,
                    total: 1000,
                },
                Progress {
                    topic: String::from(""),
                    unit: String::from(""),
                    position: 0,
                    total: 100000,
                },
            ]),
            "12.3%"
        );
        // Test total = 0
        assert_eq!(
            top_progress_entry(&[Progress {
                topic: String::from(""),
                unit: String::from(""),
                position: 0,
                total: 0,
            },]),
            "-"
        );
        // Test empty list
        assert_eq!(top_progress_entry(&[]), "-");
    }

    #[test]
    fn test_network_entry() {
        assert_eq!(
            network_entry(chrono::Duration::seconds(0), 5000),
            String::from("-"),
        );
        assert_eq!(
            network_entry(chrono::Duration::seconds(2), 400),
            String::from("1.6 kb/s"),
        );
        assert_eq!(
            network_entry(chrono::Duration::seconds(1), 37500),
            String::from("0.3 Mb/s"),
        );
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

//! This is a test. Run the test with:
//!
//! ```bash,ignore
//! cargo run --example test
//! ```

use std::process::Command;

dev_logger::init!();

const MESSAGE: &str = "hello_log_message_to_test";

fn main() {
    if std::env::args().any(|a| a == MESSAGE) {
        // Run as a child process that might print the MESSAGE.
        tracing::info!(name = MESSAGE);
        return;
    }

    // Test RUST_LOG=info enables tracing::info! messages.
    assert!(run_child("info").contains(MESSAGE));

    // Test RUST_LOG=warn disables tracing::info! messages.
    assert!(!run_child("warn").contains(MESSAGE));
}

fn run_child(rust_log: &str) -> String {
    std::env::set_var("RUST_LOG", rust_log);
    let exe = std::env::current_exe().unwrap();
    let out = Command::new(exe)
        .args(&[MESSAGE])
        .output()
        .expect("failed to execute process");
    String::from_utf8_lossy(&out.stderr).to_string()
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

//! Convenient env_logger for testing purpose.
//!
//! # Example
//!
//! ```
//! // In lib.rs:
//! #[cfg(test)]
//! dev_logger::init!();
//!
//! // In test function:
//! tracing::info!(name = "message");
//!
//! // Set RUST_LOG=info and run the test.
//! ```

use std::io;
use std::sync::Arc;
use std::sync::Mutex;

pub use ctor::ctor;
use tracing_subscriber::fmt::format::FmtSpan;
use tracing_subscriber::fmt::MakeWriter;
use tracing_subscriber::fmt::Subscriber;
use tracing_subscriber::EnvFilter;

/// Initialize tracing and env_logger for adhoc logging (ex. in a library test)
/// purpose.
pub fn init() {
    let builder = Subscriber::builder()
        .with_env_filter(EnvFilter::from_env("LOG"))
        .with_ansi(false)
        .with_target(false)
        .without_time()
        .with_span_events(FmtSpan::ACTIVE);

    builder.init();
}

/// Trace the given function using the given filter (in EnvFilter format).
/// Return strings representing the traced logs.
pub fn traced(filter: &str, func: impl FnOnce()) -> Vec<String> {
    #[derive(Clone, Default)]
    struct Output(Arc<Mutex<Vec<String>>>);

    impl MakeWriter<'_> for Output {
        type Writer = Output;
        fn make_writer(&self) -> Self::Writer {
            self.clone()
        }
    }

    impl io::Write for Output {
        fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
            let mut lines = self.0.lock().unwrap();
            let mut s = String::from_utf8_lossy(buf).trim().to_string();

            // Buck unittest targets add "_unittest" suffix to crate names
            // that will affect the "target". Workaround it by removing the
            // suffix.
            if cfg!(fbcode_build) {
                s = s.replace("_unittest: ", ": ");
                s = s.replace("_unittest::", "::");
            }

            lines.push(s);
            Ok(buf.len())
        }
        fn flush(&mut self) -> io::Result<()> {
            Ok(())
        }
    }

    let out = Output::default();
    let builder = Subscriber::builder()
        .with_env_filter(EnvFilter::new(filter))
        .with_ansi(false)
        .without_time()
        .with_writer(out.clone())
        .with_span_events(FmtSpan::ACTIVE);
    let dispatcher = builder.finish();
    tracing::subscriber::with_default(dispatcher, func);

    let lines = out.0.lock().unwrap();
    lines.clone()
}

/// Call `init` on startup. This is useful for tests.
#[macro_export]
macro_rules! init {
    () => {
        #[dev_logger::ctor]
        fn dev_logger_init_ctor() {
            dev_logger::init();
        }
    };
}

#[test]
fn test_traced() {
    let lines = traced("info", || {
        tracing::info_span!("bar", x = 1).in_scope(|| {
            tracing::info!("foo");
            tracing::debug!("foo2");
        });
    });
    assert_eq!(
        lines,
        [
            "INFO bar{x=1}: dev_logger: enter",
            "INFO bar{x=1}: dev_logger: foo",
            "INFO bar{x=1}: dev_logger: exit"
        ]
    );
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

use std::fmt;
use std::io;
use std::net::SocketAddr;
use std::net::TcpStream;
use std::net::ToSocketAddrs;
use std::time::Duration;
use std::vec;

use auth::AuthSection;
use configmodel::Config;
use configmodel::ConfigExt;
use http_client::HttpClientError;
use thiserror::Error;
use url::Host;
use url::Url;

#[derive(Debug, Error)]
pub enum HostError {
    #[error("DNS error: {0}")]
    DNS(io::Error),
    #[error("TCP error: {0}")]
    TCP(io::Error),
    #[error("invalid host config: {0}")]
    Config(String),
}

#[derive(Debug, Error)]
pub enum HttpError {
    #[error("unexpected http response: {0:?}")]
    UnexpectedResponse(HttpResponse),
    #[error("http request failure: {0:?}")]
    RequestFailure(HttpClientError),
    #[error(transparent)]
    MissingCerts(#[from] auth::MissingCerts),
    #[error("{0}")]
    InvalidCert(auth::X509Error, HttpClientError),
    #[error("invalid http config: {0}")]
    Config(String),
}

#[derive(Debug, Error)]
pub enum Diagnosis {
    #[error("invalid config: {0}")]
    BadConfig(String),
    #[error("no internet connectivity: {0}")]
    NoInternet(HostError),
    #[error("no corp connectivity: {0}")]
    NoCorp(HostError),
    #[error("x2pagentd problem: {0}")]
    AuthProxyProblem(HttpError),
    #[error("http connection problem: {0}")]
    HttpProblem(HttpError),
}

pub struct HttpResponse {
    status: http::StatusCode,
    headers: http::HeaderMap,
    body: Vec<u8>,
}

impl fmt::Debug for HttpResponse {
    // Turn body into text based on content-type.
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match header_value(&self.headers, "content-type") {
            Some("application/json") | Some("text/html") => {
                write!(
                    f,
                    "HttpResponse {{ status: {:?}, headers: {:?}, body: {:?} }}",
                    self.status,
                    self.headers,
                    String::from_utf8_lossy(&self.body)
                )
            }
            _ => {
                write!(
                    f,
                    "HttpResponse {{ status: {:?}, headers: {:?}, body: {:?} }}",
                    self.status, self.headers, self.body
                )
            }
        }
    }
}

impl Diagnosis {
    pub fn treatment(&self, config: &dyn Config) -> String {
        match self {
            Self::NoCorp(_) => {
                "Please check your VPN connection (internet okay, but can't reach corp)."
                    .to_string()
            }
            Self::NoInternet(_) => {
                "Please check your internet connection (failed external connectivity test)."
                    .to_string()
            }
            Self::BadConfig(msg) => format!("Invalid config: {}", msg),
            Self::AuthProxyProblem(_) => {
                let mut msg = "Local auth proxy problem.".to_string();
                if let Some(help) = config.get("help", "auth-proxy-help") {
                    msg = format!("{} {}", msg, help.as_ref());
                }
                msg
            }
            Self::HttpProblem(err) => diagnose_http_error(config, err),
        }
    }

    fn short_name(&self) -> String {
        match self {
            Diagnosis::BadConfig(_) => "bad_config".to_string(),
            Diagnosis::NoInternet(_) => "no_internet".to_string(),
            Diagnosis::NoCorp(_) => "no_corp".to_string(),
            Diagnosis::AuthProxyProblem(err) => format!("auth_proxy_problem({})", err.short_name()),
            Diagnosis::HttpProblem(err) => format!("http_problem({})", err.short_name()),
        }
    }
}

impl HttpError {
    fn short_name<'a>(&'a self) -> &'a str {
        match self {
            HttpError::UnexpectedResponse(res) => res.status.as_str(),
            HttpError::RequestFailure(_) => "other",
            HttpError::MissingCerts(_) => "missing_certs",
            HttpError::Config(_) => "config",
            HttpError::InvalidCert(_, _) => "invalid_cert",
        }
    }
}

fn diagnose_http_error(config: &dyn Config, err: &HttpError) -> String {
    let maybe_append_help = |mut msg: String, help_name: &str| -> String {
        if let Some(help) = config.get("help", help_name) {
            msg = format!("{}\n\n{}", msg, help.as_ref());
        }
        msg
    };

    match err {
        HttpError::UnexpectedResponse(res) => diagnose_unexpected_response(res),
        HttpError::RequestFailure(HttpClientError::Tls(err)) => maybe_append_help(
            // We weren't able to diagnose a particular cert problem,
            // so give a generic TLS message. Include the error since
            // it may have something more useful.
            format!("TLS error - please check your certificates.\n\n{}", err),
            "tlshelp",
        ),
        HttpError::InvalidCert(err, _) => maybe_append_help(format!("{}", err), "tlsauthhelp"),
        HttpError::MissingCerts(err) => maybe_append_help(format!("{}", err), "tlsauthhelp"),
        HttpError::RequestFailure(HttpClientError::Curl(err)) => diagnose_curl_error(err),

        HttpError::Config(err) => err.to_string(),
        HttpError::RequestFailure(_) => format!("{}", err),
    }
}

fn diagnose_curl_error(err: &curl::Error) -> String {
    if err.is_operation_timedout() {
        "Network timeout. Please check your connection.".to_string()
    } else {
        format!("{}", err)
    }
}

fn diagnose_unexpected_response(res: &HttpResponse) -> String {
    match res.status {
        http::StatusCode::FORBIDDEN => {
            "You lack permission for this repo. Please see https://fburl.com/svnuser, and make sure you have permission for this repo.".to_string()
        }
        _ => {
            let mut header_hints = vec![];
            if let Some(x2p_error_type) = header_value(&res.headers, "x-x2pagentd-error-type") {
                let mut hint = format!("x2pagentd: {}", x2p_error_type);
                if let Some(x2p_msg) = header_value(&res.headers, "x-x2pagentd-error-msg") {
                    hint = format!("{} ({})", hint, x2p_msg);
                }
                header_hints.push(hint);
            }

            for (n, v) in res.headers.iter() {
                if let Some(advice_type) = n.as_str().strip_prefix("x-fb-validated-x2pauth-advice-")
                {
                    if let Ok(v) = v.to_str() {
                        header_hints.push(format!("x2p auth {}: {}", advice_type, v));
                    }
                }
            }

            format!(
                "Unexpected server response: {}{}",
                res.status,
                header_hints
                    .iter()
                    .map(|h| format!("\n\t{}", h))
                    .collect::<Vec<String>>()
                    .join(""),
            )
        }
    }
}

fn header_value<'a>(h: &'a http::HeaderMap, key: &str) -> Option<&'a str> {
    h.get(key).and_then(|val| val.to_str().ok())
}

pub struct Doctor {
    // Allow tests to stub DNS responses.
    dns_lookup: Box<dyn Fn(&str) -> io::Result<vec::IntoIter<SocketAddr>>>,
    tcp_connect_timeout: Duration,
    stub_healthcheck_response:
        Option<Box<dyn Fn(&Url, bool) -> Result<HttpResponse, HttpClientError>>>,
}

fn real_dns_lookup(host_port: &str) -> io::Result<vec::IntoIter<SocketAddr>> {
    host_port.to_socket_addrs()
}

impl Doctor {
    pub fn new() -> Self {
        Doctor {
            dns_lookup: Box::new(real_dns_lookup),
            tcp_connect_timeout: Duration::from_secs(1),
            stub_healthcheck_response: None,
        }
    }

    fn check_host_tcp(&self, url: &Url) -> Result<(), HostError> {
        let port = url.port().unwrap_or(443);

        let sock_addrs: Vec<SocketAddr> = match url.host() {
            None => {
                return Err(HostError::Config(format!("url {} has no host", url)));
            }
            Some(Host::Domain(host)) => {
                // Let tests stub out host check results.
                if host == "test_fail" {
                    return Err(HostError::TCP(io::Error::new(io::ErrorKind::Other, "test")));
                } else if host == "test_succeed" {
                    return Ok(());
                }

                match (self.dns_lookup)(&format!("{}:{}", host, port)) {
                    Err(err) => return Err(HostError::DNS(err)),
                    Ok(addrs) => addrs.collect(),
                }
            }
            Some(Host::Ipv4(ip)) => vec![SocketAddr::new(ip.into(), port)],
            Some(Host::Ipv6(ip)) => vec![SocketAddr::new(ip.into(), port)],
        };

        if sock_addrs.is_empty() {
            return Err(HostError::DNS(io::Error::new(
                io::ErrorKind::Other,
                format!("{:?} resolved to 0 IPs", url.host()),
            )));
        }

        for i in 0..sock_addrs.len() {
            match TcpStream::connect_timeout(&sock_addrs[i], self.tcp_connect_timeout) {
                Err(err) => {
                    if i == sock_addrs.len() - 1 {
                        return Err(HostError::TCP(err));
                    }
                }
                Ok(_) => break,
            }
        }

        Ok(())
    }

    pub fn diagnose(&self, config: &dyn Config) -> Result<(), Diagnosis> {
        let res = || -> Result<(), Diagnosis> {
            self.check_corp_connectivity(config)?;
            self.check_http_connectivity(config)?;
            Ok(())
        }();

        let short_res = match &res {
            Err(err) => err.short_name(),
            Ok(()) => "ok".to_string(),
        };
        tracing::debug!(target: "network_doctor_diagnosis", network_doctor_diagnosis=&short_res[..]);

        res
    }

    fn check_corp_connectivity(&self, config: &dyn Config) -> Result<(), Diagnosis> {
        let repo_url = config_url(config, "edenapi", "url", None)?;

        // First check for corp connectivity.
        let corp_err = match self.check_host_tcp(&repo_url) {
            Ok(()) => return Ok(()),
            Err(err) => err,
        };

        // If we don't have corp connectivity, see if we have internet connectivity.
        let external_url = config_url(
            config,
            "doctor",
            "external-host-check-url",
            Some("https://www.facebook.com"),
        )?;

        match self.check_host_tcp(&external_url) {
            Ok(()) => Err(Diagnosis::NoCorp(corp_err)),
            Err(err) => Err(Diagnosis::NoInternet(err)),
        }
    }

    fn check_http_connectivity(&self, config: &dyn Config) -> Result<(), Diagnosis> {
        let mut url = config_url(config, "edenapi", "url", None)?;
        let repo_name = match config.get("remotefilelog", "reponame") {
            Some(name) => name.to_string(),
            None => {
                return Err(Diagnosis::BadConfig(
                    "remotefilelog.reponame is not set".to_string(),
                ));
            }
        };

        // Build edenpi/:repo/capabilities URL. This will suss out permission
        // errors (as opposed to :repo/health_check).
        match url.path_segments_mut() {
            Ok(mut path) => path.pop_if_empty().push(&repo_name).push("capabilities"),
            Err(()) => return Err(Diagnosis::BadConfig("bad edenapi.url".into())),
        };

        let mut x2pagentd_err = None;
        if use_x2pagentd(config, &url) {
            x2pagentd_err = match self.check_host_http(config, &url, true) {
                Ok(()) => return Ok(()),
                Err(err) => Some(err),
            }
        }

        match (self.check_host_http(config, &url, false), x2pagentd_err) {
            (Ok(()), Some(err)) => Err(Diagnosis::AuthProxyProblem(err)),
            (Ok(()), None) => Ok(()),
            (Err(_), Some(proxy_err)) => Err(Diagnosis::HttpProblem(proxy_err)),
            (Err(err), None) => Err(Diagnosis::HttpProblem(err)),
        }
    }

    fn check_host_http(
        &self,
        config: &dyn Config,
        url: &Url,
        use_x2pagentd: bool,
    ) -> Result<(), HttpError> {
        tracing::debug!(%url, use_x2pagentd, "check_host_http");

        let mut hc = hg_http::http_config(config, url)?;

        if !use_x2pagentd {
            hc.unix_socket_path = None;

            let auth = AuthSection::from_config(config).best_match_for(url)?;
            (hc.cert_path, hc.key_path, hc.ca_path) = auth
                .map(|auth| (auth.cert, auth.key, auth.cacerts))
                .unwrap_or_default();

            if url.scheme() == "https" && hc.cert_path.is_none() {
                return Err(HttpError::Config(format!("no auth section for {}", url)));
            }
        }

        let result = if let Some(stub) = &self.stub_healthcheck_response {
            stub(url, use_x2pagentd)
        } else {
            let mut req = hg_http::http_client("network-doctor", hc.clone()).get(url.clone());
            req.set_timeout(Duration::from_secs(3));
            req.send().map(|res| HttpResponse {
                status: res.status(),
                headers: res.headers().clone(),
                body: res.body().to_vec(),
            })
        };

        match result {
            Ok(res) if res.status.is_success() => Ok(()),
            Ok(res) => Err(HttpError::UnexpectedResponse(res)),
            Err(err) => {
                if let HttpClientError::Tls(_) = err {
                    if let Some(cert_path) = &hc.cert_path {
                        if let Err(x509_err) = auth::check_certs(cert_path) {
                            return Err(HttpError::InvalidCert(x509_err, err));
                        }
                    }
                }

                Err(HttpError::RequestFailure(err))
            }
        }
    }
}

fn use_x2pagentd(config: &dyn Config, url: &Url) -> bool {
    let hc = match hg_http::http_config(config, url) {
        Ok(hc) => hc,
        Err(_) => return false,
    };
    if hc.unix_socket_path.is_none() {
        return false;
    }

    match url.domain() {
        None => false,
        Some(domain) => hc.unix_socket_domains.contains(domain),
    }
}

// Extract and parse URL from config with optional default if no config value.
fn config_url(
    config: &dyn Config,
    section: &str,
    field: &str,
    default: Option<&str>,
) -> Result<Url, Diagnosis> {
    let url: String = match (config.get_nonempty_opt(section, field), default) {
        (Err(err), _) => return Err(Diagnosis::BadConfig(format!("config error: {}", err))),
        (Ok(Some(url)), _) => url,
        (Ok(None), Some(default)) => default.to_string(),
        (Ok(None), None) => {
            return Err(Diagnosis::BadConfig(format!(
                "no config for {}.{}",
                section, field
            )));
        }
    };

    match Url::parse(&url) {
        Err(err) => Err(Diagnosis::BadConfig(format!(
            "invalid url {}: {}",
            url, err
        ))),
        Ok(url) => Ok(url),
    }
}

#[cfg(test)]
mod tests {
    use std::collections::BTreeMap;
    use std::collections::HashMap;
    use std::net::SocketAddr;
    use std::net::TcpListener;

    use http::StatusCode;
    use tempfile::tempdir;

    use super::*;

    macro_rules! response {
        ($status:expr $(, $name:tt : $val:tt)* $(,)?) => {
            response!($status, b"", $($name : $val,)*)
        };
        ($status:expr, $body:expr $(, $name:tt : $val:tt)* $(,)?) => {
            HttpResponse {
                status: $status,
                headers: (&HashMap::<String, String>::from([$(($name.to_string(), $val.to_string()),)*])).try_into().unwrap(),
                body: $body.to_vec(),
            }
        };
    }

    #[test]
    fn test_check_host() {
        let mut doc = Doctor::new();

        let listener = TcpListener::bind("127.0.0.1:0").unwrap();
        let addr = listener.local_addr().unwrap();

        // No DNS lookup required - can connect directly to IP.
        {
            doc.dns_lookup = Box::new(move |_host_port| -> io::Result<vec::IntoIter<SocketAddr>> {
                panic!("dns lookup!")
            });

            let url = Url::parse(&format!("https://{}:{}", addr.ip(), addr.port())).unwrap();
            assert!(doc.check_host_tcp(&url).is_ok());
        }

        // DNS lookup is okay.
        {
            doc.dns_lookup = Box::new(move |_host_port| -> io::Result<vec::IntoIter<SocketAddr>> {
                Ok(vec![addr].into_iter())
            });

            let url = Url::parse(&format!("https://localhost:{}", addr.port())).unwrap();
            assert!(doc.check_host_tcp(&url).is_ok());
        }

        // DNS lookup fails.
        {
            doc.dns_lookup = Box::new(move |_host_port| -> io::Result<vec::IntoIter<SocketAddr>> {
                Err(io::Error::new(
                    io::ErrorKind::Other,
                    "some sort of DNS failure!",
                ))
            });

            let url = Url::parse(&format!("https://localhost:{}", addr.port())).unwrap();
            assert!(matches!(doc.check_host_tcp(&url), Err(HostError::DNS(_))));
        }

        // No one is listening.
        {
            doc.tcp_connect_timeout = Duration::from_millis(1);

            let url = Url::parse("https://169.254.0.1:1234").unwrap();
            assert!(matches!(doc.check_host_tcp(&url), Err(HostError::TCP(_))))
        }
    }

    #[test]
    fn test_check_corp_connectivity() {
        let non_working_url = "https://169.254.0.1:1234";

        // Both corp and external fail.
        {
            let mut cfg = BTreeMap::new();
            cfg.insert("edenapi.url", non_working_url);
            cfg.insert("doctor.external-host-check-url", non_working_url);

            let mut doc = Doctor::new();
            doc.tcp_connect_timeout = Duration::from_millis(1);

            assert!(matches!(
                doc.check_corp_connectivity(&cfg),
                Err(Diagnosis::NoInternet(_))
            ));
        }

        let listener = TcpListener::bind("127.0.0.1:0").unwrap();
        let listener_addr = listener.local_addr().unwrap();
        let working_url = format!("https://{}:{}", listener_addr.ip(), listener_addr.port());

        // External works.
        {
            let mut cfg = BTreeMap::new();
            cfg.insert("edenapi.url", non_working_url);
            cfg.insert("doctor.external-host-check-url", &working_url);

            let doc = Doctor::new();
            assert!(matches!(
                doc.check_corp_connectivity(&cfg),
                Err(Diagnosis::NoCorp(_))
            ));
        }

        // Corp works.
        {
            let mut cfg: BTreeMap<&str, &str> = BTreeMap::new();
            cfg.insert("edenapi.url", &working_url);
            cfg.insert("doctor.external-host-check-url", &working_url);

            let doc = Doctor::new();
            assert!(matches!(doc.check_corp_connectivity(&cfg), Ok(())));
        }
    }

    #[test]
    fn test_check_http_connectivity() {
        let td = tempdir().unwrap();

        let mut doc = Doctor::new();

        let fake_cert_path = td.path().join("cert");
        std::fs::write(&fake_cert_path, "foo").unwrap();

        let mut cfg = BTreeMap::new();
        cfg.insert("edenapi.url", "https://example.com/edenapi/");
        cfg.insert("remotefilelog.reponame", "some_repo");
        cfg.insert("auth.test.prefix", "*");
        cfg.insert("auth.test.cert", fake_cert_path.to_str().unwrap());
        cfg.insert("auth.test.key", fake_cert_path.to_str().unwrap());

        // Happy path - server returns 200 for /capabilities.
        {
            doc.stub_healthcheck_response = Some(Box::new(|url, _x2p| {
                assert_eq!(url.path(), "/edenapi/some_repo/capabilities");
                Ok(response!(http::StatusCode::OK))
            }));
            assert!(matches!(doc.check_http_connectivity(&cfg), Ok(())));
        }

        // Unexpected HTTP response.
        {
            doc.stub_healthcheck_response = Some(Box::new(|_url, _x2p| {
                Ok(response!(http::StatusCode::FORBIDDEN))
            }));
            assert!(matches!(
                doc.check_http_connectivity(&cfg),
                Err(Diagnosis::HttpProblem(HttpError::UnexpectedResponse(
                    HttpResponse {
                        status: http::StatusCode::FORBIDDEN,
                        ..
                    }
                )))
            ));
        }

        // Simulate x2pagentd specific problem.
        {
            let mut cfg = cfg.clone();

            cfg.insert("auth_proxy.unix_socket_path", "/dev/null");
            cfg.insert("auth_proxy.unix_socket_domains", "example.com");

            // Simulate x2pagentd specific problem.
            doc.stub_healthcheck_response = Some(Box::new(|_url, x2p| {
                if x2p {
                    Ok(response!(http::StatusCode::INTERNAL_SERVER_ERROR))
                } else {
                    Ok(response!(http::StatusCode::OK))
                }
            }));
            assert!(matches!(
                doc.check_http_connectivity(&cfg),
                Err(Diagnosis::AuthProxyProblem(HttpError::UnexpectedResponse(
                    HttpResponse {
                        status: http::StatusCode::INTERNAL_SERVER_ERROR,
                        ..
                    }
                )))
            ));
        }

        // Make sure we support non-https.
        {
            let mut cfg = cfg.clone();

            cfg.insert("edenapi.url", "http://example.com/edenapi/");

            cfg.insert("auth.test.prefix", "doesnt_match");

            doc.stub_healthcheck_response =
                Some(Box::new(|_url, _x2p| Ok(response!(http::StatusCode::OK))));
            assert!(matches!(doc.check_http_connectivity(&cfg), Ok(())));
        }

        // Give a specific error for missing auth section.
        {
            let mut cfg = cfg.clone();

            cfg.insert("auth.test.prefix", "doesnt_match");

            doc.stub_healthcheck_response =
                Some(Box::new(|_url, _x2p| Ok(response!(http::StatusCode::OK))));

            match doc.check_http_connectivity(&cfg) {
                Err(Diagnosis::HttpProblem(HttpError::Config(msg))) => {
                    assert_eq!(
                        msg,
                        "no auth section for https://example.com/edenapi/some_repo/capabilities"
                    );
                }
                _ => panic!(),
            };
        }
    }

    #[test]
    fn test_diagnose_unexpected_response() {
        assert_eq!(
            diagnose_unexpected_response(&response!(
                StatusCode::INTERNAL_SERVER_ERROR,
                "x-x2pagentd-error-type": "apple",
                "x-x2pagentd-error-msg": "not crispy",
            )),
            "Unexpected server response: 500 Internal Server Error\n\tx2pagentd: apple (not crispy)"
        );

        assert_eq!(
            diagnose_unexpected_response(&response!(
                StatusCode::INTERNAL_SERVER_ERROR,
                "x-fb-validated-x2pauth-advice-access-denied": "reboot laptop",
            )),
            "Unexpected server response: 500 Internal Server Error\n\tx2p auth access-denied: reboot laptop"
        );

        assert_eq!(
            diagnose_unexpected_response(&response!(StatusCode::FORBIDDEN)),
            "You lack permission for this repo. Please see https://fburl.com/svnuser, and make sure you have permission for this repo.",
        );
    }

    #[test]
    fn test_response_body_format() {
        assert_eq!(
            format!(
                "{:?}",
                response!(
                    StatusCode::INTERNAL_SERVER_ERROR,
                    b"some body",
                    "content-type": "application/octet-stream",
                )
            ),
            "HttpResponse { status: 500, headers: {\"content-type\": \"application/octet-stream\"}, body: [115, 111, 109, 101, 32, 98, 111, 100, 121] }"
        );

        assert_eq!(
            format!(
                "{:?}",
                response!(
                    StatusCode::INTERNAL_SERVER_ERROR,
                    b"{}",
                    "content-type": "application/json",
                )
            ),
            "HttpResponse { status: 500, headers: {\"content-type\": \"application/json\"}, body: \"{}\" }"
        );
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

//! # drawdag
//!
//! Utilities to parse ASCII revision DAG and create commits from them.

use std::collections::BTreeMap;
use std::collections::BTreeSet;
use std::collections::HashSet;

mod succ;

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
enum Direction {
    /// From bottom to top. Roots are at the bottom.
    BottomTop,

    /// From left to right. Roots are at the left.
    LeftRight,
}

/// Parse an ASCII DAG. Extract edge information.
/// Return a map from names to their parents.
///
/// The direction of the graph is automatically detected.
/// If `|` is used, then roots are at the bottom, heads are at the top side.
/// Otherwise, `-` can be used, and roots are at the left, heads are at the
/// right. `|` and `-` cannot be used together.
///
/// # Example:
///
/// ```
/// use drawdag::parse;
///
/// let edges = parse(r#"
///             E
///              \
///     C----B----A
///        /
///      D-
/// "#);
/// let expected = "{\"A\": {\"B\", \"E\"}, \"B\": {\"C\", \"D\"}, \"C\": {}, \"D\": {}, \"E\": {}}";
/// assert_eq!(format!("{:?}", edges), expected);
///
/// let edges = parse(r#"
///   A
///  /|
/// | B
/// E |
///   |\
///   C D
/// "#);
/// assert_eq!(format!("{:?}", edges), expected);
/// ```
pub fn parse(text: &str) -> BTreeMap<String, BTreeSet<String>> {
    use Direction::BottomTop;
    use Direction::LeftRight;

    // Detect direction.
    let direction = if "|:".chars().any(|c| text.contains(c)) {
        BottomTop
    } else {
        LeftRight
    };
    let lines: Vec<Vec<char>> = text.lines().map(|line| line.chars().collect()).collect();

    // (y, x) -> char. Return a space if (y, x) is out of range.
    let get = |y: isize, x: isize| -> char {
        if y < 0 || x < 0 {
            ' '
        } else {
            lines
                .get(y as usize)
                .cloned()
                .map_or(' ', |line| line.get(x as usize).cloned().unwrap_or(' '))
        }
    };

    // Like `get`, but concatenate left and right parts if they look like a word.
    let get_name = |y: isize, x: isize| -> String {
        (0..x)
            .rev()
            .map(|x| get(y, x))
            .take_while(|&ch| is_name(ch, direction))
            .collect::<Vec<_>>()
            .into_iter()
            .rev()
            .chain(
                (x..)
                    .map(|x| get(y, x))
                    .take_while(|&ch| is_name(ch, direction)),
            )
            .collect()
    };

    /// State used to visit the graph.
    #[derive(Eq, PartialEq, Ord, PartialOrd, Hash, Copy, Clone)]
    struct State {
        y: isize,
        x: isize,
        expected: &'static str,
        is_range: bool,
    }

    // Follow the ASCII edges at the given position.
    // Return a list of (parent, is_range).
    let get_parents = |y: isize, x: isize| -> Vec<(String, bool)> {
        let mut parents = Vec::new();
        let mut visited = HashSet::new();
        let mut visit = |state: State, to_visit: &mut Vec<State>| {
            if visited.insert(state) {
                let y = state.y;
                let x = state.x;
                let expected = state.expected;
                let ch = get(y, x);
                if is_name(ch, direction) && expected.contains('t') {
                    // t: text
                    parents.push((get_name(y, x), state.is_range));
                    return;
                }
                if !expected.contains(ch) {
                    return;
                }

                // Quickly construct a `State`.
                let is_range = state.is_range || ch == ':' || ch == '.';
                let s = |y, x, expected| State {
                    y,
                    x,
                    expected,
                    is_range,
                };

                match (ch, direction) {
                    (' ', _) => {}
                    ('|', BottomTop) | (':', BottomTop) => {
                        to_visit.push(s(y + 1, x - 1, "/"));
                        to_visit.push(s(y + 1, x, ":|/\\t"));
                        to_visit.push(s(y + 1, x + 1, "\\"));
                    }
                    ('\\', BottomTop) => {
                        to_visit.push(s(y + 1, x + 1, ":|\\t"));
                        to_visit.push(s(y + 1, x, ":|t"));
                    }
                    ('/', BottomTop) => {
                        to_visit.push(s(y + 1, x - 1, ":|/t"));
                        to_visit.push(s(y + 1, x, ":|t"));
                    }
                    ('-', LeftRight) | ('.', LeftRight) => {
                        to_visit.push(s(y - 1, x - 1, "\\"));
                        to_visit.push(s(y, x - 1, ".-/\\t"));
                        to_visit.push(s(y + 1, x - 1, "/"));
                    }
                    ('\\', LeftRight) => {
                        to_visit.push(s(y - 1, x - 1, ".-\\t"));
                        to_visit.push(s(y, x - 1, ".-t"));
                    }
                    ('/', LeftRight) => {
                        to_visit.push(s(y + 1, x - 1, ".-/t"));
                        to_visit.push(s(y, x - 1, ".-t"));
                    }
                    _ => unreachable!(),
                }
            }
        };

        let s = |y, x, expected| State {
            y,
            x,
            expected,
            is_range: false,
        };
        let mut to_visit: Vec<State> = match direction {
            BottomTop => [
                s(y + 1, x - 1, "/"),
                s(y + 1, x, "|:"),
                s(y + 1, x + 1, "\\"),
            ],
            LeftRight => [
                s(y - 1, x - 1, "\\"),
                s(y, x - 1, "-."),
                s(y + 1, x - 1, "/"),
            ],
        }
        .iter()
        .cloned()
        .filter(|state| state.expected.contains(get(state.y, state.x)))
        .collect();
        while let Some(state) = to_visit.pop() {
            visit(state, &mut to_visit);
        }

        parents
    };

    // Scan every character
    let mut edges: BTreeMap<String, BTreeSet<String>> = BTreeMap::new();
    for y in 0..lines.len() as isize {
        for x in 0..lines[y as usize].len() as isize {
            let ch = get(y, x);
            if is_name(ch, direction) {
                let name = get_name(y, x);
                edges.entry(name.clone()).or_default();
                for (parent, is_range) in get_parents(y, x) {
                    if !is_range {
                        edges.get_mut(&name).unwrap().insert(parent);
                    } else {
                        // Insert a chain of name -> parent. For example,
                        // name="D", parent="A", insert D -> C -> B -> A.

                        assert!(
                            parent.len() < name.len()
                                || (parent.len() == name.len() && parent < name),
                            "empty range: {:?} to {:?}",
                            parent,
                            name
                        );

                        let mut current: String = parent.clone();
                        loop {
                            let next = succ::str_succ(&current);
                            edges.entry(next.clone()).or_default().insert(current);

                            if next == name {
                                break;
                            }

                            assert!(
                                next.len() < name.len()
                                    || (next.len() == name.len() && next < name),
                                "mismatched range endpoints: {:?} to {:?}",
                                parent,
                                name
                            );

                            current = next;
                        }
                    }
                }
            }
            // Sanity check
            match (ch, direction) {
                ('-', BottomTop) => panic!("'-' is incompatible with BottomTop direction"),
                ('|', LeftRight) => panic!("'|' is incompatible with LeftRight direction"),
                _ => {}
            }
        }
    }

    edges
}

/// Commit the DAG by using the given commit function.
///
/// The commit function takes two arguments: Commit identity by the ASCII dag,
/// and parents defined by the commit function. The commit function returns the
/// identity of the committed change, and this function will use them as parents
/// passed into the future `commit_func` calls.
pub fn commit(
    dag: &BTreeMap<String, BTreeSet<String>>,
    mut commit_func: impl FnMut(String, Vec<Box<[u8]>>) -> Box<[u8]>,
) {
    let mut committed: BTreeMap<String, Box<[u8]>> = BTreeMap::new();

    while committed.len() < dag.len() {
        let mut made_progress = false;
        for (name, parents) in dag.iter() {
            if !committed.contains_key(name)
                && parents.iter().all(|name| committed.contains_key(name))
            {
                let parent_ids = parents.iter().map(|name| committed[name].clone()).collect();
                let new_id = commit_func(name.clone(), parent_ids);
                committed.insert(name.to_string(), new_id);
                made_progress = true;
            }
        }
        assert!(made_progress, "graph contains cycles");
    }
}

/// Parse the ASCII DAG and commit it. See [`parse`] and [`commit`] for details.
pub fn drawdag(text: &str, commit_func: impl FnMut(String, Vec<Box<[u8]>>) -> Box<[u8]>) {
    commit(&parse(text), commit_func)
}

fn is_name(ch: char, direction: Direction) -> bool {
    match (ch, direction) {
        ('.', Direction::BottomTop) => true,
        _ => ch.is_alphanumeric() || ",()_'\"".contains(ch),
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    struct CommitLog {
        log: String,
    }

    impl CommitLog {
        fn new() -> Self {
            Self { log: String::new() }
        }

        fn commit(&mut self, name: String, parents: Vec<Box<[u8]>>) -> Box<[u8]> {
            let new_id = self.log.chars().filter(|&ch| ch == '\n').count();
            let parents_str: Vec<String> = parents
                .into_iter()
                .map(|p| String::from_utf8(p.into_vec()).unwrap())
                .collect();
            self.log += &format!(
                "{}: {{ parents: {:?}, name: {} }}\n",
                new_id, parents_str, name
            );
            format!("{}", new_id).as_bytes().to_vec().into_boxed_slice()
        }
    }

    fn assert_drawdag(text: &str, expected: &str) {
        let mut log = CommitLog::new();
        drawdag(text, |n, p| log.commit(n, p));
        assert_eq!(log.log, expected);
    }

    /// Parse drawdag text, and return a list of strings as the parse result.
    /// Unlike `assert_drawdag`, `assert_eq!(d(t), e)` works with `cargo-fixeq`.
    fn p(text: &str) -> Vec<String> {
        parse(text)
            .into_iter()
            .map(|(k, vs)| {
                let vs = vs.into_iter().collect::<Vec<_>>().join(", ");
                format!("{} -> [{}]", k, vs)
            })
            .collect()
    }

    #[test]
    #[should_panic]
    fn test_drawdag_cycle1() {
        let mut log = CommitLog::new();
        drawdag("A-B B-A", |n, p| log.commit(n, p));
    }

    #[test]
    #[should_panic]
    fn test_drawdag_cycle2() {
        let mut log = CommitLog::new();
        drawdag("A-B-C-A", |n, p| log.commit(n, p));
    }

    #[test]
    #[should_panic]
    fn test_drawdag_mismatched_range1() {
        let mut log = CommitLog::new();
        drawdag("0..A", |n, p| log.commit(n, p));
    }

    #[test]
    #[should_panic]
    fn test_drawdag_mismatched_range2() {
        let mut log = CommitLog::new();
        drawdag("(09)..(A0)", |n, p| log.commit(n, p));
    }

    #[test]
    fn test_drawdag() {
        assert_drawdag(
            "A-C-B",
            r#"0: { parents: [], name: A }
1: { parents: ["0"], name: C }
2: { parents: ["1"], name: B }
"#,
        );

        assert_drawdag(
            r#"
    C-D-\     /--I--J--\
A-B------E-F-G-H--------K--L"#,
            r#"0: { parents: [], name: A }
1: { parents: ["0"], name: B }
2: { parents: [], name: C }
3: { parents: ["2"], name: D }
4: { parents: ["1", "3"], name: E }
5: { parents: ["4"], name: F }
6: { parents: ["5"], name: G }
7: { parents: ["6"], name: H }
8: { parents: ["6"], name: I }
9: { parents: ["8"], name: J }
10: { parents: ["7", "9"], name: K }
11: { parents: ["10"], name: L }
"#,
        );

        assert_drawdag(
            r#"
      G
      |
I D C F
 \ \| |
  H B E
   \|/
    A
"#,
            r#"0: { parents: [], name: A }
1: { parents: ["0"], name: B }
2: { parents: ["1"], name: C }
3: { parents: ["1"], name: D }
4: { parents: ["0"], name: E }
5: { parents: ["4"], name: F }
6: { parents: ["5"], name: G }
7: { parents: ["0"], name: H }
8: { parents: ["7"], name: I }
"#,
        );

        assert_drawdag(
            r#"
    A
   /|\
  H B E
 / /| |
I D C F
      |
      G
"#,
            r#"0: { parents: [], name: C }
1: { parents: [], name: D }
2: { parents: [], name: G }
3: { parents: [], name: I }
4: { parents: ["0", "1"], name: B }
5: { parents: ["2"], name: F }
6: { parents: ["3"], name: H }
7: { parents: ["5"], name: E }
8: { parents: ["4", "7", "6"], name: A }
"#,
        );
    }

    #[test]
    fn test_parse_range() {
        assert_eq!(p("A..D"), ["A -> []", "B -> [A]", "C -> [B]", "D -> [C]"]);
        assert_eq!(
            p(r"
             A1A,B23z,(9z)..A1A,B23z,(10c)
            "),
            [
                "A1A,B23z,(10a) -> [A1A,B23z,(9z)]",
                "A1A,B23z,(10b) -> [A1A,B23z,(10a)]",
                "A1A,B23z,(10c) -> [A1A,B23z,(10b)]",
                "A1A,B23z,(9z) -> []"
            ]
        );
        assert_eq!(
            p(r"
            B08
             :
            B04"),
            [
                "B04 -> []",
                "B05 -> [B04]",
                "B06 -> [B05]",
                "B07 -> [B06]",
                "B08 -> [B07]"
            ]
        );
        assert_eq!(
            p(r"
            B10
             | \
             :  C
             | /
            B08
             :
            B06"),
            [
                "B06 -> []",
                "B07 -> [B06]",
                "B08 -> [B07]",
                "B09 -> [B08]",
                "B10 -> [B09, C]",
                "C -> [B08]"
            ]
        );
        assert_eq!(
            p(r"
             AE
             | \
             :  C
             | /
             AB
             :
             X"),
            [
                "AA -> [Z]",
                "AB -> [AA]",
                "AC -> [AB]",
                "AD -> [AC]",
                "AE -> [AD, C]",
                "C -> [AB]",
                "X -> []",
                "Y -> [X]",
                "Z -> [Y]"
            ]
        );
    }

    #[test]
    fn test_parse_special_names() {
        assert_eq!(
            p("ancestor(desc(\"D\"),desc('_A'))--B"),
            [
                "B -> [ancestor(desc(\"D\"),desc('_A'))]",
                "ancestor(desc(\"D\"),desc('_A')) -> []"
            ]
        );
        assert_eq!(
            p(r#"
                B
                |
                .
              "#),
            [". -> []", "B -> [.]"]
        );
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

/// String successor by incrementing characters.
/// Similar to Ruby's `String#succ` [1], except if all characters are
/// non-alphanumeric, this function appends `1` to the end.
///
/// [1]: https://ruby-doc.org/core-3.1.0/String.html#method-i-next
pub(crate) fn str_succ(s: &str) -> String {
    let mut chars: Vec<char> = s.chars().collect();

    // The first character to be incremented is the rightmost alphanumeric.
    let index: Option<usize> = chars
        .iter()
        .enumerate()
        .rev()
        .filter_map(|(i, c)| CharRange::from_char(*c).map(|_| i))
        .next();
    match index {
        Some(index) => {
            let mut carry: char = '1';
            for i in (0..=index).rev() {
                let ch = chars[i];
                let range = match CharRange::from_char(ch) {
                    None => {
                        chars.insert(i + 1, carry);
                        break;
                    }
                    Some(range) => range,
                };
                carry = range.carry();
                let (start, end) = range.bound();
                if ch == end {
                    chars[i] = start;
                    if i == 0 {
                        chars.insert(0, range.carry());
                    }
                } else {
                    chars[i] = ((ch as u8) + 1) as char;
                    break;
                }
            }
            chars.into_iter().collect()
        }
        None => format!("{}1", s),
    }
}

/// Range of an ASCII char.
#[derive(Copy, Clone, PartialEq)]
enum CharRange {
    Digit,
    LowerLetter,
    UpperLetter,
}

impl CharRange {
    const ALL: &'static [Self] = &[Self::Digit, Self::LowerLetter, Self::UpperLetter];

    /// Get the start and end char in the range.
    fn bound(self) -> (char, char) {
        match self {
            CharRange::Digit => ('0', '9'),
            CharRange::LowerLetter => ('a', 'z'),
            CharRange::UpperLetter => ('A', 'Z'),
        }
    }

    /// Get the char used when carrying to a new char.
    fn carry(self) -> char {
        match self {
            CharRange::Digit => '1',
            CharRange::LowerLetter => 'a',
            CharRange::UpperLetter => 'A',
        }
    }

    /// Convert a char to `CharRange` if it's in a known range.
    fn from_char(ch: char) -> Option<CharRange> {
        Self::ALL.iter().copied().find(|range| {
            let (start, end) = range.bound();
            ch >= start && ch <= end
        })
    }
}

#[cfg(test)]
#[test]
fn test_str_succ() {
    // Alphanumeric cases are consistent with Ruby.
    assert_eq!(str_succ("0"), "1");
    assert_eq!(str_succ("1"), "2");
    assert_eq!(str_succ("9"), "10");
    assert_eq!(str_succ("233"), "234");
    assert_eq!(str_succ("999"), "1000");
    assert_eq!(str_succ("0099"), "0100");
    assert_eq!(str_succ("1180591620717411303424"), "1180591620717411303425");

    assert_eq!(str_succ("a-99"), "a-100");
    assert_eq!(str_succ("aa99"), "ab00");
    assert_eq!(str_succ("aa99.."), "ab00..");

    assert_eq!(str_succ("a"), "b");
    assert_eq!(str_succ("C"), "D");
    assert_eq!(str_succ("ASDF"), "ASDG");
    assert_eq!(str_succ("zz"), "aaa");

    assert_eq!(str_succ("Zz"), "AAa");
    assert_eq!(str_succ("9z9Z"), "10a0A");

    // Non-alphanumeric cases are different from Ruby.
    assert_eq!(str_succ(""), "1");
    assert_eq!(str_succ("."), ".1");
}