(*
 * 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.
 *)

open Core
open OUnit2
open Server
open ServerTest

let create_relative_source_path ~root ~relative =
  PyrePath.create_relative ~root ~relative |> SourcePath.create


let test_initialize context =
  let internal_state = ref "uninitiailzed" in
  let build_system_initializer =
    let initialize () =
      internal_state := "initialized";
      Lwt.return (BuildSystem.create_for_testing ())
    in
    let load () = failwith "saved state loading is not supported" in
    let cleanup () = Lwt.return_unit in
    BuildSystem.Initializer.create_for_testing ~initialize ~load ~cleanup ()
  in
  let test_initialize _ =
    (* Verify that the build system has indeed been initiailzed. *)
    assert_equal ~ctxt:context ~cmp:String.equal ~printer:Fn.id "initialized" !internal_state;
    Lwt.return_unit
  in
  ScratchProject.setup
    ~context
    ~include_typeshed_stubs:false
    ~include_helper_builtins:false
    ~build_system_initializer
    []
  |> ScratchProject.test_server_with ~f:test_initialize


let test_cleanup context =
  let internal_state = ref "uncleaned" in
  let build_system_initializer =
    let initialize () = Lwt.return (BuildSystem.create_for_testing ()) in
    let load () = failwith "saved state loading is not supported" in
    let cleanup () =
      internal_state := "cleaned";
      Lwt.return_unit
    in
    BuildSystem.Initializer.create_for_testing ~initialize ~load ~cleanup ()
  in
  let open Lwt.Infix in
  let start_options =
    ScratchProject.setup
      ~context
      ~build_system_initializer
      ~include_typeshed_stubs:false
      ~include_helper_builtins:false
      []
    |> ScratchProject.start_options_of
  in

  Start.start_server
    start_options
    ~on_exception:(fun exn -> raise exn)
    ~on_started:(fun _ _ ->
      (* Shutdown the server immediately after it is started. *)
      Lwt.return_unit)
  >>= fun _ ->
  (* Verify that the build system has indeed been cleaned up. *)
  assert_equal ~ctxt:context ~cmp:String.equal ~printer:Fn.id "cleaned" !internal_state;
  Lwt.return_unit


let test_type_errors context =
  let test_source_path = PyrePath.create_absolute "/foo/test.py" |> SourcePath.create in
  let test_artifact_path =
    (* The real value will be deterimend once the server starts. *)
    ref (PyrePath.create_absolute "uninitialized" |> ArtifactPath.create)
  in
  let build_system_initializer =
    let initialize () =
      let lookup_source path =
        if [%compare.equal: ArtifactPath.t] path !test_artifact_path then
          Some test_source_path
        else
          None
      in
      let lookup_artifact path =
        if [%compare.equal: SourcePath.t] path test_source_path then
          [!test_artifact_path]
        else
          []
      in
      Lwt.return (BuildSystem.create_for_testing ~lookup_source ~lookup_artifact ())
    in
    let load () = failwith "saved state loading is not supported" in
    let cleanup () = Lwt.return_unit in
    BuildSystem.Initializer.create_for_testing ~initialize ~load ~cleanup ()
  in
  let test_type_errors client =
    let open Lwt.Infix in
    let global_root =
      Client.get_server_properties client
      |> fun { ServerProperties.configuration = { Configuration.Analysis.project_root; _ }; _ } ->
      project_root
    in
    test_artifact_path := Test.relative_artifact_path ~root:global_root ~relative:"test.py";
    let test_error =
      Analysis.AnalysisError.Instantiated.of_yojson
        (`Assoc
          [
            "line", `Int 1;
            "column", `Int 0;
            "stop_line", `Int 1;
            "stop_column", `Int 11;
            "path", `String "/foo/test.py";
            "code", `Int (-1);
            "name", `String "Revealed type";
            ( "description",
              `String
                "Revealed type [-1]: Revealed type for `42` is `typing_extensions.Literal[42]`." );
            ( "long_description",
              `String
                "Revealed type [-1]: Revealed type for `42` is `typing_extensions.Literal[42]`." );
            ( "concise_description",
              `String
                "Revealed type [-1]: Revealed type for `42` is `typing_extensions.Literal[42]`." );
            "define", `String "test.$toplevel";
          ])
      |> Result.ok_or_failwith
    in
    let test2_error =
      (* `test2.py` is intentionally not tracked by the build system. The expected behavior here is
         to show its original path. *)
      let test2_artifact_path =
        create_relative_source_path ~root:global_root ~relative:"test2.py"
      in
      Analysis.AnalysisError.Instantiated.of_yojson
        (`Assoc
          [
            "line", `Int 1;
            "column", `Int 0;
            "stop_line", `Int 1;
            "stop_column", `Int 11;
            "path", `String (SourcePath.raw test2_artifact_path |> PyrePath.absolute);
            "code", `Int (-1);
            "name", `String "Revealed type";
            ( "description",
              `String
                "Revealed type [-1]: Revealed type for `43` is `typing_extensions.Literal[43]`." );
            ( "long_description",
              `String
                "Revealed type [-1]: Revealed type for `43` is `typing_extensions.Literal[43]`." );
            ( "concise_description",
              `String
                "Revealed type [-1]: Revealed type for `43` is `typing_extensions.Literal[43]`." );
            "define", `String "test2.$toplevel";
          ])
      |> Result.ok_or_failwith
    in
    Client.assert_response
      client
      ~request:(Request.DisplayTypeError [])
      ~expected:(create_type_error_response [test_error; test2_error])
    >>= fun () ->
    Client.assert_response
      client
      ~request:(Request.DisplayTypeError ["/foo/test.py"])
      ~expected:(create_type_error_response [test_error])
  in
  ScratchProject.setup
    ~context
    ~include_typeshed_stubs:false
    ~include_helper_builtins:false
    ~build_system_initializer
    ["test.py", "reveal_type(42)"; "test2.py", "reveal_type(43)"]
  |> ScratchProject.test_server_with ~f:test_type_errors


