// Copyright 2018 Google Inc. All rights reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package main

import (
 "bytes"
 "fmt"
 "hash/crc32"
 "os"
 "path/filepath"
 "strconv"
 "strings"
 "testing"
 "time"

 "android/soong/jar"
 "android/soong/third_party/zip"
)

type testZipEntry struct {
 name      string
 mode      os.FileMode
 data      []byte
 method    uint16
 timestamp time.Time
 extra     []byte
}

var (
 A     = testZipEntry{"A", 0755, []byte("foo"), zip.Deflate, jar.DefaultTime, nil}
 a     = testZipEntry{"a", 0755, []byte("foo"), zip.Deflate, jar.DefaultTime, nil}
 a2    = testZipEntry{"a", 0755, []byte("FOO2"), zip.Deflate, jar.DefaultTime, nil}
 a3    = testZipEntry{"a", 0755, []byte("Foo3"), zip.Deflate, jar.DefaultTime, nil}
 bDir  = testZipEntry{"b/", os.ModeDir | 0755, nil, zip.Deflate, jar.DefaultTime, nil}
 bbDir = testZipEntry{"b/b/", os.ModeDir | 0755, nil, zip.Deflate, jar.DefaultTime, nil}
 bbb   = testZipEntry{"b/b/b", 0755, nil, zip.Deflate, jar.DefaultTime, nil}
 ba    = testZipEntry{"b/a", 0755, []byte("foo"), zip.Deflate, jar.DefaultTime, nil}
 bc    = testZipEntry{"b/c", 0755, []byte("bar"), zip.Deflate, jar.DefaultTime, nil}
 bd    = testZipEntry{"b/d", 0700, []byte("baz"), zip.Deflate, jar.DefaultTime, nil}
 be    = testZipEntry{"b/e", 0700, []byte(""), zip.Deflate, jar.DefaultTime, nil}

 withTimestamp    = testZipEntry{"timestamped", 0755, nil, zip.Store, jar.DefaultTime.Add(time.Hour), nil}
 withoutTimestamp = testZipEntry{"timestamped", 0755, nil, zip.Store, jar.DefaultTime, nil}

 service1a        = testZipEntry{"META-INF/services/service1", 0755, []byte("class1\nclass2\n"), zip.Store, jar.DefaultTime, nil}
 service1b        = testZipEntry{"META-INF/services/service1", 0755, []byte("class1\nclass3\n"), zip.Deflate, jar.DefaultTime, nil}
 service1combined = testZipEntry{"META-INF/services/service1", 0755, []byte("class1\nclass2\nclass3\n"), zip.Store, jar.DefaultTime, nil}
 service2         = testZipEntry{"META-INF/services/service2", 0755, []byte("class1\nclass2\n"), zip.Deflate, jar.DefaultTime, nil}

 metainfDir    = testZipEntry{jar.MetaDir, os.ModeDir | 0755, nil, zip.Store, jar.DefaultTime, []byte{0xfe, 0xca, 0, 0}}
 manifestFile  = testZipEntry{jar.ManifestFile, 0644, []byte("Manifest-Version: 1.0\nmanifest"), zip.Store, jar.DefaultTime, nil}
 manifestFile2 = testZipEntry{jar.ManifestFile, 0644, []byte("manifest2"), zip.Deflate, jar.DefaultTime, nil}

 pyMainFile = testZipEntry{"__main__.py", 0700, []byte("pyMain"), zip.Store, jar.DefaultTime, nil}
)

type testInputZip struct {
 name    string
 entries []testZipEntry
 reader  *zip.Reader
}

func (tiz *testInputZip) Name() string {
 return tiz.name
}

func (tiz *testInputZip) Open() error {
 if tiz.reader == nil {
  tiz.reader = testZipEntriesToZipReader(tiz.entries)
 }
 return nil
}

func (tiz *testInputZip) Close() error {
 tiz.reader = nil
 return nil
}

