app_logic.py

import os
import re
import tempfile
import git # Used by Repository
from huggingface_hub import (
    create_repo,
    upload_folder,
    list_repo_files,
    Repository, # Used in update_space_file
    whoami,
)
import logging
from pathlib import Path
from PIL import Image # For type hinting from gr.Image(type="pil")

# Attempt to import keylock_decode
try:
    from keylock_decode import decode_from_image_pil
    KEYLOCK_DECODE_AVAILABLE = True
except ImportError:
    KEYLOCK_DECODE_AVAILABLE = False
    decode_from_image_pil = None # Placeholder
    logging.warning("keylock-decode library not found. Image decoding feature will be disabled.")


# Configure logging (ensure this is suitable for your environment, might be set in app.py too)
# If keylock-decode uses its own logger (e.g., logging.getLogger('keylock_decode')),
# you might need to adjust its level separately if you want more/less verbose output from it.
logging.basicConfig(
    level=logging.INFO, 
    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
)
logger = logging.getLogger(__name__)
if KEYLOCK_DECODE_AVAILABLE: # Example: Set keylock_decode's logger to DEBUG if you want its verbose output
    # logging.getLogger('keylock_decode.decoder').setLevel(logging.DEBUG) # Check actual logger name in keylock-decode
    pass


# --- New Helper Function to Get API Token ---
def _get_api_token(ui_token_from_textbox=None):
    """
    Retrieves the Hugging Face API token.
    Prioritizes HF_TOKEN from environment variables (potentially set by KeyLock-Decode),
    then falls back to UI input.
    Returns (token, error_message)
    """
    env_token = os.getenv('HF_TOKEN')
    if env_token:
        logger.info("Using HF_TOKEN from environment.")
        return env_token, None
    
    if ui_token_from_textbox:
        logger.info("Using API token from UI textbox.")
        return ui_token_from_textbox, None
    
    logger.warning("HF API token not found in environment or UI textbox.")
    return None, "Error: Hugging Face API token not provided. Please enter it in the textbox or load it from an image."


# --- New Function for KeyLock-Decode Integration ---
def load_token_from_image_and_set_env(image_pil_object: Image.Image, password: str):
    """
    Decodes data from a PIL Image object using keylock-decode, sets environment variables,
    and returns status messages.
    """
    if not KEYLOCK_DECODE_AVAILABLE:
        return "Error: KeyLock-Decode library is not installed. This feature is disabled."

    if image_pil_object is None:
        return "Error: No image provided for decoding."
    if not password:
        return "Error: Password cannot be empty for image decoding."

    status_messages_display = []
    
    try:
        logger.info(f"Attempting to decode from image using KeyLock-Decode...")
        # decode_from_image_pil expects a PIL Image object
        decoded_data, status_msgs_from_lib = decode_from_image_pil(
            image_pil_object,
            password,
            set_environment_variables=True # CRITICAL: This sets os.environ
        )

        status_messages_display.extend(status_msgs_from_lib)

        if decoded_data:
            status_messages_display.append("\n**Decoded Data Summary (sensitive values masked):**")
            for key, value in decoded_data.items():
                display_value = '********' if any(k_word in key.upper() for k_word in ['TOKEN', 'KEY', 'SECRET', 'PASS']) else value
                status_messages_display.append(f"- {key}: {display_value}")
            
            if os.getenv('HF_TOKEN'):
                status_messages_display.append(f"\n**SUCCESS: HF_TOKEN was found and has been set in the current process environment.** It will be prioritized for operations.")
            elif 'HF_TOKEN' in decoded_data: # Should be rare if set_env_vars=True
                 status_messages_display.append(f"\nWarning: HF_TOKEN was in decoded data but os.getenv('HF_TOKEN') is not picking it up. This is unexpected if 'set_environment_variables=True' worked.")
            else:
                status_messages_display.append("\nNote: HF_TOKEN was not specifically found in the decoded data from the image.")
        else:
            # This case is often covered by status_msgs_from_lib, but good to have a fallback.
            if not any("No payload data found" in msg for msg in status_msgs_from_lib):
                 status_messages_display.append("No data was decoded, or the decoded data was empty.")

    except ValueError as e: # Specific errors from keylock-decode (e.g., bad password, corrupted data)
        logger.error(f"KeyLock-Decode ValueError: {e}")
        status_messages_display.append(f"**Decoding Error:** {e}")
    except Exception as e:
        logger.exception("An unexpected error occurred during image decoding with KeyLock-Decode:")
        status_messages_display.append(f"**An unexpected error occurred:** {str(e)}")
    
    return "\n".join(status_messages_display)