let test_type_errors_in_multiple_artifacts context =
  let test_source_path = PyrePath.create_absolute "/foo/test.py" |> SourcePath.create in
  let test_artifact_path0 =
    (* The real value will be deterimend once the server starts. *)
    ref (PyrePath.create_absolute "uninitialized" |> ArtifactPath.create)
  in
  let test_artifact_path1 =
    (* The real value will be deterimend once the server starts. *)
    ref (PyrePath.create_absolute "uninitialized" |> ArtifactPath.create)
  in
  let build_system_initializer =
    let initialize () =
      (* We map the queried source path to both aritfact paths *)
      let lookup_source path =
        if
          [%compare.equal: ArtifactPath.t] path !test_artifact_path0
          || [%compare.equal: ArtifactPath.t] path !test_artifact_path1
        then
          Some test_source_path
        else
          None
      in
      let lookup_artifact path =
        if [%compare.equal: SourcePath.t] path test_source_path then
          [!test_artifact_path0; !test_artifact_path1]
        else
          []
      in
      Lwt.return (BuildSystem.create_for_testing ~lookup_source ~lookup_artifact ())
    in
    let load () = failwith "saved state loading is not supported" in
    let cleanup () = Lwt.return_unit in
    BuildSystem.Initializer.create_for_testing ~initialize ~load ~cleanup ()
  in
  let test_type_errors client =
    let open Lwt.Infix in
    let global_root =
      Client.get_server_properties client
      |> fun { ServerProperties.configuration = { Configuration.Analysis.project_root; _ }; _ } ->
      project_root
    in
    test_artifact_path0 := Test.relative_artifact_path ~root:global_root ~relative:"foo/test.py";
    test_artifact_path1 := Test.relative_artifact_path ~root:global_root ~relative:"bar/test.py";
    Client.send_request client (Request.DisplayTypeError ["/foo/test.py"])
    >>= fun raw_response ->
    match Yojson.Safe.from_string raw_response with
    | `List [`String "TypeErrors"; `Assoc [("errors", `List errors); _]] ->
        (* Given that `/foo/test.py` is mapped to two artifact paths, the client should see type
           errors in both files. *)
        assert_equal ~ctxt:context ~cmp:Int.equal ~printer:Int.to_string 2 (List.length errors);
        Lwt.return_unit
    | _ ->
        let message = Format.sprintf "Unexpected response message: %s" raw_response in
        assert_failure message
  in
  ScratchProject.setup
    ~context
    ~include_typeshed_stubs:false
    ~include_helper_builtins:false
    ~build_system_initializer
    ["foo/test.py", "reveal_type(42)"; "bar/test.py", "reveal_type(43)"]
  |> ScratchProject.test_server_with ~f:test_type_errors