func (tiz *testInputZip) Entries() []*zip.File {
 if tiz.reader == nil {
  panic(fmt.Errorf("%s: should be open to get entries", tiz.Name()))
 }
 return tiz.reader.File
}

func (tiz *testInputZip) IsOpen() bool {
 return tiz.reader != nil
}

func TestMergeZips(t *testing.T) {
 testCases := []struct {
  name             string
  in               [][]testZipEntry
  stripFiles       []string
  stripDirs        []string
  jar              bool
  par              bool
  sort             bool
  ignoreDuplicates bool
  stripDirEntries  bool
  zipsToNotStrip   map[string]bool
  manifest         string
  pyMain           string

  out []testZipEntry
  err string
 }{
  {
   name: "duplicates error",
   in: [][]testZipEntry{
    {a},
    {a2},
    {a3},
   },
   out: []testZipEntry{a},
   err: "duplicate",
  },
  {
   name: "duplicates take first",
   in: [][]testZipEntry{
    {a},
    {a2},
    {a3},
   },
   out: []testZipEntry{a},

   ignoreDuplicates: true,
  },
  {
   name: "duplicates identical",
   in: [][]testZipEntry{
    {a},
    {a},
   },
   out: []testZipEntry{a},
  },
  {
   name: "sort",
   in: [][]testZipEntry{
    {be, bc, bDir, bbDir, bbb, A, metainfDir, manifestFile},
   },
   out: []testZipEntry{A, metainfDir, manifestFile, bDir, bbDir, bbb, bc, be},

   sort: true,
  },
  {
   name: "jar sort",
   in: [][]testZipEntry{
    {be, bc, bDir, A, metainfDir, manifestFile},
   },
   out: []testZipEntry{metainfDir, manifestFile, A, bDir, bc, be},

   jar: true,
  },
  {
   name: "jar merge",
   in: [][]testZipEntry{
    {metainfDir, manifestFile, bDir, be},
    {metainfDir, manifestFile2, bDir, bc},
    {metainfDir, manifestFile2, A},
   },
   out: []testZipEntry{metainfDir, manifestFile, A, bDir, bc, be},

   jar: true,
  },
  {
   name: "merge",
   in: [][]testZipEntry{
    {bDir, be},
    {bDir, bc},
    {A},
   },
   out: []testZipEntry{bDir, be, bc, A},
  },
  {
   name: "strip dir entries",
   in: [][]testZipEntry{
    {a, bDir, bbDir, bbb, bc, bd, be},
   },
   out: []testZipEntry{a, bbb, bc, bd, be},

   stripDirEntries: true,
  },
  {
   name: "strip files",
   in: [][]testZipEntry{
    {a, bDir, bbDir, bbb, bc, bd, be},
   },
   out: []testZipEntry{a, bDir, bbDir, bbb, bc},

   stripFiles: []string{"b/d", "b/e"},
  },
  {
   // merge_zips used to treat -stripFile a as stripping any file named a, it now only strips a in the
   // root of the zip.
   name: "strip file name",
   in: [][]testZipEntry{
    {a, bDir, ba},
   },
   out: []testZipEntry{bDir, ba},

   stripFiles: []string{"a"},
  },
  {
   name: "strip files glob",
   in: [][]testZipEntry{
    {a, bDir, ba},
   },
   out: []testZipEntry{bDir},

   stripFiles: []string{"**/a"},
  },
  {
   name: "strip dirs",
   in: [][]testZipEntry{
    {a, bDir, bbDir, bbb, bc, bd, be},
   },
   out: []testZipEntry{a},

   stripDirs: []string{"b"},
  },
  {
   name: "strip dirs glob",
   in: [][]testZipEntry{
    {a, bDir, bbDir, bbb, bc, bd, be},
   },
   out: []testZipEntry{a, bDir, bc, bd, be},

   stripDirs: []string{"b/*"},
  },
  {
   name: "zips to not strip",
   in: [][]testZipEntry{
    {a, bDir, bc},
    {bDir, bd},
    {bDir, be},
   },
   out: []testZipEntry{a, bDir, bd},

   stripDirs: []string{"b"},
   zipsToNotStrip: map[string]bool{
    "in1": true,
   },
  },
  {
   name: "services",
   in: [][]testZipEntry{
    {service1a, service2},
    {service1b},
   },
   jar: true,
   out: []testZipEntry{service1combined, service2},
  },
  {
   name: "strip timestamps",
   in: [][]testZipEntry{
    {withTimestamp},
    {a},
   },
   out: []testZipEntry{withoutTimestamp, a},
  },
  {
   name: "emulate par",
   in: [][]testZipEntry{
    {
     testZipEntry{name: "3/main.py"},
     testZipEntry{name: "c/main.py"},
     testZipEntry{name: "a/main.py"},
     testZipEntry{name: "2/main.py"},
     testZipEntry{name: "b/main.py"},
     testZipEntry{name: "1/main.py"},
    },
   },
   out: []testZipEntry{
    testZipEntry{name: "3/__init__.py", mode: 0700, timestamp: jar.DefaultTime},
    testZipEntry{name: "c/__init__.py", mode: 0700, timestamp: jar.DefaultTime},
    testZipEntry{name: "a/__init__.py", mode: 0700, timestamp: jar.DefaultTime},
    testZipEntry{name: "2/__init__.py", mode: 0700, timestamp: jar.DefaultTime},
    testZipEntry{name: "b/__init__.py", mode: 0700, timestamp: jar.DefaultTime},
    testZipEntry{name: "1/__init__.py", mode: 0700, timestamp: jar.DefaultTime},
    testZipEntry{name: "3/main.py", timestamp: jar.DefaultTime},
    testZipEntry{name: "c/main.py", timestamp: jar.DefaultTime},
    testZipEntry{name: "a/main.py", timestamp: jar.DefaultTime},
    testZipEntry{name: "2/main.py", timestamp: jar.DefaultTime},
    testZipEntry{name: "b/main.py", timestamp: jar.DefaultTime},
    testZipEntry{name: "1/main.py", timestamp: jar.DefaultTime},
   },
   par: true,
  },
  {
   name:     "manifest",
   manifest: "Manifest-Version: 1.0\nmanifest",
   out:      []testZipEntry{metainfDir, manifestFile},
  },
  {
   name:   "pyMain",
   pyMain: "pyMain",
   out:    []testZipEntry{pyMainFile},
  },
 }

 for _, test := range testCases {
  t.Run(test.name, func(t *testing.T) {
   inputZips := make([]InputZip, len(test.in))
   for i, in := range test.in {
    inputZips[i] = &testInputZip{name: "in" + strconv.Itoa(i), entries: in}
   }

   want := testZipEntriesToBuf(test.out)

   out := &bytes.Buffer{}
   writer := zip.NewWriter(out)

   manifestFile := ""
   if test.manifest != "" {
    manifestFile = filepath.Join(t.TempDir(), "manifest.txt")
    err := os.WriteFile(manifestFile, []byte(test.manifest), 0666)
    if err != nil {
     t.Fatal(err)
    }
   }

   pyMainFile := ""
   if test.pyMain != "" {
    pyMainFile = filepath.Join(t.TempDir(), "pymain.txt")
    err := os.WriteFile(pyMainFile, []byte(test.pyMain), 0666)
    if err != nil {
     t.Fatal(err)
    }
   }

   err := mergeZips(inputZips, writer, manifestFile, pyMainFile,
    test.sort, test.jar, test.par, test.stripDirEntries, test.ignoreDuplicates,
    test.stripFiles, test.stripDirs, test.zipsToNotStrip)

   closeErr := writer.Close()
   if closeErr != nil {
    t.Fatal(closeErr)
   }

   if test.err != "" {
    if err == nil {
     t.Fatal("missing err, expected: ", test.err)
    } else if !strings.Contains(strings.ToLower(err.Error()), strings.ToLower(test.err)) {
     t.Fatal("incorrect err, want:", test.err, "got:", err)
    }
    return
   } else if err != nil {
    t.Fatal("unexpected err: ", err)
   }

   if !bytes.Equal(want, out.Bytes()) {
    t.Error("incorrect zip output")
    t.Errorf("want:\n%s", dumpZip(want))
    t.Errorf("got:\n%s", dumpZip(out.Bytes()))
    os.WriteFile("/tmp/got.zip", out.Bytes(), 0755)
    os.WriteFile("/tmp/want.zip", want, 0755)
   }
  })
 }
}

