import os
import shutil
import tempfile
import time
import fnmatch
from pathlib import Path

import gradio as gr
import httpx
from git import Repo
from huggingface_hub import CommitOperationAdd, HfApi


os.environ["HF_HUB_DOWNLOAD_TIMEOUT"] = "300"
os.environ["HF_HUB_ETAG_TIMEOUT"] = "300"
os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1"

api = HfApi()
ALLOWED_SPACE_SDKS = {"gradio", "streamlit", "docker", "static"}


def remove_broken_symlinks(directory):
    removed = []
    for root, dirs, files in os.walk(directory, topdown=True):
        dirs[:] = [d for d in dirs if d != ".git"]

        for dirname in dirs[:]:
            dirpath = os.path.join(root, dirname)
            if os.path.islink(dirpath) and not os.path.exists(dirpath):
                os.unlink(dirpath)
                dirs.remove(dirname)
                removed.append(dirpath)

        for filename in files:
            filepath = os.path.join(root, filename)
            if os.path.islink(filepath) and not os.path.exists(filepath):
                os.unlink(filepath)
                removed.append(filepath)

    return removed


def resolve_symlinks_to_copies(directory):
    converted = []
    for root, dirs, files in os.walk(directory, topdown=True):
        dirs[:] = [d for d in dirs if d != ".git"]

        for filename in files:
            filepath = os.path.join(root, filename)
            if not os.path.islink(filepath) or not os.path.exists(filepath):
                continue

            real_path = os.path.realpath(filepath)
            try:
                os.unlink(filepath)
                if os.path.isfile(real_path):
                    shutil.copy2(real_path, filepath)
                elif os.path.isdir(real_path):
                    shutil.copytree(real_path, filepath)
                converted.append(filepath)
            except Exception as e:
                print(f"Warning: could not convert symlink {filepath}: {e}")

    return converted


def count_files_and_size(directory):
    total_files = 0
    total_size = 0

    for root, dirs, files in os.walk(directory):
        dirs[:] = [d for d in dirs if d != ".git"]
        for filename in files:
            filepath = os.path.join(root, filename)
            if not os.path.exists(filepath):
                continue
            total_files += 1
            try:
                total_size += os.path.getsize(filepath)
            except OSError:
                pass

    return total_files, total_size


def summarize_tree(directory, max_depth=3, max_files_per_dir=10):
    lines = []
    base = Path(directory)

    for root, dirs, files in os.walk(directory):
        dirs[:] = [d for d in dirs if d != ".git"]
        rel = Path(root).relative_to(base)
        depth = 0 if str(rel) == "." else len(rel.parts)
        if depth > max_depth:
            dirs[:] = []
            continue

        indent = "  " * depth
        name = base.name if str(rel) == "." else rel.name
        lines.append(f"{indent}{name}/")

        file_indent = "  " * (depth + 1)
        for filename in sorted(files)[:max_files_per_dir]:
            lines.append(f"{file_indent}{filename}")
        if len(files) > max_files_per_dir:
            lines.append(f"{file_indent}... and {len(files) - max_files_per_dir} more files")

    return "\n".join(lines[:250])


def normalize_space_id(space_destination, token):
    username = api.whoami(token=token)["name"]
    if "/" not in space_destination:
        return f"{username}/{space_destination}"
    return space_destination


def normalize_space_sdk(space_sdk):
    if space_sdk is None:
        return "gradio"

    normalized = str(space_sdk).strip().lower()
    if not normalized:
        return "gradio"

    if normalized not in ALLOWED_SPACE_SDKS:
        raise ValueError(
            f"Invalid Space SDK `{space_sdk}`. Expected one of: {', '.join(sorted(ALLOWED_SPACE_SDKS))}."
        )

    return normalized


def clone_repository(git_repo_url, tmp_dir, recursive_bool):
    print(f"Cloning {git_repo_url}")

    repo = Repo.clone_from(
        git_repo_url,
        tmp_dir,
        depth=1,
        recursive=recursive_bool,
    )

    if recursive_bool:
        print("Initializing submodules recursively...")
        try:
            repo.git.submodule("update", "--init", "--recursive", "--depth", "1")
        except Exception as shallow_error:
            print(f"Shallow submodule update failed, retrying full submodule update: {shallow_error}")
            repo.git.submodule("update", "--init", "--recursive")

    return repo