let test_update context =
  let internal_state = ref "unupdated" in
  let test_source_path = PyrePath.create_absolute "/foo/test.py" |> SourcePath.create in
  let test_artifact_path =
    (* The real value will be deterimend once the server starts. *)
    ref (PyrePath.create_absolute "uninitialized" |> ArtifactPath.create)
  in
  let build_system_initializer =
    let initialize () =
      let lookup_source path =
        if [%compare.equal: ArtifactPath.t] path !test_artifact_path then
          Some test_source_path
        else
          None
      in
      let lookup_artifact path =
        if [%compare.equal: SourcePath.t] path test_source_path then
          [!test_artifact_path]
        else
          []
      in
      let update actual_path_events =
        assert_equal
          ~ctxt:context
          ~cmp:[%compare.equal: SourcePath.Event.t list]
          ~printer:(fun paths -> [%sexp_of: SourcePath.Event.t list] paths |> Sexp.to_string_hum)
          [(test_source_path |> SourcePath.Event.(create ~kind:Kind.CreatedOrChanged))]
          actual_path_events;
        internal_state := "updated";
        Lwt.return []
      in
      Lwt.return (BuildSystem.create_for_testing ~update ~lookup_source ~lookup_artifact ())
    in
    let load () = failwith "saved state loading is not supported" in
    let cleanup () = Lwt.return_unit in
    BuildSystem.Initializer.create_for_testing ~initialize ~load ~cleanup ()
  in
  let test_update client =
    let open Lwt.Infix in
    let root =
      Client.get_server_properties client
      |> fun { ServerProperties.configuration = { Configuration.Analysis.project_root; _ }; _ } ->
      project_root
    in
    test_artifact_path := Test.relative_artifact_path ~root ~relative:"test.py";

    ArtifactPath.raw !test_artifact_path |> File.create ~content:"reveal_type(42)" |> File.write;
    Client.send_request
      client
      (Request.IncrementalUpdate [SourcePath.raw test_source_path |> PyrePath.absolute])
    >>= fun _ ->
    (* Verify that the build system has indeed been updated. *)
    assert_equal ~ctxt:context ~cmp:String.equal ~printer:Fn.id "updated" !internal_state;
    (* Verify that recheck has indeed happened. *)
    let expected_error =
      Analysis.AnalysisError.Instantiated.of_yojson
        (`Assoc
          [
            "line", `Int 1;
            "column", `Int 0;
            "stop_line", `Int 1;
            "stop_column", `Int 11;
            "path", `String "/foo/test.py";
            "code", `Int (-1);
            "name", `String "Revealed type";
            ( "description",
              `String
                "Revealed type [-1]: Revealed type for `42` is `typing_extensions.Literal[42]`." );
            ( "long_description",
              `String
                "Revealed type [-1]: Revealed type for `42` is `typing_extensions.Literal[42]`." );
            ( "concise_description",
              `String
                "Revealed type [-1]: Revealed type for `42` is `typing_extensions.Literal[42]`." );
            "define", `String "test.$toplevel";
          ])
      |> Result.ok_or_failwith
    in
    Client.assert_response
      client
      ~request:(Request.DisplayTypeError [])
      ~expected:(create_type_error_response [expected_error])
    >>= fun () -> Lwt.return_unit
  in
  ScratchProject.setup
    ~context
    ~include_typeshed_stubs:false
    ~include_helper_builtins:false
    ~build_system_initializer
    ["test.py", "reveal_type(True)"]
  |> ScratchProject.test_server_with ~f:test_update


let assert_source_path ~context ~expected actual =
  assert_equal
    ~ctxt:context
    ~cmp:[%compare.equal: SourcePath.t option]
    ~printer:(Option.value_map ~default:"NONE" ~f:SourcePath.show)
    expected
    actual