func testZipEntriesToBuf(entries []testZipEntry) []byte {
 b := &bytes.Buffer{}
 zw := zip.NewWriter(b)

 for _, e := range entries {
  fh := zip.FileHeader{
   Name: e.name,
  }
  fh.SetMode(e.mode)
  fh.Method = e.method
  fh.SetModTime(e.timestamp)
  fh.UncompressedSize64 = uint64(len(e.data))
  fh.CRC32 = crc32.ChecksumIEEE(e.data)
  fh.Extra = e.extra
  if fh.Method == zip.Store {
   fh.CompressedSize64 = fh.UncompressedSize64
  }

  w, err := zw.CreateHeaderAndroid(&fh)
  if err != nil {
   panic(err)
  }

  _, err = w.Write(e.data)
  if err != nil {
   panic(err)
  }
 }

 err := zw.Close()
 if err != nil {
  panic(err)
 }

 return b.Bytes()
}

func testZipEntriesToZipReader(entries []testZipEntry) *zip.Reader {
 b := testZipEntriesToBuf(entries)
 r := bytes.NewReader(b)

 zr, err := zip.NewReader(r, int64(len(b)))
 if err != nil {
  panic(err)
 }

 return zr
}

func dumpZip(buf []byte) string {
 r := bytes.NewReader(buf)
 zr, err := zip.NewReader(r, int64(len(buf)))
 if err != nil {
  panic(err)
 }

 var ret string

 for _, f := range zr.File {
  ret += fmt.Sprintf("%v: %v %v %v %08x %s\n", f.Name, f.Mode(), f.UncompressedSize64, f.CompressedSize64, f.CRC32, f.ModTime())
 }

 return ret
}

