import gradio as gr
import os
import shutil
import tempfile
import json
from datetime import datetime
from huggingface_hub import HfApi, upload_file, list_repo_files, hf_hub_download
from PIL import Image, ImageFont, ImageDraw
import requests
import traceback

print("=" * 60)
print("🚀 FONT MANAGER FOR HF DATASET")
print("=" * 60)

# =============================================
# CONFIGURATION
# =============================================
HF_TOKEN = os.environ.get("HF_TOKEN")
HF_USERNAME = "yukee1992"
FONTS_DATASET = f"{HF_USERNAME}/video-fonts"
PREVIEW_DIR = "/tmp/font_previews"

print(f"📦 Fonts Dataset: {FONTS_DATASET}")
print(f"🔑 HF Token: {'✅ Set' if HF_TOKEN else '❌ Missing'}")

# Create preview directory
os.makedirs(PREVIEW_DIR, exist_ok=True)

# Initialize HF API
api = HfApi(token=HF_TOKEN)

# =============================================
# DATASET FUNCTIONS
# =============================================

def ensure_dataset_exists():
    """Create dataset if it doesn't exist"""
    try:
        api.dataset_info(FONTS_DATASET)
        print(f"✅ Dataset {FONTS_DATASET} exists")
    except Exception:
        print(f"📦 Creating dataset: {FONTS_DATASET}")
        api.create_repo(
            repo_id=FONTS_DATASET,
            repo_type="dataset",
            private=False,
            exist_ok=True
        )
        print(f"✅ Created dataset")

def get_font_list():
    """Get list of all fonts in dataset"""
    try:
        files = api.list_repo_files(repo_id=FONTS_DATASET, repo_type="dataset")
        fonts = []
        categories = set()
        
        for file in files:
            if file.endswith(('.ttf', '.otf')):
                # Parse category from path
                parts = file.split('/')
                category = parts[0] if len(parts) > 1 else "uncategorized"
                categories.add(category)
                
                fonts.append({
                    "path": file,
                    "name": os.path.basename(file),
                    "category": category,
                    "url": f"https://huggingface.co/datasets/{FONTS_DATASET}/resolve/main/{file}"
                })
        
        return fonts, sorted(list(categories))
    except Exception as e:
        print(f"❌ Error listing fonts: {e}")
        return [], []

# =============================================
# FIXED UPLOAD FUNCTION
# =============================================
def upload_font(file_obj, category, font_name, description):
    """Upload a font to the dataset - FIXED for both file objects and paths"""
    try:
        # Check if file is provided
        if file_obj is None:
            return "❌ Please select a font file to upload"
        
        # Check if font name is provided
        if not font_name or font_name.strip() == "":
            return "❌ Please enter a font name"
        
        print(f"📤 Uploading font: {font_name}")
        print(f"   Category: {category}")
        print(f"   File object type: {type(file_obj)}")
        
        # Handle different input types
        temp_file_path = None
        original_filename = "font.ttf"
        
        if hasattr(file_obj, 'read'):
            # It's a file-like object (Gradio file object)
            print("   Type: File-like object")
            with tempfile.NamedTemporaryFile(delete=False, suffix='.ttf') as tmp:
                shutil.copyfileobj(file_obj, tmp)
                temp_file_path = tmp.name
            original_filename = getattr(file_obj, 'name', 'font.ttf')
            
        elif isinstance(file_obj, str):
            # It's a string (file path)
            print("   Type: String path")
            if os.path.exists(file_obj):
                temp_file_path = file_obj
                original_filename = os.path.basename(file_obj)
            else:
                return f"❌ File not found at path: {file_obj}"
            
        else:
            return f"❌ Unexpected file type: {type(file_obj)}"
        
        print(f"   Original filename: {original_filename}")
        print(f"   Temp path: {temp_file_path}")
        
        # Determine file extension from original filename
        if original_filename.endswith('.otf'):
            ext = '.otf'
        else:
            ext = '.ttf'
        
        # Create safe filename
        safe_name = font_name.replace(' ', '-').lower()
        # Remove any non-alphanumeric characters except hyphen
        safe_name = ''.join(c for c in safe_name if c.isalnum() or c == '-')
        filename = f"{safe_name}{ext}"
        
        # Path in dataset: category/filename
        dataset_path = f"{category}/{filename}"
        
        print(f"   Dataset path: {dataset_path}")
        
        # Upload to dataset
        api.upload_file(
            path_or_fileobj=temp_file_path,
            path_in_repo=dataset_path,
            repo_id=FONTS_DATASET,
            repo_type="dataset"
        )
        
        print(f"✅ File uploaded successfully")
        
        # Update metadata
        try:
            update_font_metadata(category, filename, font_name, description)
        except Exception as e:
            print(f"⚠️ Metadata update failed: {e}")
        
        # Clean up temp file if we created one
        if temp_file_path and hasattr(file_obj, 'read') and os.path.exists(temp_file_path):
            os.unlink(temp_file_path)
        
        return f"✅ Successfully uploaded {font_name} to {category}/"
        
    except Exception as e:
        print(f"❌ Upload error: {str(e)}")
        traceback.print_exc()
        return f"❌ Upload failed: {str(e)}"