let assert_artifact_paths ~context ~expected actual =
  assert_equal
    ~ctxt:context
    ~cmp:[%compare.equal: ArtifactPath.t list]
    ~printer:(List.to_string ~f:ArtifactPath.show)
    expected
    actual


let test_buck_update context =
  let assert_source_path = assert_source_path ~context in
  let assert_artifact_paths = assert_artifact_paths ~context in
  let source_root = bracket_tmpdir context |> PyrePath.create_absolute in
  let artifact_root = bracket_tmpdir context |> PyrePath.create_absolute in

  let get_buck_build_system () =
    let interface =
      (* Here's the set up: we have 2 files, `foo/bar.py` and `foo/baz.py`. If `is_rebuild` is
         false, we'll only include `foo/bar.py` in the target. If `is_rebuild` is true, we'll
         include both files. The `is_rebuild` flag is initially false but will be set to true after
         the first build. This setup emulates an incremental Buck update where the user edits the
         TARGET file to include another source in the target. *)
      let is_rebuild = ref false in
      let construct_build_map _ =
        let build_mappings =
          if !is_rebuild then
            ["bar.py", "foo/bar.py"; "baz.py", "foo/baz.py"]
          else (
            is_rebuild := true;
            ["bar.py", "foo/bar.py"])
        in
        Buck.(
          BuildMap.(create (Partial.of_alist_exn build_mappings)) |> Interface.WithMetadata.create)
        |> Lwt.return
      in
      Buck.Interface.Eager.create_for_testing ~construct_build_map ()
    in
    let builder = Buck.Builder.Eager.create ~source_root ~artifact_root interface in
    BuildSystem.Initializer.buck ~builder ~artifact_root ~targets:["//foo:target"] ()
    |> BuildSystem.Initializer.run
  in
  let open Lwt.Infix in
  get_buck_build_system ()
  >>= fun buck_build_system ->
  let bar_source = create_relative_source_path ~root:source_root ~relative:"foo/bar.py" in
  let bar_artifact = Test.relative_artifact_path ~root:artifact_root ~relative:"bar.py" in
  let baz_source = create_relative_source_path ~root:source_root ~relative:"foo/baz.py" in
  let baz_artifact = Test.relative_artifact_path ~root:artifact_root ~relative:"baz.py" in

  (* Initially, we build bar.py but not baz.py. *)
  assert_source_path
    ~expected:(Some bar_source)
    (BuildSystem.lookup_source buck_build_system bar_artifact);
  assert_source_path ~expected:None (BuildSystem.lookup_source buck_build_system baz_artifact);
  assert_artifact_paths
    ~expected:[bar_artifact]
    (BuildSystem.lookup_artifact buck_build_system bar_source);
  assert_artifact_paths ~expected:[] (BuildSystem.lookup_artifact buck_build_system baz_source);

  (* Rebuild the project. The fake TARGET file is needed to force a full rebuild. *)
  let fake_target_file =
    create_relative_source_path ~root:source_root ~relative:"TARGETS"
    |> SourcePath.Event.(create ~kind:Kind.CreatedOrChanged)
  in
  BuildSystem.update
    buck_build_system
    [
      (bar_source |> SourcePath.Event.(create ~kind:Kind.CreatedOrChanged));
      (baz_source |> SourcePath.Event.(create ~kind:Kind.Deleted));
      fake_target_file;
    ]
  >>= fun _ ->
  (* After the rebuild, both bar.py and baz.py should be included in build map. *)
  assert_source_path
    ~expected:(Some bar_source)
    (BuildSystem.lookup_source buck_build_system bar_artifact);
  assert_source_path
    ~expected:(Some baz_source)
    (BuildSystem.lookup_source buck_build_system baz_artifact);
  assert_artifact_paths
    ~expected:[bar_artifact]
    (BuildSystem.lookup_artifact buck_build_system bar_source);
  assert_artifact_paths
    ~expected:[baz_artifact]
    (BuildSystem.lookup_artifact buck_build_system baz_source);

  Lwt.return_unit


