open Ast

let parse query = 
  let lexbuf = Lexing.from_string query in
  Parser.main Lexer.read_token lexbuf

let equal_ast ast1 ast2 = 
  ast1 = ast2

let query_testable =
  Alcotest.testable Ast.pp_query equal_ast

let test_simple_select () =
  let query = "SELECT a FROM t" in
  let q1 = parse query in
  let ast1 = Query(
    Select(
      [
        Column(
          Ref("a"),
          None
        )
      ], 
      Some(
        [
          Table("t")
        ]
      ), 
      None
    )
  ) in
  Alcotest.(check query_testable) query q1 ast1;
  
  let query = "SELECT * FROM t" in
  let q2 = parse query in
  let ast2 = Query(
    Select(
      [
        Asterisk
      ], 
      Some(
        [
          Table("t")
        ]
      ), 
      None
    )
  ) in
  Alcotest.(check query_testable) query q2 ast2 ;

  let query = "SELECT 'a'" in
  let q3 = parse query in
  let ast3 = Query(
    Select(
      [
        Column(
          StringLiteral("a"),
          None
        )
      ],
      None,
      None
    )
  ) in
  Alcotest.(check query_testable) query q3 ast3;
  
  let q4 = parse "SELECT DATE '2024-12-25' AS date" in
  let ast4 = Query(
    Select(
      [
        Column(
          DateLiteral("2024-12-25"),
          Some(
            As("date")
          )
        )
      ],
      None,
      None
    )
  ) in
  Alcotest.(check query_testable) "OK" q4 ast4

let test_default_join () =
  let q1 = parse "SELECT a FROM t1 JOIN t2 ON b = c" in
  let ast1 = Query(
    Select(
      [
        Column(
            Ref("a"),
            None
        )
      ],
      Some(
        [
          Join(
            Table("t1"), 
            Left, 
            Table("t2"),
            Some(
              Condition(
                Ref("b"),
                Comparison(
                  Equals, 
                  Ref("c"))
              )
            )
          )
        ]
      ),
      None
    )
  ) in
  Alcotest.(check query_testable) "Ok" q1 ast1

let test_left_join () = 
  let query =  "SELECT a FROM t1 LEFT JOIN t2 ON a = b" in
  let q1 = parse query in
  let ast1 = Query(
    Select(
      [
        Column(
          Ref("a"),
          None
        )
      ], 
      Some(
        [
          Join(
            Table("t1"), 
            Left, 
            Table("t2"),
            Some(
              Condition(
                Ref("a"),
                Comparison(Equals, Ref("b"))
              )
            )
          )
        ]
      ),
      None
    )
  ) in
  Alcotest.(check query_testable) query q1 ast1
  
let test_right_join () = 
  let q1 = parse "SELECT a FROM t1 RIGHT JOIN t2 ON a = b" in
  let ast1 = Query(
    Select(
      [
        Column(
          Ref("a"),
          None
        )
      ], 
      Some(
        [
          Join(
            Table("t1"), 
            Right, 
            Table("t2"),
            Some(
              Condition(
                Ref("a"),
                Comparison(Equals, Ref("b"))
              )
            )
          )
        ]
      ),
      None
    )
  ) in
  Alcotest.(check query_testable) "Ok" q1 ast1

let test_inner_join () = 
  let q1 = parse "SELECT a FROM t1 INNER JOIN t2 ON a = b" in
  let ast1 = Query(
    Select(
      [
        Column(
          Ref("a"),
          None
        )
      ], 
      Some(
        [
          Join(
            Table("t1"), 
            Inner, 
            Table("t2"),
            Some(
              Condition(
                Ref("a"),
                Comparison(Equals, Ref("b"))
              )
            )
          )
        ]
      ),
      None
    )
  ) in
  Alcotest.(check query_testable) "Ok" q1 ast1

let test_union_join () = 
  let q1 = parse "SELECT a FROM t1 UNION JOIN t2" in
  let ast1 = Query(
    Select(
      [
        Column(
          Ref("a"),
          None
        )
      ], 
      Some(
        [
          Join(
            Table("t1"), 
            Union, 
            Table("t2"),
            None
          )
        ]
      ),
      None
    )
  ) in
  Alcotest.(check query_testable) "Ok" q1 ast1

let test_cross_join () = 
  let q1 = parse "SELECT a FROM t1 CROSS JOIN t2" in
  let ast1 = Query(
    Select(
      [
        Column(
          Ref("a"),
          None
        )
      ], 
      Some(
        [
          Join(
            Table("t1"), 
            Cross, 
            Table("t2"),
            None
          )
        ]
      ),
      None
    )
  ) in
  Alcotest.(check query_testable) "Ok" q1 ast1

let test_natural_join () = 
  let q1 = parse "SELECT a FROM t1 NATURAL JOIN t2" in
  let ast1 = Query(
    Select(
      [
        Column(
          Ref("a"),
          None
        )
      ],
      Some(
        [
          Join(
            Table("t1"), 
            Natural, 
            Table("t2"),
            None
          )
        ]
      ),
      None
    )
  ) in
  Alcotest.(check query_testable) "Ok" q1 ast1

let test_join_join () = 
  let q1 = parse "SELECT a FROM t1 JOIN t2 ON a = b JOIN t3 ON a = c" in
  let ast1 = Query(
    Select(
      [
        Column(
          Ref("a"),
          None
        )
      ], 
      Some([
        Join(
          Join(
            Table("t1"), 
            Left, 
            Table("t2"),
            Some(
              Condition(
                Ref("a"),
                Comparison(Equals, Ref("b"))
              )
            ) 
          ),
          Left,
          Table("t3"),
          Some(
            Condition(
              Ref("a"),
              Comparison(Equals, Ref("c"))
            )
          )
        )
      ]),
      None
    )
  ) in
  Alcotest.(check query_testable) "Ok" q1 ast1

let test_where_equals () = 
  let q1 = parse "SELECT a FROM t1 WHERE a = a OR a = b" in
  let ast1 = Query(
    Select(
      [
        Column(
          Ref("a"),
          None
        )
      ],
      Some(
        [
          Table("t1")
        ]
      ),
      None
    )
  ) in
  Alcotest.(check query_testable) "Ok" q1 ast1

let simple_select_set  = [ ("Equals", `Quick, test_simple_select) ]

let simple_join_set = [
  ("Default Join", `Quick, test_default_join);
  ("Left Join", `Quick, test_left_join);
  ("Right Join", `Quick, test_right_join);
  ("Inner Join", `Quick, test_inner_join);
  ("Union Join", `Quick, test_union_join);
  ("Cross Join", `Quick, test_cross_join);
  ("Natural Join", `Quick, test_natural_join)
]

let multiple_joins_set = [ 
  ("Join Join", `Quick, test_join_join)
]

let where_clauses_set = [
  ("Where Equals", `Quick, test_where_equals)
]

let () =
  Alcotest.run "Ast tests"
    [
      ("Simple Selects", simple_select_set); 
      ("Simple Joins", simple_join_set);
      ("Multiple Joins", multiple_joins_set);
      ("Where Clauses", where_clauses_set)
    ]