def update_font_metadata(category, filename, font_name, description):
    """Update the metadata JSON file for the category"""
    try:
        metadata_path = f"{category}/metadata.json"
        existing_metadata = []
        
        # Try to download existing metadata
        try:
            metadata_file = hf_hub_download(
                repo_id=FONTS_DATASET,
                filename=metadata_path,
                repo_type="dataset",
                token=HF_TOKEN
            )
            with open(metadata_file, 'r', encoding='utf-8') as f:
                existing_metadata = json.load(f)
        except:
            existing_metadata = []
        
        # Add new font metadata
        existing_metadata.append({
            "name": font_name,
            "file": filename,
            "description": description,
            "category": category,
            "uploaded_at": datetime.now().isoformat()
        })
        
        # Save updated metadata
        with tempfile.NamedTemporaryFile(mode='w', delete=False, suffix='.json', encoding='utf-8') as f:
            json.dump(existing_metadata, f, indent=2, ensure_ascii=False)
            meta_path = f.name
        
        # Upload metadata
        api.upload_file(
            path_or_fileobj=meta_path,
            path_in_repo=metadata_path,
            repo_id=FONTS_DATASET,
            repo_type="dataset"
        )
        
        os.unlink(meta_path)
        print(f"✅ Metadata updated for {category}")
        
    except Exception as e:
        print(f"⚠️ Metadata error: {e}")
        raise

def delete_font(font_path):
    """Delete a font from the dataset"""
    try:
        api.delete_file(
            path_in_repo=font_path,
            repo_id=FONTS_DATASET,
            repo_type="dataset"
        )
        return f"✅ Deleted: {font_path}"
    except Exception as e:
        return f"❌ Delete failed: {str(e)}"

def generate_preview(font_url, text="ABCDEFGHIJKLMNOPQRSTUVWXYZ\nabcdefghijklmnopqrstuvwxyz\n0123456789\n中文字体测试", font_size=36):
    """Generate a preview image for a font"""
    try:
        # Download font file
        response = requests.get(font_url)
        if response.status_code != 200:
            return None
            
        font_path = os.path.join(PREVIEW_DIR, f"temp_{datetime.now().timestamp()}.ttf")
        with open(font_path, 'wb') as f:
            f.write(response.content)
        
        # Create preview image
        img = Image.new('RGB', (800, 400), color='white')
        d = ImageDraw.Draw(img)
        
        try:
            font = ImageFont.truetype(font_path, font_size)
        except:
            font = ImageFont.load_default()
        
        # Draw text
        d.text((50, 50), text, fill='black', font=font)
        
        # Save preview
        preview_path = os.path.join(PREVIEW_DIR, f"preview_{datetime.now().timestamp()}.png")
        img.save(preview_path)
        
        # Clean up font file
        os.unlink(font_path)
        
        return preview_path
    except Exception as e:
        print(f"❌ Preview failed: {e}")
        return None

# =============================================
# GRADIO INTERFACE
# =============================================

# Ensure dataset exists
ensure_dataset_exists()

# Get initial font list
fonts, categories = get_font_list()

