app_logic.py

import os
import re
import tempfile
import shutil
import git
import re

from huggingface_hub import (
    create_repo,
    upload_folder,
    list_repo_files,
    Repository,
    whoami,
    hf_hub_download,
)
import logging
from pathlib import Path
from PIL import Image

# --- Import functions from keylock_decode ---
try:
    from keylock_decode import (
        decode_data_from_image_pil,
        save_decoded_data_locally_encrypted,
        load_decoded_data_locally_encrypted
    )
    KEYLOCK_DECODE_AVAILABLE = True
except ImportError:
    KEYLOCK_DECODE_AVAILABLE = False
    decode_data_from_image_pil = None
    save_decoded_data_locally_encrypted = None
    load_decoded_data_locally_encrypted = None
    logging.warning("keylock-decode library not available or missing functions. KeyLock Wallet features will be disabled.")


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

# --- Global variable for in-memory keys and local storage path ---
_loaded_local_keys = {}
LOCAL_KEYS_FILE = "/data/.space_keys.enc" # Persistent storage in Hugging Face Spaces

# --- Helper Function: Load Local Keys into Memory ---
def _load_keys_into_memory(password: str) -> list[str]:
    # Loads and decrypts keys from the local file into memory
    global _loaded_local_keys
    status_messages = []

    if not KEYLOCK_DECODE_AVAILABLE or not load_decoded_data_locally_encrypted:
        status_messages.append("KeyLock-Decode library not available for local loading.")
        _loaded_local_keys = {}
        return status_messages

    if not password:
        status_messages.append("Error: Password required to load local keys.")
        _loaded_local_keys = {}
        return status_messages

    if not Path(LOCAL_KEYS_FILE).exists():
         status_messages.append(f"Info: Local key file not found at `{LOCAL_KEYS_FILE}`. No keys loaded.")
         _loaded_local_keys = {}
         return status_messages

    try:
        _loaded_local_keys = load_decoded_data_locally_encrypted(password, LOCAL_KEYS_FILE)

        if _loaded_local_keys:
            status_messages.append(f"Successfully loaded {len(_loaded_local_keys)} keys from local storage.")
            masked_keys = {k: ('********' if any(k_word in k.upper() for k_word in ['TOKEN', 'KEY', 'SECRET', 'PASS']) else v) for k, v in _loaded_local_keys.items()}
            status_messages.append(f"Loaded keys: {masked_keys}")
            logger.info(f"Keys loaded into memory from {LOCAL_KEYS_FILE}")
        else:
             status_messages.append(f"Info: Local key file `{LOCAL_KEYS_FILE}` was empty or contained no valid data after decryption.")
             _loaded_local_keys = {}

    except FileNotFoundError:
        status_messages.append(f"Info: Local key file not found at {LOCAL_KEYS_FILE}. No keys loaded.")
        _loaded_local_keys = {}
    except ValueError as e:
        status_messages.append(f"Error loading local keys: {e}. Password may be incorrect.")
        _loaded_local_keys = {}
    except IOError as e:
         status_messages.append(f"IO Error loading local keys: {e}")
         _loaded_local_keys = {}
    except Exception as e:
        logger.exception(f"Unexpected error loading local keys from {LOCAL_KEYS_FILE}:")
        status_messages.append(f"Unexpected error loading local keys: {str(e)}")
        _loaded_local_keys = {}

    return status_messages


# --- Helper Function: Get API Token ---
def _get_api_token(ui_token_from_textbox=None):
    # Retrieves token from in-memory keys, env vars, or textbox
    in_memory_token = _loaded_local_keys.get('HF_TOKEN')
    if in_memory_token:
        logger.info("Using HF_TOKEN from in-memory loaded keys.")
        return in_memory_token, None

    env_token = os.getenv('HF_TOKEN')
    if env_token:
        logger.info("Using HF_TOKEN from environment variable.")
        return env_token, None

    if ui_token_from_textbox:
        logger.info("Using HF_TOKEN from UI textbox.")
        return ui_token_from_textbox, None

    return None, "Error: Hugging Face API token not provided."