let assert_paths_no_order ~context ~expected actual =
  let compare = [%compare: ArtifactPath.Event.t] in
  assert_equal
    ~ctxt:context
    ~cmp:[%compare.equal: ArtifactPath.Event.t list]
    ~printer:(fun paths -> [%sexp_of: ArtifactPath.Event.t list] paths |> Sexp.to_string_hum)
    (List.sort ~compare expected)
    (List.sort ~compare actual)


let test_buck_update_without_rebuild context =
  let assert_paths_no_order = assert_paths_no_order ~context in
  let source_root = bracket_tmpdir context |> PyrePath.create_absolute in
  let artifact_root = bracket_tmpdir context |> PyrePath.create_absolute in

  let get_buck_build_system () =
    let interface =
      let is_rebuild = ref false in
      let construct_build_map _ =
        let build_mappings =
          if not !is_rebuild then (
            is_rebuild := true;
            ["bar.py", "foo/bar.py"; "baz.py", "foo/baz.py"])
          else
            assert_failure
              "Build map construction is not expected to be invoked again after the initial build"
        in
        Buck.(
          BuildMap.(create (Partial.of_alist_exn build_mappings)) |> Interface.WithMetadata.create)
        |> Lwt.return
      in
      Buck.Interface.Eager.create_for_testing ~construct_build_map ()
    in
    let builder = Buck.Builder.Eager.create ~source_root ~artifact_root interface in
    BuildSystem.Initializer.buck ~builder ~artifact_root ~targets:["//foo:target"] ()
    |> BuildSystem.Initializer.run
  in
  let open Lwt.Infix in
  get_buck_build_system ()
  >>= fun buck_build_system ->
  let bar_source = create_relative_source_path ~root:source_root ~relative:"foo/bar.py" in
  let baz_source = create_relative_source_path ~root:source_root ~relative:"foo/baz.py" in
  File.create (SourcePath.raw bar_source) ~content:"" |> File.write;
  File.create (SourcePath.raw baz_source) ~content:"" |> File.write;
  BuildSystem.update
    buck_build_system
    [
      (bar_source |> SourcePath.Event.(create ~kind:Kind.CreatedOrChanged));
      (baz_source |> SourcePath.Event.(create ~kind:Kind.CreatedOrChanged));
    ]
  >>= fun changed_artifacts ->
  assert_paths_no_order changed_artifacts ~expected:[];
  (* After the rebuild, both bar.py and baz.py should be included in build map. *)
  let bar_artifact = Test.relative_artifact_path ~root:artifact_root ~relative:"bar.py" in
  let baz_artifact = Test.relative_artifact_path ~root:artifact_root ~relative:"baz.py" in
  assert_source_path
    ~context
    ~expected:(Some bar_source)
    (BuildSystem.lookup_source buck_build_system bar_artifact);
  assert_source_path
    ~context
    ~expected:(Some baz_source)
    (BuildSystem.lookup_source buck_build_system baz_artifact);
  Lwt.return_unit


let test_unwatched_dependency_no_failure_on_initialize context =
  let bucket_root = bracket_tmpdir context |> PyrePath.create_absolute in
  let bucket = "BUCKET" in
  let wheel_root = bracket_tmpdir context |> PyrePath.create_absolute in
  let checksum_path = "CHECKSUM" in
  let open Lwt.Infix in
  let test_initializer =
    BuildSystem.Initializer.track_unwatched_dependency
      {
        Configuration.UnwatchedDependency.change_indicator =
          { Configuration.ChangeIndicator.root = bucket_root; relative = bucket };
        files = { Configuration.UnwatchedFiles.root = wheel_root; checksum_path };
      }
  in
  BuildSystem.Initializer.run test_initializer
  >>= fun build_system ->
  (* Initialization should not crash, even when the checksum path does not exist *)
  let bucket_path =
    create_relative_source_path ~root:bucket_root ~relative:bucket
    |> SourcePath.Event.(create ~kind:Kind.CreatedOrChanged)
  in
  Lwt.catch
    (fun () ->
      (* Update should crash, if the checksum path does not exist *)
      BuildSystem.update build_system [bucket_path]
      >>= fun _ -> assert_failure "should not reach here")
    (function
      | ChecksumMap.LoadError _ -> Lwt.return_unit
      | exn ->
          let exn = Exception.wrap exn in
          assert_failure (Format.sprintf "wrong exception raised: `%s`" (Exception.to_string exn)))


