import os
import sys
import subprocess
import time

# Auto-install huggingface_hub if not present
try:
    from huggingface_hub import HfApi, create_repo, upload_file, get_token
except ImportError:
    print("[*] Installing required library 'huggingface_hub'...")
    subprocess.check_call([sys.executable, "-m", "pip", "install", "huggingface_hub"])
    from huggingface_hub import HfApi, create_repo, upload_file, get_token

def clear_screen():
    os.system('cls' if os.name == 'nt' else 'clear')

def main():
    clear_screen()
    print("=" * 60)
    print("  HUGGING FACE - OPEN WEBUI AUTO DEPLOYER & CHAT SYSTEM")
    print("=" * 60)
    print("\nThis script will automatically create a Hugging Face Space,")
    print("install Open WebUI, and connect it to Hugging Face's API")
    print("so you can chat with any Hugging Face model in real time.\n")

    # 1. Ask for Hugging Face Token / Check cache
    print("Step 1: Authenticating with Hugging Face")
    print("-----------------------------------------")
    hf_token = get_token()
    
    if hf_token:
        print("[+] Active Hugging Face Token found in cache.")
    else:
        print("Please generate a Write Token from: https://huggingface.co/settings/tokens")
        hf_token = input("Enter your Hugging Face Write Token: ").strip()
        if not hf_token:
            print("[!] Token cannot be empty. Exiting.")
            return

    # Initialize HF API client
    api = HfApi(token=hf_token)

    try:
        # Validate token and get username
        user_info = api.whoami()
        username = user_info['name']
        print(f"[+] Successfully authenticated as user: {username}\n")
    except Exception as e:
        print(f"[!] Authentication failed: {e}")
        print("[!] Please check your token and try again.")
        return

    # 2. Ask for Space Name
    print("Step 2: Configure Space Name")
    print("----------------------------")
    default_space_name = "open-webui-chat"
    space_name = input(f"Enter Space Name (default: '{default_space_name}'): ").strip()
    if not space_name:
        space_name = default_space_name
    
    # Clean space name (lowercase, replace spaces/special chars with hyphens)
    space_name = "".join(c if c.isalnum() else "-" for c in space_name.lower())
    while "--" in space_name:
        space_name = space_name.replace("--", "-")
    space_name = space_name.strip("-")

    repo_id = f"{username}/{space_name}"
    print(f"[*] Target Space: https://huggingface.co/spaces/{repo_id}\n")

    # 3. Create Space (Docker SDK)
    print("Step 3: Creating Space on Hugging Face")
    print("--------------------------------------")
    try:
        create_repo(
            repo_id=repo_id,
            token=hf_token,
            repo_type="space",
            space_sdk="docker",
            private=False,
            exist_ok=True
        )
        print(f"[+] Space '{repo_id}' successfully prepared/created.")
    except Exception as e:
        print(f"[!] Error creating space: {e}")
        return

    # 4. Generate Dockerfile and README.md
    print("\nStep 4: Writing configuration files")
    print("---------------------------------")
    
    # Dockerfile contents optimized for Hugging Face Spaces (runs as non-root UID 1000)
    dockerfile_content = """FROM ghcr.io/open-webui/open-webui:main

# Environment variables for Hugging Face Spaces compatibility
ENV DATA_DIR=/tmp/open-webui
ENV HF_HOME=/tmp/hf_home
ENV WEBUI_PORT=7860
ENV PORT=7860
ENV WEBUI_AUTH=True

# Set permissions for the writable /tmp directory
USER root
RUN mkdir -p /tmp/open-webui /tmp/hf_home && chmod -R 777 /tmp/open-webui /tmp/hf_home

# Switch back to the Space user (UID 1000)
USER 1000
EXPOSE 7860
"""

    # README.md metadata configured for Hugging Face Space Docker
    readme_content = f"""---
title: Open WebUI Chat
emoji: 💬
colorFrom: blue
colorTo: green
sdk: docker
app_port: 7860
pinned: false
---

# Open WebUI Chat Space

This Space runs [Open WebUI](https://github.com/open-webui/open-webui) connected directly to the Hugging Face Serverless Inference API.

## Features
- Real-time chat with Hugging Face Models
- Select any available model from the list
- No local GPU/Ollama setup required
"""

    # Save files locally temporarily
    dockerfile_path = "Dockerfile_temp"
    readme_path = "README_temp.md"

    with open(dockerfile_path, "w", encoding="utf-8") as f:
        f.write(dockerfile_content)
    with open(readme_path, "w", encoding="utf-8") as f:
        f.write(readme_content)

    # 5. Upload files to Hugging Face
    print("[*] Uploading Dockerfile...")
    try:
        upload_file(
            path_or_fileobj=dockerfile_path,
            path_in_repo="Dockerfile",
            repo_id=repo_id,
            repo_type="space",
            token=hf_token
        )
        print("[+] Dockerfile uploaded.")
    except Exception as e:
        print(f"[!] Dockerfile upload failed: {e}")

    print("[*] Uploading README.md...")
    try:
        upload_file(
            path_or_fileobj=readme_path,
            path_in_repo="README.md",
            repo_id=repo_id,
            repo_type="space",
            token=hf_token
        )
        print("[+] README.md uploaded.")
    except Exception as e:
        print(f"[!] README.md upload failed: {e}")

    # Remove temporary files
    if os.path.exists(dockerfile_path):
        os.remove(dockerfile_path)
    if os.path.exists(readme_path):
        os.remove(readme_path)

    # 6. Configure Secrets & Variables
    print("\nStep 5: Configuring API connections (Hugging Face Router)")
    print("---------------------------------------------------------")
    try:
        # Set Hugging Face API key as a secret (represented as OPENAI_API_KEY for Open WebUI)
        print("[*] Configuring secret: OPENAI_API_KEY...")
        api.add_space_secret(
            repo_id=repo_id,
            key="OPENAI_API_KEY",
            value=hf_token
        )
        
        # Set API Base URL pointing to Hugging Face OpenAI-compatible API
        print("[*] Configuring variable: OPENAI_API_BASE_URL...")
        api.add_space_variable(
            repo_id=repo_id,
            key="OPENAI_API_BASE_URL",
            value="https://router.huggingface.co/v1"
        )

        # Disable local Ollama inside the space container to save memory
        print("[*] Configuring variable: ENABLE_OLLAMA...")
        api.add_space_variable(
            repo_id=repo_id,
            key="ENABLE_OLLAMA",
            value="False"
        )
        
        print("[+] All API configurations saved successfully.")
    except Exception as e:
        print(f"[!] Failed to set space secrets/variables: {e}")
        print("[!] You may need to manually add OPENAI_API_KEY and OPENAI_API_BASE_URL under Space Settings.")

    # 7. Complete Deployment info
    space_url = f"https://huggingface.co/spaces/{repo_id}"
    print("\n" + "=" * 60)
    print("  DEPLOYMENT INITIATED SUCCESSFULLY!")
    print("=" * 60)
    print(f"\nYour Open WebUI instance is building and will be live shortly.")
    print(f"👉 Access URL: {space_url}")
    print("\nHow to use:")
    print("1. Click the URL above and wait a few minutes for the build to finish.")
    print("2. Once loaded, sign up / log in to your Open WebUI space (the first user created is automatically the Admin).")
    print("3. In the chat interface, click the model selection dropdown at the top.")
    print("4. You will see Hugging Face models available (e.g., DeepSeek, Qwen, etc.) in real time!")
    print("5. Type your message and hit send to start chatting.\n")
    print("=" * 60)

if __name__ == "__main__":
    main()