# --- Main Function: Process KeyLock Image and Store Locally ---
def process_keylock_image_and_store_locally(image_pil_object: Image.Image, password: str):
    # Decodes image, saves data encrypted locally, and loads into memory
    status_messages_display = []

    if not KEYLOCK_DECODE_AVAILABLE or not decode_data_from_image_pil or \
       not save_decoded_data_locally_encrypted or not load_decoded_data_locally_encrypted:
        status_messages_display.append("Error: KeyLock-Decode library not available.")
        global _loaded_local_keys
        _loaded_local_keys = {}
        return "\n".join(status_messages_display)

    if image_pil_object is None:
        status_messages_display.append("Error: No KeyLock Wallet image provided.")
        return "\n".join(status_messages_display)

    if not password:
        status_messages_display.append("Error: Password cannot be empty.")
        return "\n".join(status_messages_display)

    decoded_data = None
    try:
        logger.info(f"Attempting to decode from KeyLock Wallet image...")
        decoded_data, status_msgs_from_lib = decode_data_from_image_pil(image_pil_object, password)
        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}")

            try:
                save_decoded_data_locally_encrypted(decoded_data, password, LOCAL_KEYS_FILE)
                status_messages_display.append(f"\n**SUCCESS: Decoded data saved to encrypted local file:** `{LOCAL_KEYS_FILE}`")

                load_status = _load_keys_into_memory(password)
                status_messages_display.extend(load_status)
                if not _loaded_local_keys:
                    status_messages_display.append("Warning: No keys were loaded into memory after saving.")


            except Exception as e:
                 status_messages_display.append(f"Error saving or loading data locally: {str(e)}")
                 logger.error(f"Error during local save/load after decoding: {e}")
                 global _loaded_local_keys
                 _loaded_local_keys = {}

        elif not status_msgs_from_lib:
             status_messages_display.append("Info: Decoding process completed, but no data was extracted.")


    except ValueError as e:
        status_messages_display.append(f"**Decoding Error:** {e}. Please check password and image.")
        global _loaded_local_keys
        _loaded_local_keys = {}
    except Exception as e:
        status_messages_display.append(f"**Unexpected error during processing:** {str(e)}")
        logger.exception("Unexpected error during keylock image processing:")
        global _loaded_local_keys
        _loaded_local_keys = {}


    return "\n".join(status_messages_display)

# --- Function: Load Keys from Local File ---
def load_keys_from_local_file(password: str) -> str:
    # Triggers loading keys from the encrypted local file into memory
    status_messages = _load_keys_into_memory(password)
    if not status_messages:
         if _loaded_local_keys:
             return f"Keys successfully loaded from `{LOCAL_KEYS_FILE}`."
         else:
             return f"Attempted to load keys from `{LOCAL_KEYS_FILE}`, but no keys were loaded. Check password."

    return "\n".join(status_messages)


# --- Markdown Processing Functions ---
def process_commented_markdown(commented_input):
    # Processes markdown by stripping comments if a specific marker is present
    lines = commented_input.strip().split("\n")
    if any( "# # Space:" in line.strip() for line in lines):
        cleaned_lines = [line[2:] if line.startswith("# ") else line for line in lines]
        return cleaned_lines
    return lines

def parse_markdown(markdown_input):
    # Parses markdown input to extract space info and file content
    space_info = {"repo_name_md": "", "owner_md": "", "files": []}
    current_file_path = None; current_file_content_lines = []
    in_file_definition = False; in_code_block = False
    lines = process_commented_markdown(markdown_input)

    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:
            if line_content_stripped.startswith("# Space:"):
                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
            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

# --- Helper Function: Determine Repository ID ---
def _determine_repo_id(ui_api_token_from_textbox, space_name_ui, owner_ui):
    # Determines the full repo ID (owner/space_name)
    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 Owner field."
    final_owner = owner_ui; error_message = None
    if not final_owner:
        resolved_api_token, token_err = _get_api_token(ui_api_token_from_textbox)
        if token_err: return None, token_err
        if not resolved_api_token: return None, "Error: API token required for auto owner determination."
        try:
            user_info = whoami(token=resolved_api_token)
            if user_info and 'name' in user_info: final_owner = user_info['name']
            else: error_message = "Error: Could not retrieve username. Check token/permissions."
        except Exception as e: error_message = f"Error retrieving username: {str(e)}."
        if error_message: return None, error_message
    if not final_owner: return None, "Error: Owner could not be determined."
    return f"{final_owner}/{space_name_ui}", None

