app_logic.py

import os
import tempfile
import gradio as gr
from huggingface_hub import create_repo, upload_folder, whoami
import logging
from pathlib import Path
from PIL import Image
from keylock import core as keylock_core
from repo_to_md.core import markdown_to_files

logging.basicConfig(
    level=logging.INFO, 
    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
)
logger = logging.getLogger(__name__)

def get_username_from_token(api_token: str):
    if not api_token:
        logger.info("API token is empty, cannot fetch username.")
        return ""
    try:
        logger.info("Attempting to fetch username from provided API token.")
        user_info = whoami(token=api_token)
        username = user_info.get('name')
        if username:
            logger.info(f"Successfully fetched username: {username}")
            return username
        else:
            logger.warning("Token was valid, but no username found in whoami() response.")
            return ""
    except Exception as e:
        logger.error(f"Failed to validate token and fetch username: {e}")
        return ""

def _get_api_token(ui_token_from_textbox=None):
    env_token = os.getenv('HF_TOKEN')
    if env_token: return env_token, None
    if ui_token_from_textbox: return ui_token_from_textbox, None
    return None, "Error: Hugging Face API token not provided."

def extract_markdown_from_image(repo_image: Image.Image, image_password: str):
    if not repo_image:
        return None, "Error: Please upload a repository image."
    password_to_use = image_password if image_password is not None else ""
    try:
        logger.info("Starting extraction of markdown from image...")
        result = keylock_core.extract_repo_from_image(repo_image, password_to_use)
        markdown_content = result.get('markdown')
        if not markdown_content:
            return None, "Error: Extraction succeeded, but no markdown content was found in the image."
        logger.info("Successfully extracted markdown content.")
        return markdown_content, result.get('status', "Extraction successful.")
    except Exception as e:
        logger.error(f"Failed to extract from image: {e}")
        return None, f"Error during extraction: {e}"

def parse_markdown_for_build(markdown_input):
    space_info = {"files": []}
    current_file_path = None; current_file_content_lines = []
    in_file_definition = False; in_code_block = False
    lines = markdown_input.strip().split("\n")
    for line_content_orig in lines:
        line_content_stripped = line_content_orig.strip()
        if line_content_stripped.startswith("### File:"):
            if current_file_path and in_file_definition:
                space_info["files"].append({"path": current_file_path, "content": "\n".join(current_file_content_lines)})
            current_file_path = line_content_stripped.replace("### File:", "").strip()
            current_file_content_lines = []; in_file_definition = True; in_code_block = False
            continue
        if not in_file_definition: continue
        if line_content_stripped.startswith("```"):
            in_code_block = not in_code_block
            continue
        current_file_content_lines.append(line_content_orig)
    if current_file_path and in_file_definition:
        space_info["files"].append({"path": current_file_path, "content": "\n".join(current_file_content_lines)})
    space_info["files"] = [f for f in space_info["files"] if f.get("path")]
    return space_info

def _determine_repo_id(ui_api_token_from_textbox, space_name_ui, owner_ui):
    if not space_name_ui: return None, "Error: Space Name cannot be empty."
    if "/" in space_name_ui: return None, "Error: Space Name should not contain '/'. Use the Owner field."
    final_owner = owner_ui
    if not final_owner:
        resolved_api_token, token_err = _get_api_token(ui_api_token_from_textbox)
        if token_err: return None, token_err
        try:
            user_info = whoami(token=resolved_api_token)
            final_owner = user_info.get('name')
        except Exception as e:
            return None, f"Error retrieving username from token: {e}. Please specify the Owner manually."
    if not final_owner: return None, "Error: Owner could not be determined."
    return f"{final_owner}/{space_name_ui}", None

def create_space(ui_api_token_from_textbox, space_name_ui, owner_ui, sdk_ui, markdown_input):
    repo_id, err = _determine_repo_id(ui_api_token_from_textbox, space_name_ui, owner_ui)
    if err: return err
    try:
        resolved_api_token, token_err = _get_api_token(ui_api_token_from_textbox)
        if token_err: return token_err
        space_info = parse_markdown_for_build(markdown_input)
        if not space_info["files"]: return "Error: No files were found after parsing the markdown from the image."
        with tempfile.TemporaryDirectory() as temp_dir:
            repo_staging_path = Path(temp_dir) / "repo_staging"
            repo_staging_path.mkdir()
            for file_info in space_info["files"]:
                file_path_abs = repo_staging_path / file_info["path"]
                file_path_abs.parent.mkdir(parents=True, exist_ok=True)
                with open(file_path_abs, "w", encoding="utf-8") as f:
                    f.write(file_info["content"])
            create_repo(repo_id=repo_id, token=resolved_api_token, repo_type="space", space_sdk=sdk_ui, exist_ok=True)
            upload_folder(
                repo_id=repo_id,
                folder_path=str(repo_staging_path),
                path_in_repo=".",
                token=resolved_api_token,
                repo_type="space",
                commit_message="Initial Space creation from repo-embedded image"
            )
            return f"**Success!** Space created or updated: [{repo_id}](https://huggingface.co/spaces/{repo_id})"
    except Exception as e:
        logger.exception(f"Error in create_space for {repo_id}:")
        return f"Error during Space creation: {str(e)}"

def build_space_from_image(api_token, repo_image, image_password, new_space_name, owner, sdk):
    markdown_string, status_msg = extract_markdown_from_image(repo_image, image_password)
    if not markdown_string:
        return status_msg
    result = create_space(api_token, new_space_name, owner, sdk, markdown_string)
    return result

def preview_image_contents(repo_image, image_password):
    markdown_string, status_msg = extract_markdown_from_image(repo_image, image_password)
    if not markdown_string:
        return status_msg, gr.update(visible=False), gr.update(choices=[], value=None), "", []
    try:
        file_data, _ = markdown_to_files(markdown_string)
        if isinstance(file_data, str) or not file_data:
            raise ValueError("Markdown parsing failed or returned no files.")
        
        filenames = [f['filename'] for f in file_data]
        if not filenames:
            status = "✅ Preview generated, but no files were found in the markdown structure."
            return status, gr.update(visible=False), gr.update(choices=[], value=None), "", []

        first_filename = filenames
        first_content = file_data[0]['content']
        
        status = f"✅ Preview generated successfully. Found {len(filenames)} file(s)."
        
        return (
            status,
            gr.update(visible=True),
            gr.update(choices=filenames, value=first_filename),
            first_content,
            file_data
        )
    except Exception as e:
        logger.error(f"Error parsing markdown for preview: {e}", exc_info=True)
        error_msg = f"❌ Error: Could not parse files from the image's content. {e}"
        return error_msg, gr.update(visible=False), gr.update(choices=[], value=None), "", []

def update_file_preview(selected_filename, all_file_data):
    if not selected_filename or not all_file_data:
        return ""
    selected_file = next((f for f in all_file_data if f['filename'] == selected_filename), None)
    if selected_file:
        return selected_file.get('content', 'File content not found.')
    return f"Select file to view contents."