with gr.Blocks(title="Font Manager", theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 🎨 Font Manager for Video Projects")
    gr.Markdown(f"### Dataset: `{FONTS_DATASET}`")
    
    with gr.Tabs():
        # =====================================
        # TAB 1: UPLOAD FONTS
        # =====================================
        with gr.TabItem("📤 Upload Fonts"):
            gr.Markdown("### Upload New Fonts")
            
            with gr.Row():
                with gr.Column(scale=1):
                    font_file = gr.File(
                        label="Select Font File",
                        file_types=[".ttf", ".otf"],
                        file_count="single"
                    )
                    
                    font_name = gr.Textbox(
                        label="Font Name",
                        placeholder="e.g., TikTok Sans Bold",
                        info="Enter a descriptive name for the font"
                    )
                    
                    font_category = gr.Dropdown(
                        label="Category",
                        choices=["english", "chinese", "japanese", "korean", "display", "handwriting", "other"],
                        value="english"
                    )
                    
                    font_description = gr.Textbox(
                        label="Description (Optional)",
                        placeholder="Describe the font style...",
                        lines=2
                    )
                    
                    with gr.Row():
                        upload_btn = gr.Button("🚀 Upload Font", variant="primary")
                        clear_btn = gr.Button("🗑️ Clear", variant="secondary")
                    
                    upload_output = gr.Textbox(label="Upload Status", interactive=False)
                
                with gr.Column(scale=1):
                    gr.Markdown("### Preview Area")
                    preview_image = gr.Image(label="Font Preview", height=300)
                    
                    preview_text = gr.Textbox(
                        label="Preview Text",
                        value="ABCDEFGHIJKLMNOPQRSTUVWXYZ\nabcdefghijklmnopqrstuvwxyz\n0123456789\n中文字体测试",
                        lines=3
                    )
                    
                    preview_size = gr.Slider(
                        label="Preview Size",
                        minimum=12,
                        maximum=72,
                        value=36,
                        step=2
                    )
                    
                    preview_btn = gr.Button("👁️ Preview Font")
            
            # Preview function
            def update_preview(file, text, size):
                if file is None:
                    return None
                try:
                    # Handle different file types
                    if hasattr(file, 'read'):
                        # It's a file-like object
                        with tempfile.NamedTemporaryFile(delete=False, suffix='.ttf') as tmp:
                            shutil.copyfileobj(file, tmp)
                            tmp_path = tmp.name
                    elif isinstance(file, str) and os.path.exists(file):
                        # It's a file path
                        tmp_path = file
                    else:
                        return None
                    
                    # Create preview image
                    img = Image.new('RGB', (800, 400), color='white')
                    d = ImageDraw.Draw(img)
                    try:
                        font = ImageFont.truetype(tmp_path, size)
                    except:
                        font = ImageFont.load_default()
                    
                    d.text((50, 50), text, fill='black', font=font)
                    
                    preview_path = os.path.join(PREVIEW_DIR, f"preview_{datetime.now().timestamp()}.png")
                    img.save(preview_path)
                    
                    # Clean up temp file if we created one
                    if hasattr(file, 'read') and os.path.exists(tmp_path):
                        os.unlink(tmp_path)
                    
                    return preview_path
                except Exception as e:
                    print(f"Preview error: {e}")
                    return None
            
            preview_btn.click(
                fn=update_preview,
                inputs=[font_file, preview_text, preview_size],
                outputs=[preview_image]
            )
            
            # Upload function
            upload_btn.click(
                fn=upload_font,
                inputs=[font_file, font_category, font_name, font_description],
                outputs=[upload_output]
            ).then(
                fn=lambda: (None, "", "", "", None),
                inputs=[],
                outputs=[font_file, font_name, font_description, upload_output, preview_image]
            )
            
            # Clear button
            clear_btn.click(
                fn=lambda: (None, "", "", "", None),
                inputs=[],
                outputs=[font_file, font_name, font_description, upload_output, preview_image]
            )
        
        # =====================================
        # TAB 2: BROWSE FONTS
        # =====================================
        with gr.TabItem("🔍 Browse Fonts"):
            gr.Markdown("### Font Library")
            
            refresh_btn = gr.Button("🔄 Refresh List", variant="secondary")
            
            category_filter_dropdown = gr.Dropdown(
                label="Filter by Category",
                choices=["all"] + categories,
                value="all"
            )
            
            font_gallery = gr.Gallery(
                label="Fonts",
                columns=4,
                object_fit="contain",
                height="auto"
            )
            
            with gr.Row():
                selected_font = gr.Dropdown(
                    label="Select Font",
                    choices=[f["name"] for f in fonts],
                    interactive=True
                )
                
                delete_btn = gr.Button("🗑️ Delete Selected Font", variant="stop")
                delete_output = gr.Textbox(label="Delete Status", interactive=False)
            
            font_preview_big = gr.Image(label="Font Preview", height=300)
            font_info = gr.JSON(label="Font Details")
            
            def load_font_gallery(filter_cat):
                fonts, _ = get_font_list()
                images = []
                
                for font in fonts:
                    if filter_cat != "all" and font["category"] != filter_cat:
                        continue
                    
                    preview = generate_preview(font["url"], font_size=24)
                    if preview:
                        images.append(preview)
                
                return images
        
            def show_font_details(font_name):
                fonts, _ = get_font_list()
                for font in fonts:
                    if font["name"] == font_name:
                        preview = generate_preview(font["url"], font_size=36)
                        return preview, font
                return None, {}
            
            def delete_selected_font(font_name):
                fonts, _ = get_font_list()
                for font in fonts:
                    if font["name"] == font_name:
                        return delete_font(font["path"])
                return "❌ Font not found"
            
            def update_font_dropdown():
                fonts, _ = get_font_list()
                return gr.Dropdown(choices=[f["name"] for f in fonts])
            
            refresh_btn.click(
                fn=load_font_gallery,
                inputs=[category_filter_dropdown],
                outputs=[font_gallery]
            ).then(
                fn=update_font_dropdown,
                outputs=[selected_font]
            )
            
            category_filter_dropdown.change(
                fn=load_font_gallery,
                inputs=[category_filter_dropdown],
                outputs=[font_gallery]
            )
            
            selected_font.change(
                fn=show_font_details,
                inputs=[selected_font],
                outputs=[font_preview_big, font_info]
            )
            
            delete_btn.click(
                fn=delete_selected_font,
                inputs=[selected_font],
                outputs=[delete_output]
            ).then(
                fn=load_font_gallery,
                inputs=[category_filter_dropdown],
                outputs=[font_gallery]
            ).then(
                fn=update_font_dropdown,
                outputs=[selected_font]
            )
        
        # =====================================
        # TAB 3: CATEGORY MANAGEMENT
        # =====================================
        with gr.TabItem("📁 Categories"):
            gr.Markdown("### Manage Font Categories")
            
            def get_categories():
                _, cats = get_font_list()
                return {"categories": cats}
            
            current_categories = gr.JSON(
                label="Current Categories",
                value=get_categories()
            )
            
            with gr.Row():
                new_category = gr.Textbox(
                    label="New Category Name",
                    placeholder="e.g., handwriting, decorative"
                )
                add_category_btn = gr.Button("➕ Add Category", variant="primary")
            
            add_category_output = gr.Textbox(label="Result", interactive=False)
            
            def add_category(cat_name):
                if not cat_name:
                    return "❌ Please enter a category name"
                
                try:
                    placeholder_path = f"{cat_name}/.keep"
                    with tempfile.NamedTemporaryFile(mode='w', delete=False) as f:
                        f.write("# This directory exists")
                        temp_path = f.name
                    
                    api.upload_file(
                        path_or_fileobj=temp_path,
                        path_in_repo=placeholder_path,
                        repo_id=FONTS_DATASET,
                        repo_type="dataset"
                    )
                    os.unlink(temp_path)
                    return f"✅ Category '{cat_name}' created"
                except Exception as e:
                    return f"❌ Failed: {str(e)}"
            
            add_category_btn.click(
                fn=add_category,
                inputs=[new_category],
                outputs=[add_category_output]
            ).then(
                fn=None,
                inputs=None,
                outputs=None,
                js="() => { setTimeout(() => window.location.reload(), 2000); }"
            )
        
        # =====================================
        # TAB 4: DATASET INFO
        # =====================================
        with gr.TabItem("ℹ️ Dataset Info"):
            gr.Markdown("### Dataset Information")
            
            def get_dataset_stats():
                fonts, cats = get_font_list()
                total_size = 0
                category_counts = {}
                
                for cat in cats:
                    category_counts[cat] = len([f for f in fonts if f["category"] == cat])
                
                for font in fonts:
                    try:
                        response = requests.head(font["url"])
                        total_size += int(response.headers.get('content-length', 0))
                    except:
                        pass
                
                return {
                    "total_fonts": len(fonts),
                    "categories": cats,
                    "fonts_per_category": category_counts,
                    "total_size_mb": round(total_size / (1024 * 1024), 2),
                    "dataset_url": f"https://huggingface.co/datasets/{FONTS_DATASET}"
                }
            
            stats_btn = gr.Button("📊 Refresh Stats")
            stats_display = gr.JSON(label="Dataset Statistics")
            
            stats_btn.click(
                fn=get_dataset_stats,
                inputs=[],
                outputs=[stats_display]
            )
            
            gr.Markdown("### Quick Links")
            gr.Markdown(f"- [View Dataset on HF](https://huggingface.co/datasets/{FONTS_DATASET})")
    
    # =========================================
    # INITIAL LOAD
    # =========================================
    demo.load(
        fn=lambda: load_font_gallery("all"),
        inputs=[],
        outputs=[font_gallery]
    )

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)