def ensure_space_exists(space_destination, token, space_sdk="gradio"):
    resolved_space_sdk = normalize_space_sdk(space_sdk)
    authenticated_api = HfApi(token=token)
    try:
        authenticated_api.repo_info(repo_id=space_destination, repo_type="space")
        print(f"Space exists: {space_destination}")
        return False
    except Exception as e:
        print(f"Space not found or not accessible, creating it: {e}")

    authenticated_api.create_repo(
        repo_id=space_destination,
        repo_type="space",
        private=True,
        space_sdk=resolved_space_sdk,
        exist_ok=True,
    )
    print(f"Space created: {space_destination} with SDK: {resolved_space_sdk}")
    return True


def upload_small_folder(tmp_dir, space_destination, token, ignore_patterns):
    max_retries = 3
    for attempt in range(max_retries):
        try:
            print(f"Upload attempt {attempt + 1}/{max_retries}")
            api.upload_folder(
                folder_path=tmp_dir,
                repo_id=space_destination,
                repo_type="space",
                token=token,
                commit_message="Migrated from Git repository",
                ignore_patterns=ignore_patterns,
            )
            return
        except (httpx.ReadTimeout, httpx.ConnectTimeout) as e:
            if attempt == max_retries - 1:
                raise RuntimeError(f"Upload timeout after {max_retries} attempts: {e}") from e
            print(f"Timeout, retrying in 5s: {e}")
            time.sleep(5)


def should_ignore_path(relative_path, ignore_patterns):
    rel = relative_path.replace(os.sep, "/")
    name = os.path.basename(rel)
    parts = set(Path(rel).parts)

    ignored_dirs = {".git", ".cache", "__pycache__", ".pytest_cache", "node_modules", ".venv", "venv"}
    if parts.intersection(ignored_dirs):
        return True

    for pattern in ignore_patterns:
        normalized = pattern.replace(os.sep, "/")
        if fnmatch.fnmatch(rel, normalized) or fnmatch.fnmatch(name, normalized):
            return True

        if normalized.endswith("/**"):
            prefix = normalized[:-3].rstrip("/")
            if rel == prefix or rel.startswith(prefix + "/"):
                return True

    return False


def iter_upload_files(folder_path, ignore_patterns):
    folder = Path(folder_path)
    for root, dirs, files in os.walk(folder_path):
        dirs[:] = [
            dirname
            for dirname in dirs
            if not should_ignore_path(str((Path(root) / dirname).relative_to(folder)), ignore_patterns)
        ]

        for filename in files:
            local_path = Path(root) / filename
            relative_path = str(local_path.relative_to(folder)).replace(os.sep, "/")
            if should_ignore_path(relative_path, ignore_patterns):
                continue
            yield local_path, relative_path


def upload_folder_in_batches(tmp_dir, space_destination, token, ignore_patterns, max_files=75, max_bytes=50 * 1024 * 1024):
    authenticated_api = HfApi(token=token)
    batch = []
    batch_size = 0
    batch_index = 1
    total_files = 0

    def commit_batch(files, index):
        if not files:
            return

        operations = [
            CommitOperationAdd(path_or_fileobj=str(local_path), path_in_repo=path_in_repo)
            for local_path, path_in_repo in files
        ]
        print(f"Uploading batch {index}: {len(operations)} files")
        authenticated_api.create_commit(
            repo_id=space_destination,
            repo_type="space",
            operations=operations,
            commit_message=f"Migrated files batch {index}",
        )

    for local_path, relative_path in iter_upload_files(tmp_dir, ignore_patterns):
        file_size = local_path.stat().st_size
        should_flush = batch and (len(batch) >= max_files or batch_size + file_size > max_bytes)

        if should_flush:
            commit_batch(batch, batch_index)
            batch_index += 1
            total_files += len(batch)
            batch = []
            batch_size = 0

        batch.append((local_path, relative_path))
        batch_size += file_size

    if batch:
        commit_batch(batch, batch_index)
        total_files += len(batch)

    print(f"Uploaded {total_files} files in {batch_index} batch(es)")