# --- Function: Fetch File Content from Hub ---
def get_space_file_content(ui_api_token_from_textbox, space_name_ui, owner_ui, file_path_in_repo):
    # Fetches content of a specific file from a Hugging Face Space
    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 None, token_err

        repo_id, err = _determine_repo_id(ui_api_token_from_textbox, space_name_ui, owner_ui)
        if err:
            return None, err
        repo_id_for_error_logging = repo_id

        if not file_path_in_repo:
            return None, "Error: File path cannot be empty."

        logger.info(f"Attempting to download file: {file_path_in_repo} from Space: {repo_id}")
        downloaded_file_path = hf_hub_download(
            repo_id=repo_id,
            filename=file_path_in_repo,
            repo_type="space",
            token=resolved_api_token,
        )

        content = Path(downloaded_file_path).read_text(encoding="utf-8")
        logger.info(f"Successfully downloaded and read file: {file_path_in_repo}")
        return content, None

    except Exception as e:
        if "404" in str(e) or "not found" in str(e).lower():
             logger.warning(f"File not found {file_path_in_repo} in {repo_id_for_error_logging}: {e}")
             return None, f"Error: File '{file_path_in_repo}' not found in Space '{repo_id_for_error_logging}'."
        logger.exception(f"Error fetching file content for {file_path_in_repo}:")
        return None, f"Error fetching file content: {str(e)}"

# --- Function: List Space Files ---
def list_space_files_for_browsing(ui_api_token_from_textbox, space_name_ui, owner_ui):
    # Lists files in a Hugging Face Space
    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 None, token_err

        repo_id, err = _determine_repo_id(ui_api_token_from_textbox, space_name_ui, owner_ui)
        if err: return None, err
        repo_id_for_error_logging = repo_id

        files = list_repo_files(repo_id=repo_id, token=resolved_api_token, repo_type="space")
        if not files:
            return [], f"No files found in Space `{repo_id}`."
        return files, None
    except Exception as e:
        logger.exception(f"Error listing files for {repo_id_for_error_logging}:")
        return None, f"Error listing files for `{repo_id_for_error_logging}`: {str(e)}"

# --- Core Function: Create/Update Space ---
def create_space(ui_api_token_from_textbox, space_name_ui, owner_ui, sdk_ui, markdown_input):
    # Creates or updates a Hugging Face Space with files from markdown
    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
        space_info = parse_markdown(markdown_input)
        if not space_info["files"]: return "Error: No files found in markdown."
        with tempfile.TemporaryDirectory() as temp_dir:
            repo_staging_path = Path(temp_dir) / "repo_staging_content"
            repo_staging_path.mkdir(exist_ok=True)
            for file_info in space_info["files"]:
                if not file_info.get("path"): continue
                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"])
            try:
                create_repo(repo_id=repo_id, token=resolved_api_token, repo_type="space", space_sdk=sdk_ui, private=False)
            except Exception as e:
                err_str = str(e).lower()
                if not ("already exists" in err_str or "you already created this repo" in err_str or "exists" in err_str):
                    return f"Error creating Space '{repo_id}': {str(e)}"
            upload_folder(repo_id=repo_id, folder_path=str(repo_staging_path), path_in_repo=".", token=resolved_api_token, repo_type="space", commit_message=f"Initial Space setup of {repo_id} via Builder")
            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)}"

# --- Core Function: Update Space File ---
def update_space_file(ui_api_token_from_textbox, space_name_ui, owner_ui, file_path_in_repo, file_content, commit_message_ui):
    # Updates a specific file in a Hugging Face Space
    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 to update cannot be empty."
        file_path_in_repo = file_path_in_repo.lstrip('/').replace(os.sep, '/')
        commit_message_ui = commit_message_ui or f"Update {file_path_in_repo} via Space Builder"
        with tempfile.TemporaryDirectory() as temp_dir_for_update:
            repo_local_clone_path = Path(temp_dir_for_update) / "update_clone"
            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="[email protected]")
            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)
            cloned_repo.push_to_hub(commit_message=commit_message_ui)
            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}, file {file_path_in_repo}:")
        return f"Error updating file for `{repo_id_for_error_logging}`: {str(e)}"