type DummyInputZip struct {
 isOpen bool
}

func (diz *DummyInputZip) Name() string {
 return "dummy"
}

func (diz *DummyInputZip) Open() error {
 diz.isOpen = true
 return nil
}

func (diz *DummyInputZip) Close() error {
 diz.isOpen = false
 return nil
}

func (DummyInputZip) Entries() []*zip.File {
 panic("implement me")
}

func (diz *DummyInputZip) IsOpen() bool {
 return diz.isOpen
}

func TestInputZipsManager(t *testing.T) {
 const nInputZips = 20
 const nMaxOpenZips = 10
 izm := NewInputZipsManager(20, 10)
 managedZips := make([]InputZip, nInputZips)
 for i := 0; i < nInputZips; i++ {
  managedZips[i] = izm.Manage(&DummyInputZip{})
 }

 t.Run("InputZipsManager", func(t *testing.T) {
  for i, iz := range managedZips {
   if err := iz.Open(); err != nil {
    t.Fatalf("Step %d: open failed: %s", i, err)
    return
   }
   if izm.nOpenZips > nMaxOpenZips {
    t.Errorf("Step %d: should be <=%d open zips", i, nMaxOpenZips)
   }
  }
  if !managedZips[nInputZips-1].IsOpen() {
   t.Error("The last input should stay open")
  }
  for _, iz := range managedZips {
   iz.Close()
  }
  if izm.nOpenZips > 0 {
   t.Error("Some input zips are still open")
  }
 })
}