def process(git_repo_url, recursive_bool, space_destination, user_token, space_sdk, progress=gr.Progress()):
    if not git_repo_url or not space_destination or not user_token:
        return "Error: please fill in the Git URL, Space ID, and HF token."

    tmp_dir = tempfile.mkdtemp(prefix="clone-to-space-")

    try:
        progress(0.02, desc="Verifying Hugging Face token...")
        resolved_space_sdk = normalize_space_sdk(space_sdk)
        space_destination = normalize_space_id(space_destination.strip(), user_token.strip())
        print(f"Target Space: {space_destination}")
        print(f"Space SDK: {resolved_space_sdk}")

        progress(0.08, desc="Cloning Git repository...")
        clone_repository(git_repo_url.strip(), tmp_dir, bool(recursive_bool))

        progress(0.18, desc="Cleaning Git metadata...")
        git_dir = os.path.join(tmp_dir, ".git")
        if os.path.exists(git_dir):
            shutil.rmtree(git_dir)

        progress(0.22, desc="Processing symbolic links...")
        broken_links = remove_broken_symlinks(tmp_dir)
        converted_links = resolve_symlinks_to_copies(tmp_dir)
        print(f"Removed {len(broken_links)} broken symlinks")
        print(f"Converted {len(converted_links)} valid symlinks")

        progress(0.28, desc="Processing README...")
        readme_path = os.path.join(tmp_dir, "README.md")
        original_readme_path = os.path.join(tmp_dir, "ORIGINAL_README.md")
        if os.path.exists(readme_path) and not os.path.exists(original_readme_path):
            os.rename(readme_path, original_readme_path)

        file_count, total_size = count_files_and_size(tmp_dir)
        size_mb = total_size / (1024 * 1024)
        print(f"Files: {file_count}, size: {size_mb:.2f} MB")
        print("Repository preview:")
        print(summarize_tree(tmp_dir))

        progress(0.36, desc="Ensuring destination Space exists...")
        ensure_space_exists(space_destination, user_token.strip(), resolved_space_sdk)

        ignore_patterns = [
            ".git",
            ".git/**",
            ".gitignore",
            ".gitattributes",
            ".DS_Store",
            "*.env",
            "**/__pycache__/**",
            "**/*.pyc",
            "**/.pytest_cache/**",
            ".cache",
            ".cache/**",
            "node_modules",
            "node_modules/**",
            ".venv",
            ".venv/**",
            "venv",
            "venv/**",
        ]

        if size_mb > 50 or file_count > 500:
            progress(0.48, desc=f"Large upload: {file_count} files, {size_mb:.1f} MB")
            print("Using chunked create_commit upload for Space compatibility")
            upload_folder_in_batches(tmp_dir, space_destination, user_token.strip(), ignore_patterns)
        else:
            progress(0.48, desc=f"Uploading: {file_count} files, {size_mb:.1f} MB")
            print("Using HfApi.upload_folder()")
            upload_small_folder(tmp_dir, space_destination, user_token.strip(), ignore_patterns)

        progress(1.0, desc="Complete")
        space_url = f"https://huggingface.co/spaces/{space_destination}"
        return f"Success! View Space: {space_url}"

    except Exception as e:
        error_msg = str(e)
        print(f"Error: {error_msg}")
        return f"Error occurred: {error_msg}"

    finally:
        if os.path.exists(tmp_dir):
            shutil.rmtree(tmp_dir)
            print("Cleaned up temporary directory")


css = """
.gradio-container {
  max-width: 920px !important;
  margin: 0 auto !important;
}

#main-container {
  padding: 28px 18px;
}

#main-container h1,
#main-container > .markdown p {
  text-align: center;
}

.form-panel {
  border: 1px solid var(--border-color-primary);
  border-radius: 8px;
  padding: 20px;
}

.primary-btn button,
button.primary {
  border-radius: 8px !important;
}
"""


with gr.Blocks(theme=gr.themes.Soft(), css=css, title="Clone Git Repo to HF Space") as demo:
    with gr.Column(elem_id="main-container"):
        gr.Markdown(
            """
            # Clone Git Repo to HF Space
            Clone a Git repository into a Hugging Face Space, including submodules when Recursive mode is enabled.
            """
        )

        with gr.Column(elem_classes=["form-panel"]):
            with gr.Row():
                git_repo_url = gr.Textbox(
                    label="Git Repository URL",
                    placeholder="https://github.com/username/repo.git",
                    scale=3,
                )
                recursive_mode = gr.Checkbox(
                    label="Recursive mode",
                    value=True,
                    info="Clone Git submodules too. Recommended for repos with third_party dependencies.",
                    scale=1,
                )

            with gr.Row():
                space_destination = gr.Textbox(
                    label="Destination Space ID",
                    placeholder="your_hf_username/space_name",
                    info="If you only enter a Space name, your username is added automatically.",
                )
                user_token = gr.Textbox(
                    label="HF Token with WRITE permission",
                    type="password",
                )

            space_sdk = gr.Radio(
                label="Space SDK",
                choices=["gradio", "streamlit", "docker", "static"],
                value="gradio",
            )

            submit_btn = gr.Button("Clone to Space", variant="primary", elem_classes=["primary-btn"])
            status = gr.Textbox(label="Status", lines=4, interactive=False)

        submit_btn.click(
            fn=process,
            inputs=[git_repo_url, recursive_mode, space_destination, user_token, space_sdk],
            outputs=[status],
        )


if __name__ == "__main__":
    demo.queue(default_concurrency_limit=1, max_size=10)
    demo.launch()