let test_unwatched_dependency_update context =
  let assert_paths_no_order = assert_paths_no_order ~context in
  let bucket_root = bracket_tmpdir context |> PyrePath.create_absolute in
  let bucket = "BUCKET" in
  let bucket_path_event =
    create_relative_source_path ~root:bucket_root ~relative:bucket
    |> SourcePath.Event.(create ~kind:Kind.CreatedOrChanged)
  in
  let wheel_root = bracket_tmpdir context |> PyrePath.create_absolute in
  let checksum_path = "CHECKSUM" in
  let checksum_full_path = PyrePath.create_relative ~root:wheel_root ~relative:checksum_path in
  let content =
    {|
      {
        "a.py": "checksum0",
        "b/c.py": "checksum1",
        "d/e/f.pyi": "checksum2"
      }
    |}
  in
  File.create checksum_full_path ~content |> File.write;
  let open Lwt.Infix in
  let instagram_initializer =
    BuildSystem.Initializer.track_unwatched_dependency
      {
        Configuration.UnwatchedDependency.change_indicator =
          { Configuration.ChangeIndicator.root = bucket_root; relative = bucket };
        files = { Configuration.UnwatchedFiles.root = wheel_root; checksum_path };
      }
  in
  BuildSystem.Initializer.run instagram_initializer
  >>= fun build_system ->
  let test_path_event =
    PyrePath.create_absolute "/some/source/file.py"
    |> SourcePath.create
    |> SourcePath.Event.(create ~kind:Kind.CreatedOrChanged)
  in
  (* Normal update does not yield additional changed paths. *)
  BuildSystem.update build_system [test_path_event]
  >>= fun updated ->
  assert_paths_no_order updated ~expected:[];

  (* Touching the change indicator but not the checksum file does not yield additional changed
     paths. *)
  BuildSystem.update build_system [bucket_path_event]
  >>= fun updated ->
  assert_paths_no_order updated ~expected:[];

  (* Checksum file update will lead to additional changed paths. *)
  let content =
    {|
      {
        "a.py": "checksum3",
        "d/e/f.pyi": "checksum2",
        "g.py": "checksum4"
      }
    |}
  in
  File.create checksum_full_path ~content |> File.write;
  BuildSystem.update build_system [test_path_event; bucket_path_event]
  >>= fun updated ->
  assert_paths_no_order
    updated
    ~expected:
      [
        (Test.relative_artifact_path ~root:wheel_root ~relative:"a.py"
        |> ArtifactPath.Event.(create ~kind:Kind.CreatedOrChanged));
        (Test.relative_artifact_path ~root:wheel_root ~relative:"b/c.py"
        |> ArtifactPath.Event.(create ~kind:Kind.Deleted));
        (Test.relative_artifact_path ~root:wheel_root ~relative:"g.py"
        |> ArtifactPath.Event.(create ~kind:Kind.CreatedOrChanged));
      ];

  Lwt.return_unit


let () =
  "build_system_test"
  >::: [
         "initialize" >:: OUnitLwt.lwt_wrapper test_initialize;
         "cleanup" >:: OUnitLwt.lwt_wrapper test_cleanup;
         "type_errors" >:: OUnitLwt.lwt_wrapper test_type_errors;
         "type_errors_in_multiple_artifacts"
         >:: OUnitLwt.lwt_wrapper test_type_errors_in_multiple_artifacts;
         "update" >:: OUnitLwt.lwt_wrapper test_update;
         "buck_update" >:: OUnitLwt.lwt_wrapper test_buck_update;
         "buck_update_without_rebuild" >:: OUnitLwt.lwt_wrapper test_buck_update_without_rebuild;
         "unwatched_dependency_no_failure_on_initialize"
         >:: OUnitLwt.lwt_wrapper test_unwatched_dependency_no_failure_on_initialize;
         "unwatched_dependency_update" >:: OUnitLwt.lwt_wrapper test_unwatched_dependency_update;
       ]
  |> Test.run
