import argparse
import os

import httpx
import requests
from huggingface_hub import HfApi, upload_folder, create_repo
from huggingface_hub.utils import set_client_factory

try:
    from huggingface_hub import get_token
except ImportError:
    def get_token():
        return None


import time
from urllib3.util.retry import Retry
from requests.adapters import HTTPAdapter

# Configure Hugging Face Hub retry settings globally
os.environ["HF_HUB_ENABLE_EMERGENCY_RETRY"] = "True"
os.environ["HF_HUB_EMERGENCY_RETRY_WAIT_TIME"] = "5"

# Create a shared requests.Session to reuse TCP/TLS connections (Connection Pooling)
_session = requests.Session()

# Configure automatic retries at the urllib3 adapter level
retry_strategy = Retry(
    total=5,
    backoff_factor=1,
    status_forcelist=[500, 502, 503, 504],
    raise_on_status=False
)
adapter = HTTPAdapter(max_retries=retry_strategy)
_session.mount("https://", adapter)
_session.mount("http://", adapter)


# Custom transport to route httpx requests through requests library to bypass Windows/ISP TLS connection issues (WinError 10054)
class RequestsTransport(httpx.BaseTransport):
    def handle_request(self, request: httpx.Request) -> httpx.Response:
        content = request.read()
        max_attempts = 5
        backoff = 1.5
        
        headers = dict(request.headers)
        # Ensure a valid User-Agent is set
        if "user-agent" not in {k.lower() for k in headers.keys()}:
            headers["User-Agent"] = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36"
            
        for attempt in range(max_attempts):
            try:
                resp = _session.request(
                    method=request.method,
                    url=str(request.url),
                    headers=headers,
                    data=content,
                    timeout=60.0,
                )
                
                resp_headers = {k.lower(): v for k, v in resp.headers.items()}
                resp_headers.pop('content-encoding', None)
                resp_headers.pop('transfer-encoding', None)
                resp_headers['content-length'] = str(len(resp.content))
                
                return httpx.Response(
                    status_code=resp.status_code,
                    headers=list(resp_headers.items()),
                    content=resp.content,
                    request=request,
                )
            except (requests.exceptions.ConnectionError, requests.exceptions.Timeout) as e:
                if attempt == max_attempts - 1:
                    print(f"Error: Connection to Hugging Face failed after {max_attempts} attempts due to network restrictions.")
                    raise e
                print(f"Network reset (WinError 10054) or timeout encountered. Retrying in {backoff}s... (Attempt {attempt + 1}/{max_attempts})")
                time.sleep(backoff)
                backoff *= 2.0


# Register the requests-based transport
set_client_factory(lambda: httpx.Client(transport=RequestsTransport()))



def parse_args():
    parser = argparse.ArgumentParser(
        description="Upload a local folder to Hugging Face (Model, Space, or Dataset)."
    )
    parser.add_argument(
        "--repo-id",
        required=True,
        help="Hugging Face repo name, e.g. username/MCQ_model",
    )
    parser.add_argument(
        "--folder-path",
        default="models/ner_model",
        help="Local folder to upload. Default: models/ner_model",
    )
    parser.add_argument(
        "--repo-type",
        default="model",
        choices=["model", "dataset", "space"],
        help="Repository type: 'model', 'dataset', or 'space'. Default: model",
    )
    parser.add_argument(
        "--path-in-repo",
        default="",
        help="Destination path inside the repo. Default uploads files to the repo root.",
    )
    parser.add_argument(
        "--private",
        action="store_true",
        help="Create a private repo instead of public.",
    )
    parser.add_argument(
        "--space-sdk",
        choices=["gradio", "streamlit", "docker", "static"],
        default="streamlit",
        help="Space SDK to use if repo-type is 'space'. Default: streamlit",
    )
    parser.add_argument(
        "--token",
        help="Hugging Face token. If not set, the script checks HF_TOKEN and local cached login credentials.",
    )
    return parser.parse_args()


def main():
    args = parse_args()
    token = args.token or os.environ.get("HF_TOKEN") or get_token()
    if not token:
        raise SystemExit(
            "Hugging Face token not found. Please set the HF_TOKEN environment variable,\n"
            "run 'huggingface-cli login' to authenticate, or pass the token using --token."
        )

    folder_path = os.path.abspath(args.folder_path)
    if not os.path.isdir(folder_path):
        raise SystemExit(f"Local folder not found: {folder_path}")

    # Warn if uploading a Space without a README.md (required for Space metadata)
    if args.repo_type == "space" and not os.path.exists(os.path.join(folder_path, "README.md")):
        print("WARNING: Uploading a Space but no 'README.md' was found at the root of the folder.")
        print("   Hugging Face Spaces require a 'README.md' containing YAML configuration metadata (SDK, Title, etc.).\n")

    api = HfApi(token=token)
    print(f"Creating or checking repository '{args.repo_id}' (type: {args.repo_type})...")
    create_repo_kwargs = {
        "repo_id": args.repo_id,
        "token": token,
        "private": args.private,
        "repo_type": args.repo_type,
        "exist_ok": True,
    }
    if args.repo_type == "space":
        create_repo_kwargs["space_sdk"] = "static"
    create_repo(**create_repo_kwargs)

    print(f"Uploading folder '{folder_path}' to repository '{args.repo_id}' (type: {args.repo_type})...")
    print("Note: For large models (e.g. >500MB), the upload might take several minutes. Progress will be shown below.")
    upload_folder(
        repo_id=args.repo_id,
        folder_path=folder_path,
        path_in_repo=args.path_in_repo,
        token=token,
        repo_type=args.repo_type,
        ignore_patterns=["*.pyc", "**/__pycache__/*", "**/venv/*", "venv/*", "**/.git/*", "*.DS_Store", ".env", "models/*", "data/*"],
        delete_patterns=["**/venv/*", "venv/*", "question_generation/venv/*"],
    )

    print("Upload complete!")


if __name__ == "__main__":
    main()