import os
import unittest
import uuid

import tos

from tosfsspec import TosFileSystem


class TosFileSystemTestMixin:
    bucket_type = 'fns'

    @classmethod
    def setUpClass(cls):
        cls.region = os.environ.get('TOS_REGION')
        cls.endpoint = os.environ.get('TOS_ENDPOINT', '')
        cls.access_key = os.environ.get('TOS_ACCESS_KEY', '')
        cls.secret_key = os.environ.get('TOS_SECRET_KEY', '')

        if not all([cls.region, cls.access_key, cls.secret_key]):
            raise unittest.SkipTest("TOS environment variables (REGION, ACCESS_KEY, SECRET_KEY) not set")

        cls.fs = TosFileSystem(
            region=cls.region,
            endpoint=cls.endpoint,
            key=cls.access_key,
            secret=cls.secret_key,
        )
        cls.bucket_name = f"test-tos-fs-{cls.bucket_type}-{uuid.uuid4().hex}"
        try:
            cls.fs.mkdir(cls.bucket_name, create_bucket=True, bucket_type=cls.bucket_type)
        except Exception as e:
            print(f"Failed to create bucket {cls.bucket_name}: {e}")
            raise

        # Initialize native tos client for cleanup operations not supported by fs
        cls.tos_client = tos.TosClientV2(
            ak=cls.access_key,
            sk=cls.secret_key,
            endpoint=cls.endpoint,
            region=cls.region,
        )

    @classmethod
    def tearDownClass(cls):
        try:
            # Clean up bucket contents including multipart uploads
            if cls.bucket_type == 'hns':
                cls._cleanup_hns()
            else:
                cls._cleanup_fns()

            # 3. Delete the bucket
            cls.fs.rmdir(cls.bucket_name, delete_bucket=True)
        except Exception as e:
            print(f"Failed to delete bucket {cls.bucket_name}: {e}")
            raise

    @classmethod
    def _cleanup_fns(cls):
        try:
            # 1. Abort all multipart uploads
            upload_id_marker = None
            while True:
                output = cls.tos_client.list_multipart_uploads(cls.bucket_name, upload_id_marker=upload_id_marker)
                if output.uploads:
                    for upload in output.uploads:
                        cls.tos_client.abort_multipart_upload(cls.bucket_name, upload.key, upload.upload_id)

                if not output.is_truncated:
                    break
                upload_id_marker = output.next_upload_id_marker

            # 2. Delete all objects (files and directories)
            files = cls.fs.ls(cls.bucket_name, detail=False)
            for f in files:
                cls.fs.rm(f, recursive=True)

        except Exception as e:
            print(f"Error cleaning up FNS bucket contents: {e}")
            raise

    @classmethod
    def _cleanup_hns(cls):
        try:
            # 1. Abort all multipart uploads
            upload_id_marker = None
            while True:
                output = cls.tos_client.list_multipart_uploads(cls.bucket_name, upload_id_marker=upload_id_marker)
                if output.uploads:
                    for upload in output.uploads:
                        cls.tos_client.abort_multipart_upload(cls.bucket_name, upload.key, upload.upload_id)

                if not output.is_truncated:
                    break
                upload_id_marker = output.next_upload_id_marker

            # 2. Recursive delete using TOS SDK directly
            def delete_recursive(prefix=''):
                marker = ''
                while True:
                    output = cls.tos_client.list_objects(cls.bucket_name, prefix=prefix, delimiter='/', marker=marker)

                    # Delete files
                    if output.contents:
                        for obj in output.contents:
                            if obj.key.endswith('/'):
                                continue
                            cls.tos_client.delete_object(cls.bucket_name, obj.key)

                    # Process subdirectories
                    if output.common_prefixes:
                        for cp in output.common_prefixes:
                            # cp.prefix contains the subdirectory path ending with '/'
                            delete_recursive(cp.prefix)
                            # Delete the empty directory itself
                            cls.tos_client.delete_object(cls.bucket_name, cp.prefix)

                    if not output.is_truncated:
                        break
                    marker = output.next_marker

            delete_recursive()

        except Exception as e:
            print(f"Error cleaning up HNS bucket contents: {e}")
            raise

    def setUp(self):
        # Clean bucket before each test? Or just ensure unique paths?
        # Using unique paths per test is better for performance.
        self.test_dir = f"{self.bucket_name}/test-{uuid.uuid4().hex}"
        self.fs.makedirs(self.test_dir)

    def tearDown(self):
        try:
            self.fs.rm(self.test_dir, recursive=True)
        except Exception:
            pass

    def test_simple_write_read(self):
        assert self.fs.isdir(f"{self.test_dir}/")
        assert self.fs.isdir(f"{self.test_dir}")

        path = f"{self.test_dir}/file.txt"
        data = b"hello world"

        # Write
        with self.fs.open(path, "wb") as f:
            f.write(data)

        # Read
        with self.fs.open(path, "rb") as f:
            read_data = f.read()

        self.assertEqual(data, read_data)

    def test_ls(self):
        path1 = f"{self.test_dir}/file1.txt"
        path2 = f"{self.test_dir}/file2.txt"
        self.fs.touch(path1)
        self.fs.touch(path2)

        files = self.fs.ls(self.test_dir, detail=False)
        self.assertTrue(len(files) >= 2)
        self.assertIn(path1, files)
        self.assertIn(path2, files)

    def test_exists(self):
        path = f"{self.test_dir}/exists.txt"
        self.assertFalse(self.fs.exists(path))

        self.fs.touch(path)
        self.assertTrue(self.fs.exists(path))

    def test_info(self):
        path = f"{self.test_dir}/info.txt"
        data = b"content"
        with self.fs.open(path, "wb") as f:
            f.write(data)

        info = self.fs.info(path)
        self.assertEqual(info['name'], path)
        self.assertEqual(info['size'], len(data))
        self.assertEqual(info['type'], 'file')

    def test_rm(self):
        path = f"{self.test_dir}/rm.txt"
        self.fs.touch(path)
        self.assertTrue(self.fs.exists(path))

        self.fs.rm(path)
        self.assertFalse(self.fs.exists(path))

    def test_cp(self):
        src = f"{self.test_dir}/src.txt"
        dst = f"{self.test_dir}/dst.txt"
        data = b"copy me"
        with self.fs.open(src, "wb") as f:
            f.write(data)

        self.fs.copy(src, dst)
        self.assertTrue(self.fs.exists(dst))

        with self.fs.open(dst, "rb") as f:
            self.assertEqual(f.read(), data)

    def test_mv(self):
        src = f"{self.test_dir}/src_mv.txt"
        dst = f"{self.test_dir}/dst_mv.txt"
        data = b"move me"
        with self.fs.open(src, "wb") as f:
            f.write(data)

        self.fs.mv(src, dst)
        self.assertFalse(self.fs.exists(src))
        self.assertTrue(self.fs.exists(dst))

        with self.fs.open(dst, "rb") as f:
            self.assertEqual(f.read(), data)

    def test_walk(self):
        # Create structure:
        # /dir1/f1
        # /dir1/subdir/f2

        dir1 = f"{self.test_dir}/dir1"
        subdir = f"{dir1}/subdir"
        f1 = f"{dir1}/f1"
        f2 = f"{subdir}/f2"

        self.fs.makedirs(subdir, exist_ok=True)
        self.fs.touch(f1)
        self.fs.touch(f2)

        walk_res = list(self.fs.walk(dir1))
        # Expected:
        # (dir1, [subdir], [f1])
        # (subdir, [], [f2]) or similar depending on implementation details

        # Note: walk output format depends on fsspec version and implementation.
        # Typically: (root, dirs, files)

        self.assertTrue(len(walk_res) >= 1)

        found_f1 = False
        found_f2 = False

        for root, dirs, files in walk_res:
            if root == dir1:
                # files might be full paths or relative names depending on implementation
                # Based on TosFileSystem.walk:
                # files = ['/'.join([bucket, key]) ...] if detail=False
                # So they are full paths.
                # Wait, let's check core.py line 338:
                # files = sorted(['/'.join([tos_object.bucket, tos_object.key]) ...])
                # So they are full paths.

                # Check f1
                for file in files:
                    if root + "/" + file == f1:
                        found_f1 = True
                        break
            elif root == subdir:
                for file in files:
                    if root + "/" + file == f2:
                        found_f2 = True
                        break

        # Since implementation might return all files in one go if it's just recursive list transformed?
        # Let's check `walk` implementation in core.py again.
        # It calls `self._client.walk`.
        # It yields `result[0], dirs, files`.

        self.assertTrue(found_f1, "f1 not found in walk")
        self.assertTrue(found_f2, "f2 not found in walk")

        walk_res = list(self.fs.walk(dir1, use_full_path=True))
        found_f1 = False
        found_f2 = False
        for root, dirs, files in walk_res:
            if root == dir1:
                # files might be full paths or relative names depending on implementation
                # Based on TosFileSystem.walk:
                # files = ['/'.join([bucket, key]) ...] if detail=False
                # So they are full paths.
                # Wait, let's check core.py line 338:
                # files = sorted(['/'.join([tos_object.bucket, tos_object.key]) ...])
                # So they are full paths.

                # Check f1
                if f1 in files:
                    found_f1 = True
            elif root == subdir:
                if f2 in files:
                    found_f2 = True

        # Since implementation might return all files in one go if it's just recursive list transformed?
        # Let's check `walk` implementation in core.py again.
        # It calls `self._client.walk`.
        # It yields `result[0], dirs, files`.

        self.assertTrue(found_f1, "f1 not found in walk")
        self.assertTrue(found_f2, "f2 not found in walk")

    def test_glob(self):
        # glob is implemented in AbstractFileSystem using find/ls
        # TosFileSystem overrides find, but glob usually calls find.

        p1 = f"{self.test_dir}/g1.txt"
        p2 = f"{self.test_dir}/g2.txt"
        p3 = f"{self.test_dir}/other.log"

        self.fs.touch(p1)
        self.fs.touch(p2)
        self.fs.touch(p3)

        txt_files = self.fs.glob(f"{self.test_dir}/*.txt")
        self.assertEqual(len(txt_files), 2)
        self.assertIn(p1, txt_files)
        self.assertIn(p2, txt_files)

    def test_append_mode(self):
        # TosFile supports append mode 'ab' if supported by backend
        path = f"{self.test_dir}/append.txt"
        data1 = b"part1"
        data2 = b"part2"

        with self.fs.open(path, "wb") as f:
            f.write(data1)

        try:
            with self.fs.open(path, "ab") as f:
                f.write(data2)

            with self.fs.open(path, "rb") as f:
                self.assertEqual(f.read(), data1 + data2)
        except NotImplementedError:
            # Append might not be supported or configured
            pass
        except Exception as e:
            # Append support depends on TOS capabilities for the bucket type or client
            print(f"Append test failed: {e}")


class TestTosFileSystemFNS(TosFileSystemTestMixin, unittest.TestCase):
    bucket_type = 'fns'


class TestTosFileSystemHNS(TosFileSystemTestMixin, unittest.TestCase):
    bucket_type = 'hns'


if __name__ == "__main__":
    unittest.main()