# --- Original `parse_markdown` function (unchanged) ---
def parse_markdown(markdown_input):
    """Parse markdown input to extract space details and file structure."""
    space_info = {"repo_name_md": "", "owner_md": "", "files": []}
    current_file = None
    file_content = []
    in_file_content = False 
    in_code_block = False 

    lines = markdown_input.strip().split("\n")
    for line_idx, line_content_orig in enumerate(lines):
        line_content_stripped = line_content_orig.strip()

        if in_file_content:
            if line_content_stripped.startswith("```"):
                if in_code_block: 
                    file_content.append(line_content_orig) 
                    in_code_block = False
                else: 
                    in_code_block = True
                    file_content.append(line_content_orig)
            elif in_code_block: 
                file_content.append(line_content_orig)
            elif not in_code_block: 
                if line_content_stripped.startswith("### File:") or line_content_stripped.startswith("## File Structure") or line_content_stripped.startswith("# Space:"):
                    if current_file and file_content:
                        space_info["files"].append({"path": current_file, "content": "\n".join(file_content)})
                    current_file = None 
                    file_content = []
                    in_file_content = False 
                else: 
                     file_content.append(line_content_orig)

        if line_content_stripped.startswith("# Space:"):
            if current_file and file_content: 
                space_info["files"].append({"path": current_file, "content": "\n".join(file_content)})
            full_space_name_md = line_content_stripped.replace("# Space:", "").strip()
            if "/" in full_space_name_md:
                space_info["owner_md"], space_info["repo_name_md"] = full_space_name_md.split("/", 1)
            else:
                space_info["repo_name_md"] = full_space_name_md
            current_file = None
            file_content = []
            in_file_content = False
            in_code_block = False

        elif line_content_stripped.startswith("## File Structure"):
            if current_file and file_content: 
                space_info["files"].append({"path": current_file, "content": "\n".join(file_content)})
            current_file = None
            file_content = []
            in_file_content = False
            in_code_block = False
            continue 

        elif line_content_stripped.startswith("### File:"):
            if current_file and file_content: 
                space_info["files"].append({"path": current_file, "content": "\n".join(file_content)})
            
            current_file = line_content_stripped.replace("### File:", "").strip()
            file_content = [] 
            in_file_content = True 
            in_code_block = False 
            
    if current_file and file_content:
        space_info["files"].append({"path": current_file, "content": "\n".join(file_content)})
    
    space_info["files"] = [f for f in space_info["files"] if f.get("path")]
    return space_info


# --- Updated `_determine_repo_id` to use `_get_api_token` ---
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 '/'. Please use the Owner field for the namespace."

    final_owner = owner_ui
    error_message = None

    if not final_owner:
        # Get the actual token to use for whoami
        resolved_api_token, token_err = _get_api_token(ui_api_token_from_textbox)
        if token_err:
            return None, token_err 
        # If resolved_api_token is None here but token_err is None, it means _get_api_token had an issue not returning error string
        # which _get_api_token is designed not to do, it should always return error if no token.
        if not resolved_api_token: 
             return None, "Error: API token is required to automatically determine owner when Owner field is empty (internal check failed)."

        try:
            user_info = whoami(token=resolved_api_token)
            if user_info and 'name' in user_info:
                final_owner = user_info['name']
                logger.info(f"Determined owner: {final_owner} using API token.")
            else:
                logger.error(f"whoami(token=...) returned: {user_info} - 'name' field missing or user_info is None.")
                error_message = "Error: Could not retrieve username from API token. Ensure token is valid and has 'Read profile' permissions. Or, specify Owner manually."
        except Exception as e:
            logger.error(f"Error calling whoami for owner: {str(e)}")
            error_message = f"Error retrieving username from API token: {str(e)}. Please specify Owner manually."
        
        if error_message:
            return None, error_message

    if not final_owner:
        return None, "Error: Owner could not be determined. Please provide an owner or ensure your API token is valid."

    return f"{final_owner}/{space_name_ui}", None


# --- Updated Core Functions to use `_get_api_token` ---

def create_space(ui_api_token_from_textbox, space_name_ui, owner_ui, sdk_ui, markdown_input):
    repo_id_for_error_logging = f"{owner_ui}/{space_name_ui}" if owner_ui else space_name_ui
    try:
        resolved_api_token, token_err = _get_api_token(ui_api_token_from_textbox)
        if token_err:
            return token_err

        # _determine_repo_id will also call _get_api_token internally if owner_ui is blank
        repo_id, err = _determine_repo_id(ui_api_token_from_textbox, space_name_ui, owner_ui)
        if err:
            return err
        repo_id_for_error_logging = repo_id # Update for more accurate logging if determination succeeds
        
        space_info = parse_markdown(markdown_input)
        if not space_info["files"]:
            return "Error: No files found in the markdown input. Ensure '### File: path/to/file.ext' markers are used correctly."

        with tempfile.TemporaryDirectory() as temp_dir:
            repo_local_path = Path(temp_dir) / "repo_upload_content"
            repo_local_path.mkdir(exist_ok=True)

            for file_info in space_info["files"]:
                if not file_info.get("path"): continue
                file_path_abs = repo_local_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"])
                logger.info(f"Wrote file: {file_path_abs}")

            try:
                create_repo(
                    repo_id=repo_id, token=resolved_api_token, repo_type="space",
                    space_sdk=sdk_ui, private=False,
                )
                logger.info(f"Created Space repo: {repo_id}")
            except Exception as e:
                err_str = str(e).lower()
                if "already exists" in err_str or "you already created this repo" in err_str or "exists" in err_str:
                    logger.info(f"Space {repo_id} already exists. Proceeding to upload/update files.")
                else:
                    return f"Error creating Space '{repo_id}': {str(e)}"
            
            upload_folder(
                repo_id=repo_id, folder_path=str(repo_local_path), path_in_repo=".",
                token=resolved_api_token, repo_type="space",
                commit_message=f"Initial Space setup of {repo_id} via Builder",
            )
            logger.info(f"Uploaded files to Space: {repo_id}")
            return f"Successfully created/updated Space: [{repo_id}](https://huggingface.co/spaces/{repo_id})"

    except Exception as e:
        logger.exception(f"Error in create_space for {repo_id_for_error_logging}:")
        return f"Error during Space creation/update: {str(e)}"


def view_space_files(ui_api_token_from_textbox, space_name_ui, owner_ui):
    repo_id_for_error_logging = f"{owner_ui}/{space_name_ui}" if owner_ui else space_name_ui
    try:
        resolved_api_token, token_err = _get_api_token(ui_api_token_from_textbox)
        if token_err:
            return token_err
        
        repo_id, err = _determine_repo_id(ui_api_token_from_textbox, space_name_ui, owner_ui)
        if err:
            return err
        repo_id_for_error_logging = repo_id
            
        files = list_repo_files(repo_id=repo_id, token=resolved_api_token, repo_type="space")
        if files:
            return f"Files in `{repo_id}`:\n\n" + "\n".join([f"- `{f}`" for f in files])
        return f"No files found in the Space `{repo_id}`."
    except Exception as e:
        logger.exception(f"Error in view_space_files for {repo_id_for_error_logging}:")
        return f"Error listing files for `{repo_id_for_error_logging}`: {str(e)}"


def update_space_file(ui_api_token_from_textbox, space_name_ui, owner_ui, file_path_in_repo, file_content, commit_message_ui):
    repo_id_for_error_logging = f"{owner_ui}/{space_name_ui}" if owner_ui else space_name_ui
    try:
        resolved_api_token, token_err = _get_api_token(ui_api_token_from_textbox)
        if token_err:
            return token_err

        repo_id, err = _determine_repo_id(ui_api_token_from_textbox, space_name_ui, owner_ui)
        if err:
            return err
        repo_id_for_error_logging = repo_id
        
        if not file_path_in_repo: return "Error: File Path cannot be empty."
        commit_message_ui = commit_message_ui or f"Update {file_path_in_repo} via Space Builder"

        with tempfile.TemporaryDirectory() as temp_dir:
            repo_local_clone_path = Path(temp_dir) / "cloned_space_repo_for_update"
            
            cloned_repo = Repository(
                local_dir=str(repo_local_clone_path),
                clone_from=f"https://huggingface.co/spaces/{repo_id}",
                repo_type="space", use_auth_token=resolved_api_token,
                git_user="Space Builder Bot", git_email="space-builder@huggingface.co"
            )
            logger.info(f"Cloned Space {repo_id} to {repo_local_clone_path}")

            full_local_file_path = Path(cloned_repo.local_dir) / file_path_in_repo
            full_local_file_path.parent.mkdir(parents=True, exist_ok=True)
            with open(full_local_file_path, "w", encoding="utf-8") as f: f.write(file_content)
            logger.info(f"Wrote updated file {file_path_in_repo} locally.")

            cloned_repo.push_to_hub(commit_message=commit_message_ui)
            logger.info(f"Pushed changes for {file_path_in_repo} to {repo_id}")
            
            return f"Successfully updated `{file_path_in_repo}` in Space [{repo_id}](https://huggingface.co/spaces/{repo_id})"

    except Exception as e:
        logger.exception(f"Error in update_space_file for {repo_id_for_error_logging}:")
        return f"Error updating file for `{repo_id_for_error_logging}`: {str(e)}"