app.py

import gradio as gr
import re
import json
import os
import tempfile
import shlex
from huggingface_hub import HfApi

try:
    from build_logic import (
        # build_logic_create_space, # This is now handled by apply_staged_changes for AI
        _get_api_token as build_logic_get_api_token,
        whoami as build_logic_whoami,
        list_space_files_for_browsing,
        get_space_repository_info, # Used for initial load and status
        get_space_file_content,
        update_space_file, # Keep for manual editing
        parse_markdown as build_logic_parse_markdown,
        delete_space_file as build_logic_delete_space_file, # Keep for manual deletion
        get_space_runtime_status,
        apply_staged_changes, # NEW: Main function for applying AI changes
        build_logic_set_space_privacy, # Moved from app.py
        build_logic_delete_space # Moved from app.py
    )
    print("build_logic.py loaded successfully.")

    from model_logic import (
        get_available_providers,
        get_models_for_provider,
        get_default_model_for_provider,
        generate_stream
    )
    print("model_logic.py loaded successfully.")
except ImportError as e:
    print(f"Warning: Local modules (build_logic.py, model_logic.py) not found. Using dummy functions. Error: {e}")
    def get_available_providers(): return ["DummyProvider", "Groq"] # Added Groq for testing
    def get_models_for_provider(p):
        if p == 'Groq': return ["llama3-8b-8192", "gemma-7b-it"]
        return ["dummy-model"]
    def get_default_model_for_provider(p):
        if p == 'Groq': return "llama3-8b-8192"
        return "dummy-model"
    # The dummy function already accepts the api_key argument ('a')
    def generate_stream(p, m, a, msgs):
        yield f"Using dummy model. API Key provided: {'Yes' if a else 'No'}. This is a dummy response as local modules were not found.\n" + bbb + "text\n### File: dummy.txt\nHello from dummy model!\n" + bbb
    # Dummy build_logic functions
    def build_logic_create_space(*args, **kwargs): return "Error: build_logic not found (Dummy)."
    def build_logic_get_api_token(key): return (key or os.getenv("HF_TOKEN"), None)
    def build_logic_whoami(token): return {"name": "dummy_user"}
    def list_space_files_for_browsing(*args): return ([], "Error: build_logic not found (Dummy).")
    def get_space_repository_info(*args): return (None, [], "Error: build_logic not found (Dummy).")
    def get_space_file_content(*args): return ("", "Error: build_logic not found (Dummy).")
    def update_space_file(*args, **kwargs): return "Error: build_logic not found (Dummy)."
    def build_logic_parse_markdown(md):
        # Dummy parser attempts to find files for testing
        files = []
        file_pattern = re.compile(r"### File:\s*(?P<filename_line>[^\n]+)\n(?:```(?P<lang>[\w\.\-\+]*)\n(?P<code>[\s\S]*?)\n```|(?P<binary_msg>\[Binary file(?: - [^\]]+)?\]))")
        for match in file_pattern.finditer(md):
            filename = _clean_filename(match.group("filename_line"))
            if filename: files.append({"path": filename, "content": match.group("code") or match.group("binary_msg") or ""})
        return {"repo_name_md": "dummy/space", "owner_md": "dummy", "files": files}
    def build_logic_delete_space_file(*args): return "Error: build_logic not found (Dummy)."
    def get_space_runtime_status(*args): return ({"stage": "DUMMY", "hardware": "dummy", "status": "dummy"}, "Error: build_logic not found (Dummy).")
    def apply_staged_changes(*args, **kwargs): return "Error: apply_staged_changes not found (Dummy).", [], None
    def build_logic_set_space_privacy(*args): return "Error: build_logic_set_space_privacy not found (Dummy)."
    def build_logic_delete_space(*args): return "Error: build_logic_delete_space not found (Dummy)."

# --- END: Dummy functions ---


# --- CORE FIX: Define triple backticks safely to prevent Markdown rendering issues ---
backtick = chr(96)
bbb = f'{backtick}{backtick}{backtick}'

# State variable to hold the *current* representation of the Space's files and structure.
# This is populated on load and updated by AI outputs or manual edits/deletes.
parsed_code_blocks_state_cache = []
BOT_ROLE_NAME = "assistant"

DEFAULT_SYSTEM_PROMPT = f"""You are an expert AI programmer and Hugging Face assistant. Your role is to generate code and file structures based on user requests, or to modify existing code provided by the user.

You operate in the context of a Hugging Face Space. You can propose file changes and trigger specific actions.

**File and Code Formatting:**
When you provide NEW code for a file, or MODIFIED code for an existing file, use the following format exactly:
### File: path/to/filename.ext
(You can add a short, optional, parenthesized description after the filename on the SAME line)
{bbb}language
# Your full code here
{bbb}

If the file is binary, or you cannot show its content, use this format:
### File: path/to/binaryfile.ext
[Binary file - approximate_size bytes]

When you provide a project file structure, use this format (for informational purposes, not parsed for changes):
## File Structure
{bbb}
📁 Root
  📄 file1.py
  📁 subfolder
    📄 file2.js
{bbb}

**Instructions and Rules:**
- The role name for your responses in the chat history must be '{BOT_ROLE_NAME}'.
- Adhere strictly to these formatting instructions.
- If you update a file, provide the FULL file content again under the same filename.
- Only the latest version of each file mentioned throughout the chat will be used for the final output. The system will merge your changes with the prior state before applying.

**Hugging Face Space Actions:**
To perform direct actions on the Hugging Face Space, use the `### HF_ACTION: COMMAND arguments...` command on a single line.
The system will parse these actions from your response and present them to the user for confirmation before executing.

Available commands:
- `CREATE_SPACE <owner>/<repo_name> --sdk <sdk> --private <true|false>`: Creates a new, empty space. SDK can be gradio, streamlit, docker, or static. Private is optional and defaults to false. Use this ONLY when the user explicitly asks to create a *new* space with a specific name. When using this action, also provide the initial file structure (e.g., app.py, README.md) using `### File:` blocks in the same response. The system will apply these files to the new space.
- `DELETE_FILE path/to/file.ext`: Deletes a specific file from the current space.
- `SET_PRIVATE <true|false>`: Sets the privacy for the current space.
- `DELETE_SPACE`: Deletes the entire current space. THIS IS PERMANENT AND REQUIRES CAUTION. Only use this if the user explicitly and clearly asks to delete the space.

You can issue multiple file updates and action commands in a single response. The system will process all of them into a single change plan.

**Current Space Context:**
You will be provided with the current state of the files in the Space the user is interacting with. Use this information to understand the current project structure and content before proposing changes or actions. This context will appear after the user's message, starting with "## Current Space Context:". Do NOT include this context in your response. Only generate your response based on the user's request and the formatting rules above.

If no code or actions are requested, respond conversationally and help the user understand the Space Commander's capabilities.
"""

# --- Helper Functions ---
# Keep existing helper functions (_infer_lang_from_filename, _clean_filename, etc.)
# Refine _parse_chat_stream_logic to integrate with the state cache

def escape_html_for_markdown(text):
    if not isinstance(text, str): return ""
    return text.replace("&", "&amp;").replace("<", "&lt;").replace(">", "&gt;") # Use &amp; for safety

def _infer_lang_from_filename(filename):
    if not filename: return "plaintext"
    if '.' in filename:
        ext = filename.split('.')[-1].lower()
        mapping = {
            'py': 'python', 'js': 'javascript', 'ts': 'typescript', 'jsx': 'javascript', 'tsx': 'typescript',
            'html': 'html', 'htm': 'html', 'css': 'css', 'scss': 'scss', 'sass': 'sass', 'less': 'less',
            'json': 'json', 'xml': 'xml', 'yaml': 'yaml', 'yml': 'yaml', 'toml': 'toml',
            'md': 'markdown', 'rst': 'rst',
            'sh': 'bash', 'bash': 'bash', 'zsh': 'bash', 'bat': 'batch', 'cmd': 'batch', 'ps1': 'powershell',
            'c': 'c', 'h': 'c', 'cpp': 'cpp', 'hpp': 'cpp', 'cs': 'csharp', 'java': 'java',
            'rb': 'ruby', 'php': 'php', 'go': 'go', 'rs': 'rust', 'swift': 'swift', 'kt': 'kotlin', 'kts': 'kotlin',
            'sql': 'sql', 'dockerfile': 'docker', 'tf': 'terraform', 'hcl': 'terraform',
            'txt': 'plaintext', 'log': 'plaintext', 'ini': 'ini', 'conf': 'plaintext', 'cfg': 'plaintext',
            'csv': 'plaintext', 'tsv': 'plaintext', 'err': 'plaintext',
            '.env': 'plaintext', '.gitignore': 'plaintext', '.npmrc': 'plaintext', '.gitattributes': 'plaintext',
            'makefile': 'makefile',
        }
        return mapping.get(ext, "plaintext")
    base_filename = os.path.basename(filename)
    if base_filename == 'Dockerfile': return 'docker'
    if base_filename == 'Makefile': return 'makefile'
    if base_filename.startswith('.'): return 'plaintext'
    return "plaintext"

def _clean_filename(filename_line_content):
    text = filename_line_content.strip()
    # Remove markdown formatting characters aggressively
    text = re.sub(r'[`\*_#]+', '', text).strip()
    # Remove parenthesized descriptions
    text = re.split(r'\s*\(', text, 1)[0].strip()
    # Remove leading/trailing quotes or colons sometimes added by models
    text = text.strip('\'":;,')
    # Ensure it doesn't start with '/'
    text = text.lstrip('/')
    return text

def _parse_and_update_state_cache(latest_bot_message_content, current_files_state):
    """
    Parses the latest bot message content for file blocks and updates the
    global state cache with the latest version of each file.
    Returns the updated state cache and a list of filenames proposed in this turn.
    """
    # Start with a dictionary representation of the current state for easy updates
    current_files_dict = {f["filename"]: f.copy() for f in current_files_state if not f.get("is_structure_block")}
    structure_block_state = next((b for b in current_files_state if b.get("is_structure_block")), None)

    content = latest_bot_message_content or ""

    file_pattern = re.compile(r"### File:\s*(?P<filename_line>[^\n]+)\n(?:```(?P<lang>[\w\.\-\+]*)\n(?P<code>[\s\S]*?)\n```|(?P<binary_msg>\[Binary file(?: - [^\]]+)?\]))", re.MULTILINE)
    structure_pattern = re.compile(r"## File Structure\n```(?:(?P<struct_lang>[\w.-]*)\n)?(?P<structure_code>[\s\S]*?)\n```", re.MULTILINE)

    # Parse File Structure block if present in the latest message (overwrites previous structure block)
    structure_match = structure_pattern.search(content)
    if structure_match:
        structure_block_state = {"filename": "File Structure (from AI)", "language": structure_match.group("struct_lang") or "plaintext", "code": structure_match.group("structure_code").strip(), "is_binary": False, "is_structure_block": True}

    current_message_proposed_filenames = []
    # Parse file blocks from the latest message
    for match in file_pattern.finditer(content):
        filename = _clean_filename(match.group("filename_line"))
        if not filename:
            print(f"Warning: Skipped file block due to empty/invalid filename parsing: '{match.group('filename_line').strip()}'")
            continue # Skip if filename couldn't be parsed

        lang, code_block, binary_msg = match.group("lang"), match.group("code"), match.group("binary_msg")

        item_data = {"filename": filename, "is_binary": False, "is_structure_block": False}

        if code_block is not None:
            item_data["code"] = code_block.strip()
            item_data["language"] = (lang.strip().lower() if lang else _infer_lang_from_filename(filename))
            item_data["is_binary"] = False # Ensure explicit False if it's a code block
        elif binary_msg is not None:
            item_data["code"] = binary_msg.strip()
            item_data["language"] = "binary"
            item_data["is_binary"] = True
        else:
             # This case shouldn't be hit with the current regex, but as a safeguard
             print(f"Warning: Skipped file block for '{filename}' due to missing code or binary marker.")
             continue # Skip if content is neither code nor binary marker

        # Update or add the file in the dictionary state
        current_files_dict[filename] = item_data
        current_message_proposed_filenames.append(filename)


    # Convert dictionary back to a list, add structure block if present
    updated_parsed_blocks = list(current_files_dict.values())
    if structure_block_state:
        updated_parsed_blocks.insert(0, structure_block_state) # Add structure block at the beginning

    # Sort for consistent ordering
    updated_parsed_blocks.sort(key=lambda b: (0, b["filename"]) if b.get("is_structure_block") else (1, b["filename"]))

    return updated_parsed_blocks, current_message_proposed_filenames


def _export_selected_logic(selected_filenames, space_line_name_for_md, parsed_blocks_for_export):
    """Generates the Markdown representation of the space state or selected files."""
    results = {"output_str": "", "error_message": None, "download_filepath": None}
    # Only include blocks that are files (not structure) for content export/display
    file_blocks_for_export = [b for b in parsed_blocks_for_export if not b.get("is_structure_block")]

    # Determine filenames present in the state that can potentially be exported/listed
    all_filenames_in_state = sorted(list(set(b["filename"] for b in file_blocks_for_export)))

    if not all_filenames_in_state:
        results["output_str"] = f"# Space: {space_line_name_for_md}\n## File Structure\n{bbb}\n📁 Root\n{bbb}\n\n*No files in state to list structure or export.*"
        # Even if no files, create a temp file for download button functionality
        try:
            with tempfile.NamedTemporaryFile(mode="w", delete=False, suffix=".md", encoding='utf-8') as tmpfile:
                tmpfile.write(results["output_str"]); results["download_filepath"] = tmpfile.name
        except Exception as e: print(f"Error creating temp file for empty state export: {e}")
        return results

    output_lines = [f"# Space: {space_line_name_for_md}"]

    # Add File Structure block if it exists in the state
    structure_block = next((b for b in parsed_blocks_for_export if b.get("is_structure_block")), None)
    if structure_block:
        output_lines.extend(["## File Structure", bbb, structure_block["code"].strip(), bbb, ""])
    else:
        # If no AI-generated structure block, create a basic one from file list
        output_lines.extend(["## File Structure", bbb, "📁 Root"])
        if all_filenames_in_state:
            for fname in all_filenames_in_state: output_lines.append(f"  📄 {fname}") # Basic flattening
        output_lines.extend([bbb, ""])

    output_lines.append("Below are the contents of all files in the space:\n")

    # Filter blocks to export content based on selection
    # If selected_filenames is None or empty, export all file blocks
    blocks_to_export_content = sorted([b for b in file_blocks_for_export if not selected_filenames or b["filename"] in selected_filenames], key=lambda b: b["filename"])

    exported_content_count = 0
    for block in blocks_to_export_content:
        output_lines.append(f"### File: {block['filename']}")
        content = block.get('code', '')
        if block.get('is_binary') or content.startswith(("[Binary file", "[Error loading content:", "[Binary or Skipped file]")):
             # For binary/error placeholders, just print the marker line
             output_lines.append(content)
        else:
            # For text content, wrap in code block
            lang = block.get('language', 'plaintext') or 'plaintext' # Ensure language is not None or empty
            output_lines.extend([f"{bbb}{lang}", content, bbb])
        output_lines.append("") # Add blank line after each file block definition
        exported_content_count += 1

    if not exported_content_count:
        if selected_filenames:
             output_lines.append("*No selected files have editable content in the state.*")
        # else: already handled by the initial check for all_filenames_in_state

    final_output_str = "\n".join(output_lines)
    results["output_str"] = final_output_str
    try:
        # Create a temporary file for the download button
        with tempfile.NamedTemporaryFile(mode="w", delete=False, suffix=".md", encoding='utf-8') as tmpfile:
            tmpfile.write(final_output_str)
            results["download_filepath"] = tmpfile.name
    except Exception as e:
        print(f"Error creating temp file for download: {e}")
        results["error_message"] = "Could not prepare file for download."
        results["download_filepath"] = None # Ensure download is disabled on error

    return results

def _convert_gr_history_to_api_messages(system_prompt, gr_history, current_user_message=None):
    messages = [{"role": "system", "content": system_prompt}] if system_prompt else []
    for user_msg, bot_msg in gr_history:
        if user_msg: messages.append({"role": "user", "content": user_msg})
        if bot_msg and isinstance(bot_msg, str): messages.append({"role": BOT_ROLE_NAME, "content": bot_msg})
    # Append the current user message last if provided
    if current_user_message: messages.append({"role": "user", "content": current_user_message})
    return messages

def _generate_ui_outputs_from_cache(owner, space_name):
    """Generates the Markdown displays and download link from the global state cache."""
    global parsed_code_blocks_state_cache
    # Markdown preview displays the *latest* version of each file from the cache
    preview_md_val = "*No files in cache to display.*"
    # Formatted markdown is the full representation including structure and all files
    space_line_name = f"{owner}/{space_name}" if owner and space_name else (owner or space_name or "your-space")
    export_result = _export_selected_logic(None, space_line_name, parsed_code_blocks_state_cache)
    formatted_md_val = export_result["output_str"]
    download_file = export_result["download_filepath"]
    formatted_md_val = formatted_md_val or "*Load or define a Space to see its Markdown structure.*" # Fallback text

    if parsed_code_blocks_state_cache:
        preview_md_lines = ["## Detected/Updated Files & Content (Latest Versions):"]
        file_blocks_in_cache = sorted([b for b in parsed_code_blocks_state_cache if not b.get("is_structure_block")], key=lambda b: b["filename"])
        structure_block_in_cache = next((b for b in parsed_code_blocks_state_cache if b.get("is_structure_block")), None)

        if structure_block_in_cache:
             preview_md_lines.append(f"\n----\n**File Structure:** (from AI)\n{bbb}\n{escape_html_for_markdown(structure_block_in_cache['code'].strip())}\n{bbb}\n")

        if file_blocks_in_cache:
            for block in file_blocks_in_cache:
                preview_md_lines.append(f"\n----\n**File:** `{escape_html_for_markdown(block['filename'])}`")
                if block.get('is_binary'): preview_md_lines.append(f" (Binary File or Skipped)\n")
                else: preview_md_lines.append(f" (Language: `{block['language']}`)\n")

                content = block.get('code', '')
                if block.get('is_binary') or content.startswith(("[Binary file", "[Error loading content:", "[Binary or Skipped file]")):
                     preview_md_lines.append(f"\n`{escape_html_for_markdown(content.strip())}`\n") # Strip for preview markdown
                else:
                     # Use 3 backticks for code block in preview
                     lang = block.get('language', 'plaintext') or 'plaintext'
                     preview_md_lines.append(f"\n{bbb}{lang}\n{content.strip()}\n{bbb}\n") # Strip for preview markdown
        preview_md_val = "\n".join(preview_md_lines)


    return formatted_md_val, preview_md_val, gr.update(value=download_file, interactive=download_file is not None)

# --- NEW: Core logic for Change Staging and Confirmation ---

def generate_and_stage_changes(ai_response_content, current_files_state, hf_owner_name, hf_repo_name):
    """
    Parses AI response, compares with current state (from cache),
    and generates a structured changeset and a markdown summary.
    """
    changeset = []
    current_files_dict = {f["filename"]: f for f in current_files_state if not f.get("is_structure_block")}

    # 1. Parse AI response for actions and file blocks
    # Use build_logic_parse_markdown to get file blocks in the AI's desired format
    ai_parsed_md = build_logic_parse_markdown(ai_response_content)
    ai_proposed_files_list = ai_parsed_md.get("files", []) # List of {"path": ..., "content": ...}

    # Convert AI proposed files list to dict for easier lookup and comparison
    ai_proposed_files_dict = {f["path"]: f for f in ai_proposed_files_list}


    # Parse HF_ACTION commands from AI response using regex on the raw content
    action_pattern = re.compile(r"### HF_ACTION:\s*(?P<command_line>[^\n]+)", re.MULTILINE)
    for match in action_pattern.finditer(ai_response_content):
        cmd_parts = shlex.split(match.group("command_line").strip())
        if not cmd_parts: continue
        command, args = cmd_parts[0].upper(), cmd_parts[1:]

        # Add actions to the changeset
        if command == "CREATE_SPACE" and args:
            # The AI command specifies the target repo_id
            repo_id = args[0]
            sdk = "gradio"  # default
            private = False # default
            if '--sdk' in args:
                try: sdk = args[args.index('--sdk') + 1]
                except IndexError: print("Warning: CREATE_SPACE --sdk requires an argument.")
            if '--private' in args:
                 try: private_str = args[args.index('--private') + 1].lower()
                 except IndexError: print("Warning: CREATE_SPACE --private requires an argument.")
                 else: private = private_str == 'true'
            # Action includes target repo, sdk, and private setting
            changeset.append({"type": "CREATE_SPACE", "repo_id": repo_id, "sdk": sdk, "private": private})
            print(f"Staged CREATE_SPACE action for {repo_id}")

        elif command == "DELETE_FILE" and args:
            file_path = args[0]
            changeset.append({"type": "DELETE_FILE", "path": file_path})
            print(f"Staged DELETE_FILE action for {file_path}")
            # Note: AI might propose deleting a file it *just* created/updated in the same turn.
            # The application logic handle_confirm_changes should process deletes *before* adds/updates,
            # or the commit operation itself should handle the conflict gracefully (delete then re-add/update).
            # The current `apply_staged_changes` does deletes first, then uploads.

        elif command == "SET_PRIVATE" and args:
            private = args[0].lower() == 'true'
            # Action applies to the currently loaded space
            changeset.append({"type": "SET_PRIVACY", "private": private, "repo_id": f"{hf_owner_name}/{hf_repo_name}"})
            print(f"Staged SET_PRIVACY action for {hf_owner_name}/{hf_repo_name} to {private}")

        elif command == "DELETE_SPACE":
            # Action applies to the currently loaded space
            changeset.append({"type": "DELETE_SPACE", "owner": hf_owner_name, "space_name": hf_repo_name})
            print(f"Staged DELETE_SPACE action for {hf_owner_name}/{hf_repo_name}")
        # Add other actions here as needed (e.g., `SET_HARDWARE`, `RESTART_SPACE`)

    # 3. Compare proposed files from AI with current files to determine CREATE/UPDATE
    # Iterate through files proposed by the AI in this turn
    for file_info in ai_proposed_files_list:
        filename = file_info["path"]
        proposed_content = file_info["content"]

        if filename in current_files_dict:
            # File exists, check if content changed
            current_content = current_files_dict[filename]["code"]
            if proposed_content != current_content:
                 # Check if current file state is a binary/error placeholder before marking as update
                 # If current is placeholder and proposed is content, treat as update (content is now known)
                 # If both are placeholders or proposed is placeholder, maybe skip or special flag?
                 is_current_placeholder = current_content.startswith(("[Binary file", "[Error loading content:", "[Binary or Skipped file]"))
                 is_proposed_placeholder = proposed_content.startswith(("[Binary file", "[Error loading content:", "[Binary or Skipped file]"))

                 if not is_proposed_placeholder: # Only stage update if AI provides actual content
                    # Determine language for potential new/updated file block representation
                    # Use the language if the file was already in cache, otherwise infer from filename
                    lang = current_files_dict[filename].get("language") or _infer_lang_from_filename(filename)
                    changeset.append({"type": "UPDATE_FILE", "path": filename, "content": proposed_content, "lang": lang})
                    print(f"Staged UPDATE_FILE action for {filename}")
                 elif is_current_placeholder and is_proposed_placeholder:
                     print(f"Skipping staging update for {filename}: both current and proposed content are placeholders.")
                 elif not is_current_placeholder and is_proposed_placeholder:
                      print(f"Warning: AI proposed placeholder content for existing file {filename}. Staging ignored.")


        else:
            # File does not exist, stage as CREATE
            # Only stage creation if AI provides actual content, not just a placeholder
            proposed_content = file_info["content"]
            if not (proposed_content.startswith("[Binary file") or proposed_content.startswith("[Error loading content:") or proposed_content.startswith("[Binary or Skipped file]")):
                lang = _infer_lang_from_filename(filename) # Infer language for a new file
                changeset.append({"type": "CREATE_FILE", "path": filename, "content": proposed_content, "lang": lang})
                print(f"Staged CREATE_FILE action for {filename}")
            else:
                 print(f"Skipping staging create for {filename}: Proposed content is a placeholder.")


    # 4. Format the changeset into a human-readable Markdown string
    if not changeset:
        md_summary = ["### 📋 Proposed Changes Plan", "\nThe AI did not propose any specific changes to files or the space.\n"]
    else:
        md_summary = ["### 📋 Proposed Changes Plan\n"]
        md_summary.append("The AI has proposed the following changes. Please review and confirm.")

        # Separate action types for clearer display
        file_changes = [c for c in changeset if c['type'] in ['CREATE_FILE', 'UPDATE_FILE', 'DELETE_FILE']]
        space_actions = [c for c in changeset if c['type'] not in ['CREATE_FILE', 'UPDATE_FILE', 'DELETE_FILE']]

        if space_actions:
             md_summary.append("\n**Space Actions:**")
             for change in space_actions:
                 if change["type"] == "CREATE_SPACE":
                     md_summary.append(f"- **🚀 Create New Space:** `{change.get('repo_id', '...')}` (SDK: {change.get('sdk', 'gradio')}, Private: {change.get('private', False)})")
                 elif change["type"] == "SET_PRIVACY":
                     md_summary.append(f"- **🔒 Set Privacy:** Set `{change.get('repo_id', '...')}` to `private={change.get('private', False)}`")
                 elif change["type"] == "DELETE_SPACE":
                     md_summary.append(f"- **💥 DELETE ENTIRE SPACE:** `{change.get('owner', '...')}/{change.get('space_name', '...')}` **(DESTRUCTIVE ACTION)**")
             md_summary.append("") # Add newline after actions

        if file_changes:
             md_summary.append("**File Changes:**")
             for change in file_changes:
                 if change["type"] == "CREATE_FILE":
                     md_summary.append(f"- **➕ Create File:** `{change['path']}`")
                 elif change["type"] == "UPDATE_FILE":
                     md_summary.append(f"- **🔄 Update File:** `{change['path']}`")
                 elif change["type"] == "DELETE_FILE":
                     md_summary.append(f"- **➖ Delete File:** `{change['path']}`")

    return changeset, "\n".join(md_summary)


# --- Gradio Event Handlers ---

def handle_chat_submit(user_message, chat_history, hf_api_key_input, provider_api_key_input, provider_select, model_select, system_prompt, hf_owner_name, hf_repo_name):
    global parsed_code_blocks_state_cache
    _chat_msg_in, _chat_hist = "", list(chat_history)

    # Hide confirmation UI while AI is thinking
    yield (
        _chat_msg_in, _chat_hist, "Initializing...",
        gr.update(), gr.update(), gr.update(interactive=False), gr.update(value="*No changes proposed.*"), # Clear summary
        [], gr.update(visible=False), gr.update(visible=False), gr.update(visible=False) # Hide confirm UI
    )

    if not user_message.strip():
        yield (
            _chat_msg_in, _chat_hist, "Cannot send an empty message.",
            gr.update(), gr.update(), gr.update(), gr.update(),
            [], gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
        )
        return

    _chat_hist.append((user_message, None))
    yield (
        _chat_msg_in, _chat_hist, f"Sending to {model_select}...",
        gr.update(), gr.update(), gr.update(), gr.update(),
        [], gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
    )

    # Prepare context for the AI - Export current state to Markdown
    current_sys_prompt = system_prompt.strip() or DEFAULT_SYSTEM_PROMPT
    space_id_for_context = f"{hf_owner_name}/{hf_repo_name}" if hf_owner_name and hf_repo_name else "your-space"
    export_result = _export_selected_logic(None, space_id_for_context, parsed_code_blocks_state_cache)
    current_files_context_md = export_result['output_str'] or "*No files currently loaded or defined.*"
    current_files_context = f"\n\n## Current Space Context: {space_id_for_context}\n{current_files_context_md}"

    user_message_with_context = user_message.strip() + current_files_context
    api_msgs = _convert_gr_history_to_api_messages(current_sys_prompt, _chat_hist[:-1], user_message_with_context)

    try:
        full_bot_response_content = ""
        # Pass the provider API key from the UI to the generation logic
        streamer = generate_stream(provider_select, model_select, provider_api_key_input, api_msgs)
        for chunk in streamer:
            if chunk is None: continue
            if isinstance(chunk, str):
                 # Check for error indicators early
                 if full_bot_response_content == "" and (chunk.startswith("Error:") or chunk.startswith("API HTTP Error")):
                     full_bot_response_content = chunk; break
                 full_bot_response_content += str(chunk)

            _chat_hist[-1] = (user_message, full_bot_response_content)
            yield (
                _chat_msg_in, _chat_hist, f"Streaming from {model_select}...",
                gr.update(), gr.update(), gr.update(), gr.update(),
                [], gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
            )

        # Handle potential errors from the streamer
        if full_bot_response_content.startswith("Error:") or full_bot_response_content.startswith("API HTTP Error"):
            _status = full_bot_response_content
            yield (_chat_msg_in, _chat_hist, _status, gr.update(), gr.update(), gr.update(), gr.update(), [], gr.update(visible=False), gr.update(visible=False), gr.update(visible=False))
            return

        # --- Post-streaming: Parse AI output, update cache, stage changes ---

        _status = "Stream complete. Parsing response and staging changes..."
        yield (_chat_msg_in, _chat_hist, _status, gr.update(), gr.update(), gr.update(), gr.update(), [], gr.update(visible=False), gr.update(visible=False), gr.update(visible=False))

        # 1. Update the state cache based on the *full* AI response
        # This reflects the AI's understanding and proposed file content *in the UI*.
        parsed_code_blocks_state_cache, proposed_filenames_in_turn = _parse_and_update_state_cache(full_bot_response_content, parsed_code_blocks_state_cache)

        # Regenerate UI previews based on the updated cache
        _formatted, _detected, _download = _generate_ui_outputs_from_cache(hf_owner_name, hf_repo_name)

        # 2. Generate the changeset and summary based on the AI output and current cache state
        # Pass the updated cache to generate_and_stage_changes for comparison
        staged_changeset, summary_md = generate_and_stage_changes(full_bot_response_content, parsed_code_blocks_state_cache, hf_owner_name, hf_repo_name)


        if not staged_changeset:
            _status = summary_md # Will be "No changes proposed" message
            yield (
                _chat_msg_in, _chat_hist, _status,
                _detected, _formatted, _download,
                [], # Clear changeset state
                gr.update(value=summary_md), # Display summary
                gr.update(visible=False), gr.update(visible=False), gr.update(visible=False) # Hide confirm UI
            )
        else:
            _status = "Change plan generated. Please review and confirm below."
            yield (
                _chat_msg_in, _chat_hist, _status,
                _detected, _formatted, _download,
                staged_changeset,  # Send changeset to state
                gr.update(value=summary_md), # Display summary
                gr.update(visible=True), # Show the accordion
                gr.update(visible=True), # Show confirm button
                gr.update(visible=True)  # Show cancel button
            )

    except Exception as e:
        error_msg = f"An unexpected error occurred: {e}"
        print(f"Error in handle_chat_submit: {e}")
        import traceback
        traceback.print_exc()
        if _chat_hist:
            # Ensure the last message is not None before updating
            if _chat_hist[-1] and _chat_hist[-1][0] == user_message:
                 _chat_hist[-1] = (user_message, (full_bot_response_content + "\n\n" if full_bot_response_content and full_bot_response_content != user_message else "") + error_msg)
            else: # Should not happen, but as fallback
                 _chat_hist.append((user_message, error_msg))


        # Regenerate UI previews based on the updated cache (even if there was an error after streaming)
        _formatted, _detected, _download = _generate_ui_outputs_from_cache(hf_owner_name, hf_repo_name)

        yield (
            _chat_msg_in, _chat_hist, error_msg,
            _detected, _formatted, _download,
            [], # Clear changeset state on error
            gr.update(value="*Error occurred, changes plan cleared.*"), # Clear summary display
            gr.update(visible=False), gr.update(visible=False), gr.update(visible=False) # Hide confirm UI
        )


def handle_confirm_changes(hf_api_key, owner_name, space_name, changeset):
    """Applies the staged changes from the changeset."""
    global parsed_code_blocks_state_cache

    # Hide the confirmation UI immediately
    yield "Applying changes...", gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
    # Keep the summary visible potentially with a loading indicator? Or clear? Let's clear it.
    yield "Applying changes...", gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(value="*Applying changes...*")

    if not changeset:
        # This shouldn't happen if the button is hidden, but as a safeguard
        return "No changes to apply.", gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(value="No changes were staged.")

    # Call the build_logic function to apply the changes
    # The build_logic function will return a status message string
    status_message = apply_staged_changes(hf_api_key, owner_name, space_name, changeset)

    # After applying changes, reload the space state to reflect the actual state on the Hub
    # This is important because the build_logic might fail partially, or the AI's cache might be outdated.
    # Reloading ensures the UI reflects the reality on the Hub.
    _status_reload = f"{status_message} | Reloading Space state..."
    yield _status_reload, gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(value="*Reloading Space state...*")

    # Need to call handle_load_existing_space or similar logic here
    # Let's replicate the core logic from handle_load_existing_space needed to refresh the cache and UI
    # Note: This doesn't update the chat history or other parts of handle_load_existing_space,
    # just the file-related UI elements.
    refreshed_file_list = []
    reload_error = None
    repo_id_for_reload = f"{owner_name}/{space_name}" if owner_name and space_name else None

    if repo_id_for_reload:
        sdk, file_list, err_list = get_space_repository_info(hf_api_key, space_name, owner_name)
        if err_list:
            reload_error = f"Error reloading file list after changes: {err_list}"
            parsed_code_blocks_state_cache = [] # Clear cache if list fails
        else:
            refreshed_file_list = file_list
            loaded_files = []
            for file_path in refreshed_file_list:
                content, err_get = get_space_file_content(hf_api_key, space_name, owner_name, file_path)
                lang = _infer_lang_from_filename(file_path)
                is_binary = lang == "binary" or err_get # Assume error indicates binary/unreadable
                code = f"[Error loading content: {err_get}]" if err_get else content
                loaded_files.append({"filename": file_path, "code": code, "language": lang, "is_binary": is_binary, "is_structure_block": False})
            parsed_code_blocks_state_cache = loaded_files # Update cache with refreshed state
            # Add back the structure block if it was in the cache before reload (it's AI generated, not from Hub)
            # This might be wrong - the cache should represent the *actual* state + AI's last proposed file changes.
            # Let's keep AI's proposed structure block in the cache until a new one replaces it.
            # But reloading from Hub should overwrite the *file content*. The structure block is separate.
            # If we reload from Hub, the cache should be *only* Hub files + the last AI structure block.
            last_ai_structure_block = next((b for b in parsed_code_blocks_state_cache if b.get("is_structure_block")), None) # Check cache *before* clearing
            if last_ai_structure_block:
                 # Find it in the *new* loaded_files list if it exists there (e.g. README.md could be structure?)
                 # Or just re-add the AI structure block if it was in the cache previously?
                 # Let's stick to the simpler model: AI structure block is just for display/context in the markdown tab,
                 # the actual files are what's loaded/applied. Reloading files replaces the file blocks in cache, structure block is kept or removed based on AI output.
                 # On reload, we only get actual files from the Hub. The AI structure block is NOT on the hub.
                 # So, clearing and adding only Hub files is correct.
                 pass # last_ai_structure_block is not added back.

    else:
        reload_error = "Cannot reload Space state: Owner or Space Name missing."


    # Regenerate UI previews based on the refreshed cache state
    _formatted, _detected, _download = _generate_ui_outputs_from_cache(owner_name, space_name)

    final_overall_status = status_message + (f" | Reload Status: {reload_error}" if reload_error else " | Reload Status: Space state refreshed.")

    # Clear the changeset state after application attempt
    cleared_changeset = []

    # Return updated UI elements and hide confirmation UI
    return (
        final_overall_status,
        _formatted,
        _detected,
        _download,
        gr.update(visible=False), # Hide accordion
        gr.update(visible=False), # Hide confirm button
        gr.update(visible=False), # Hide cancel button
        cleared_changeset, # Clear changeset state
        gr.update(value="*No changes proposed.*") # Clear summary display
    )


def handle_cancel_changes():
    """Clears the staged changeset and hides the confirmation UI."""
    global parsed_code_blocks_state_cache # Cancel doesn't affect the cache state
    return (
        "Changes cancelled.",
        [], # Clear changeset state
        gr.update(value="*No changes proposed.*"), # Clear summary display
        gr.update(visible=False), # Hide accordion
        gr.update(visible=False), # Hide confirm button
        gr.update(visible=False) # Hide cancel button
    )


def update_models_dropdown(provider_select):
    if not provider_select: return gr.update(choices=[], value=None)
    models = get_models_for_provider(provider_select)
    default_model = get_default_model_for_provider(provider_select)
    selected_value = default_model if default_model in models else (models[0] if models else None)
    return gr.update(choices=models, value=selected_value)

def handle_load_existing_space(hf_api_key_ui, ui_owner_name, ui_space_name):
    global parsed_code_blocks_state_cache
    _formatted_md_val, _detected_preview_val, _status_val = "*Loading files...*", "*Loading files...*", f"Loading Space: {ui_owner_name}/{ui_space_name}..."
    _file_browser_update, _iframe_html_update, _download_btn_update = gr.update(visible=False, choices=[], value=None), gr.update(value=None, visible=False), gr.update(interactive=False, value=None)
    _build_status_clear, _edit_status_clear, _runtime_status_clear = "*Build status...*", "*Select a file...*", "*Runtime status...*"
    _chat_history_clear = [] # Don't clear chat history on load
    _changeset_clear = [] # Clear staged changes on load
    _changeset_summary_clear = "*No changes proposed.*" # Clear summary on load
    _confirm_ui_hidden = gr.update(visible=False) # Hide confirm UI on load


    # Initial yield to show loading state
    outputs = [
        _formatted_md_val, _detected_preview_val, _status_val, _file_browser_update,
        gr.update(value=ui_owner_name), gr.update(value=ui_space_name), # Update owner/space fields immediately
        _iframe_html_update, _download_btn_update, _build_status_clear,
        _edit_status_clear, _runtime_status_clear, _chat_history_clear,
        _changeset_clear, _changeset_summary_clear, _confirm_ui_hidden, _confirm_ui_hidden, _confirm_ui_hidden # Hide confirmation UI
    ]
    yield outputs

    owner_to_use = ui_owner_name
    if not owner_to_use:
        token, err = build_logic_get_api_token(hf_api_key_ui)
        if err or not token:
            _status_val = f"Error: {err or 'Cannot determine owner from token.'}"
            outputs[2] = _status_val; yield outputs; return
        try:
            user_info = build_logic_whoami(token=token)
            owner_to_use = user_info.get('name')
            if not owner_to_use: raise Exception("Could not find user name from token.")
            outputs[4] = gr.update(value=owner_to_use) # Update UI owner field
            _status_val += f" (Auto-detected owner: {owner_to_use})"
        except Exception as e:
            _status_val = f"Error auto-detecting owner: {e}"; outputs[2] = _status_val; yield outputs; return

    if not owner_to_use or not ui_space_name:
        _status_val = "Error: Owner and Space Name are required."; outputs[2] = _status_val; yield outputs; return

    sdk, file_list, err = get_space_repository_info(hf_api_key_ui, ui_space_name, owner_to_use)

    # Always update owner/space inputs even on error, as user entered them
    outputs[4] = gr.update(value=owner_to_use)
    outputs[5] = gr.update(value=ui_space_name)

    if err:
        _status_val = f"Load Error: {err}"
        parsed_code_blocks_state_cache = [] # Clear cache on load error
        _formatted, _detected, _download = _generate_ui_outputs_from_cache(owner_to_use, ui_space_name)
        outputs[0], outputs[1], outputs[2], outputs[7] = _formatted, _detected, _status_val, _download # Update markdown, preview, status, download
        outputs[3] = gr.update(visible=False, choices=[], value=None) # Hide file browser
        outputs[6] = gr.update(value=None, visible=False) # Hide iframe
        yield outputs; return

    # Success case: Populate cache and UI
    loaded_files = []
    for file_path in file_list:
        content, err_get = get_space_file_content(hf_api_key_ui, ui_space_name, owner_to_use, file_path)
        lang = _infer_lang_from_filename(file_path)
        is_binary = lang == "binary" or (err_get is not None) # err_get will be a string if error
        code = f"[Error loading content: {err_get}]" if err_get else (content or "") # Ensure code is empty string if content is None
        loaded_files.append({"filename": file_path, "code": code, "language": lang, "is_binary": is_binary, "is_structure_block": False})

    # When loading, the cache should only contain the actual files from the Hub.
    # Any previous AI-generated structure block is discarded.
    parsed_code_blocks_state_cache = loaded_files

    _formatted, _detected, _download = _generate_ui_outputs_from_cache(owner_to_use, ui_space_name)
    _status_val = f"Successfully loaded {len(file_list)} files from {owner_to_use}/{ui_space_name}. SDK: {sdk or 'unknown'}."
    outputs[0], outputs[1], outputs[2], outputs[7] = _formatted, _detected, _status_val, _download # Update markdown, preview, status, download

    # Update file browser dropdown
    outputs[3] = gr.update(visible=True, choices=sorted(file_list or []), value=None)

    # Update iframe preview
    if owner_to_use and ui_space_name:
        sub_owner = re.sub(r'[^a-z0-9\-]+', '-', owner_to_use.lower()).strip('-') or 'owner'
        sub_repo = re.sub(r'[^a-z0-9\-]+', '-', ui_space_name.lower()).strip('-') or 'space'
        iframe_url = f"https://{sub_owner}-{sub_repo}{'.static.hf.space' if sdk == 'static' else '.hf.space'}"
        outputs[6] = gr.update(value=f'<iframe src="{iframe_url}?__theme=light&embed=true" width="100%" height="500px"></iframe>', visible=True)
    else:
         outputs[6] = gr.update(value=None, visible=False)


    yield outputs

# This manual build button now uses the formatted_space_output_display content
# It's separate from the AI-driven apply_staged_changes
def handle_build_space_button(hf_api_key_ui, ui_space_name_part, ui_owner_name_part, space_sdk_ui, is_private_ui, formatted_markdown_content):
    _build_status, _iframe_html, _file_browser_update = "Starting manual space build process...", gr.update(value=None, visible=False), gr.update(visible=False, choices=[], value=None)
    # Also hide confirmation UI and clear state on manual build
    _changeset_clear = []
    _changeset_summary_clear = "*Manual build initiated, changes plan cleared.*"
    _confirm_ui_hidden = gr.update(visible=False)

    yield (_build_status, _iframe_html, _file_browser_update, gr.update(value=ui_owner_name_part), gr.update(value=ui_space_name_part),
           _changeset_clear, _changeset_summary_clear, _confirm_ui_hidden, _confirm_ui_hidden, _confirm_ui_hidden)

    if not ui_space_name_part or "/" in ui_space_name_part:
        _build_status = f"Build Error: Invalid Space Name '{ui_space_name_part}'."
        yield (_build_status, _iframe_html, _file_browser_update, gr.update(), gr.update(),
               gr.update(), gr.update(), gr.update(), gr.update(), gr.update()); return

    # Use build_logic_create_space directly for manual build
    result_message = build_logic_create_space(ui_api_token_from_textbox=hf_api_key_ui, space_name_ui=ui_space_name_part, owner_ui=ui_owner_name_part, sdk_ui=space_sdk_ui, markdown_input=formatted_markdown_content, private=is_private_ui)
    _build_status = f"Manual Build Process: {result_message}"

    if "Successfully" in result_message:
        owner_to_use = ui_owner_name_part # Assume determined correctly by build_logic_create_space
        space_to_use = ui_space_name_part

        # Attempt to update UI with the new state after successful build
        # This is similar to handle_load_existing_space post-success
        sdk_built, file_list, err_list = get_space_repository_info(hf_api_key_ui, space_to_use, owner_to_use)
        if err_list:
            _build_status += f" | Error reloading file list after build: {err_list}"
            global parsed_code_blocks_state_cache
            parsed_code_blocks_state_cache = [] # Clear cache
            _file_browser_update = gr.update(visible=False, choices=[], value=None)
            _iframe_html = gr.update(value=None, visible=False)
        else:
            global parsed_code_blocks_state_cache
            loaded_files = []
            for file_path in file_list:
                 content, err_get = get_space_file_content(hf_api_key_ui, space_to_use, owner_to_use, file_path)
                 lang = _infer_lang_from_filename(file_path)
                 is_binary = lang == "binary" or (err_get is not None)
                 code = f"[Error loading content: {err_get}]" if err_get else (content or "")
                 loaded_files.append({"filename": file_path, "code": code, "language": lang, "is_binary": is_binary, "is_structure_block": False})
            parsed_code_blocks_state_cache = loaded_files # Update cache

            _file_browser_update = gr.update(visible=True, choices=sorted(file_list or []), value=None)

            sub_owner = re.sub(r'[^a-z0-9\-]+', '-', owner_to_use.lower()).strip('-') or 'owner'
            sub_repo = re.sub(r'[^a-z0-9\-]+', '-', space_to_use.lower()).strip('-') or 'space'
            iframe_url = f"https://{sub_owner}-{sub_repo}{'.static.hf.space' if sdk_built == 'static' else '.hf.space'}"
            _iframe_html = gr.update(value=f'<iframe src="{iframe_url}?__theme=light&embed=true" width="100%" height="700px"></iframe>', visible=True)


    # Need to update formatted/detected markdown displays after manual build as well
    _formatted_md, _detected_preview, _download = _generate_ui_outputs_from_cache(ui_owner_name_part, ui_space_name_part)


    yield (_build_status, _iframe_html, _file_browser_update, gr.update(value=ui_owner_name_part), gr.update(value=ui_space_name_part),
           _changeset_clear, _changeset_summary_clear, _confirm_ui_hidden, _confirm_ui_hidden, _confirm_ui_hidden,
           _formatted_md, _detected_preview, _download # Include these outputs
           )


def handle_load_file_for_editing(hf_api_key_ui, ui_space_name_part, ui_owner_name_part, selected_file_path):
    if not selected_file_path:
        yield gr.update(value=""), "Select a file.", gr.update(value=""), gr.update(language="plaintext")
        return

    content, err = get_space_file_content(hf_api_key_ui, ui_space_name_part, ui_owner_name_part, selected_file_path)
    if err:
        # If load fails, clear editor and show error
        yield "", f"Error loading '{selected_file_path}': {err}", "", gr.update(language="plaintext")
        return

    lang = _infer_lang_from_filename(selected_file_path)
    commit_msg = f"Update {selected_file_path}"
    yield content, f"Loaded `{selected_file_path}`", commit_msg, gr.update(language=lang)

def handle_commit_file_changes(hf_api_key_ui, ui_space_name_part, ui_owner_name_part, file_to_edit_path, edited_content, commit_message):
    if not file_to_edit_path:
         return "Error: No file selected for commit.", gr.update(), gr.update(), gr.update(), gr.update()

    status_msg = update_space_file(hf_api_key_ui, ui_space_name_part, ui_owner_name_part, file_to_edit_path, edited_content, commit_message)
    file_list, _ = list_space_files_for_browsing(hf_api_key_ui, ui_space_name_part, ui_owner_name_part) # Refresh file list dropdown
    global parsed_code_blocks_state_cache
    if "Successfully" in status_msg:
        # Update cache on success
        # Find the block in the cache and update its content and potentially language
        found = False
        for block in parsed_code_blocks_state_cache:
            if block["filename"] == file_to_edit_path and not block.get("is_structure_block"):
                block["code"] = edited_content
                block["language"] = _infer_lang_from_filename(file_to_edit_path) # Re-infer language in case extension changed
                block["is_binary"] = False # Assume edited content is text
                found = True
                break
        if not found:
             # If the file wasn't in cache (e.g., loaded after AI chat), add it
             parsed_code_blocks_state_cache.append({
                 "filename": file_to_edit_path,
                 "code": edited_content,
                 "language": _infer_lang_from_filename(file_to_edit_path),
                 "is_binary": False,
                 "is_structure_block": False
             })
             # Re-sort cache
             parsed_code_blocks_state_cache.sort(key=lambda b: (0, b["filename"]) if b.get("is_structure_block") else (1, b["filename"]))


    # Regenerate markdown displays from updated cache
    _formatted, _detected, _download = _generate_ui_outputs_from_cache(ui_owner_name_part, ui_space_name_part)
    return status_msg, gr.update(choices=sorted(file_list or [])), _formatted, _detected, _download

def handle_delete_file(hf_api_key_ui, ui_space_name_part, ui_owner_name_part, file_to_delete_path):
    if not file_to_delete_path:
        return "No file selected to delete.", gr.update(), "", "", "plaintext", gr.update(), gr.update(), gr.update()

    status_msg = build_logic_delete_space_file(hf_api_key_ui, ui_space_name_part, ui_owner_name_part, file_to_delete_path)
    file_list, _ = list_space_files_for_browsing(hf_api_key_ui, ui_space_name_part, ui_owner_name_part) # Refresh file list dropdown
    global parsed_code_blocks_state_cache
    if "Successfully" in status_msg:
        # Update cache: remove the deleted file
        parsed_code_blocks_state_cache = [b for b in parsed_code_blocks_state_cache if b["filename"] != file_to_delete_path]
        # Clear the editor if the deleted file was currently loaded
        file_content_editor_update = gr.update(value="")
        commit_message_update = gr.update(value="")
        editor_lang_update = gr.update(language="plaintext")
    else:
        # If deletion failed, keep the editor content and status as they were
        file_content_editor_update = gr.update()
        commit_message_update = gr.update()
        editor_lang_update = gr.update()


    # Regenerate markdown displays from updated cache
    _formatted, _detected, _download = _generate_ui_outputs_from_cache(ui_owner_name_part, ui_space_name_part)
    return (
        status_msg,
        gr.update(choices=sorted(file_list or []), value=None), # Clear selected value after delete
        file_content_editor_update,
        commit_message_update,
        editor_lang_update,
        _formatted,
        _detected,
        _download
    )

def handle_refresh_space_status(hf_api_key_ui, ui_owner_name, ui_space_name):
    if not ui_owner_name or not ui_space_name:
        return "*Owner and Space Name must be provided to get status.*"

    status_details, err = get_space_runtime_status(hf_api_key_ui, ui_space_name, ui_owner_name)
    if err: return f"**Error:** {err}"
    if not status_details: return "*Could not retrieve status details.*"

    md = f"### Status for {ui_owner_name}/{ui_space_name}\n"
    # Display key status details
    md += f"- **Stage:** `{status_details.get('stage', 'N/A')}`\n"
    md += f"- **Status:** `{status_details.get('status', 'N/A')}`\n" # More detailed status
    md += f"- **Hardware:** `{status_details.get('hardware', 'N/A')}`\n"
    requested_hw = status_details.get('requested_hardware')
    if requested_hw: md += f"- **Requested Hardware:** `{requested_hw}`\n"
    error_msg = status_details.get('error_message')
    if error_msg: md += f"- **Error:** `{error_msg}`\n"
    log_link = status_details.get('full_log_link')
    if log_link and log_link != "#": md += f"- [View Full Logs]({log_link})\n"

    return md


# --- UI Theming and CSS (Unchanged) ---
custom_theme = gr.themes.Base(primary_hue="teal", secondary_hue="purple", neutral_hue="zinc", text_size="sm", spacing_size="md", radius_size="sm", font=["System UI", "sans-serif"])
custom_css = """
body { background: linear-gradient(to bottom right, #2c3e50, #34495e); color: #ecf0f1; }
.gradio-container { background: transparent !important; }
.gr-box, .gr-panel, .gr-pill { background-color: rgba(44, 62, 80, 0.8) !important; border-color: rgba(189, 195, 199, 0.2) !important; }
.gr-textbox, .gr-dropdown, .gr-button, .gr-code, .gr-chat-message { border-color: rgba(189, 195, 199, 0.3) !important; background-color: rgba(52, 73, 94, 0.9) !important; color: #ecf0f1 !important; }
.gr-button.gr-button-primary { background-color: #1abc9c !important; color: white !important; border-color: #16a085 !important; }
.gr-button.gr-button-secondary { background-color: #9b59b6 !important; color: white !important; border-color: #8e44ad !important; }
.gr-button.gr-button-stop { background-color: #e74c3c !important; color: white !important; border-color: #c0392b !important; }
.gr-markdown { background-color: rgba(44, 62, 80, 0.7) !important; padding: 10px; border-radius: 5px; overflow-x: auto; } /* Added overflow-x */
.gr-markdown h1, .gr-markdown h2, .gr-markdown h3, .gr-markdown h4, .gr-markdown h5, .gr-markdown h6 { color: #ecf0f1 !important; border-bottom-color: rgba(189, 195, 199, 0.3) !important; }
.gr-markdown pre code { background-color: rgba(52, 73, 94, 0.95) !important; border-color: rgba(189, 195, 199, 0.3) !important; }
.gr-chatbot { background-color: rgba(44, 62, 80, 0.7) !important; border-color: rgba(189, 195, 199, 0.2) !important; }
.gr-chatbot .message { background-color: rgba(52, 73, 94, 0.9) !important; color: #ecf0f1 !important; border-color: rgba(189, 195, 199, 0.3) !important; }
.gr-chatbot .message.user { background-color: rgba(46, 204, 113, 0.9) !important; color: black !important; }
/* Custom styles for Proposed Changes Accordion */
.gradio-container .gr-accordion { border-color: rgba(189, 195, 199, 0.3) !important; }
.gradio-container .gr-accordion.closed { background-color: rgba(52, 73, 94, 0.9) !important; }
.gradio-container .gr-accordion.open { background-color: rgba(44, 62, 80, 0.8) !important; }

"""


# --- Gradio UI Definition ---
with gr.Blocks(theme=custom_theme, css=custom_css) as demo:
    # State to hold the parsed change plan from the AI
    changeset_state = gr.State([])
    # State to hold the *current* representation of the Space's files and structure.
    # This is populated on load and updated by AI outputs or manual edits/deletes.
    # It's shared globally across handlers that modify/read the Space state.
    # parsed_code_blocks_state_cache = gr.State([]) # Global variable is simpler for now

    gr.Markdown("# 🤖 AI-Powered Hugging Face Space Commander")
    gr.Markdown("Use an AI assistant to create, modify, build, and manage your Hugging Face Spaces directly from this interface.")
    gr.Markdown("## ❗ This will cause changes to your huggingface spaces if you give it your Huggingface Key")
    gr.Markdown("⚒ Under Development - Use with Caution")


    with gr.Sidebar():
        with gr.Column(scale=1):
            with gr.Accordion("⚙️ Configuration", open=True):
                hf_api_key_input = gr.Textbox(label="Hugging Face Token", type="password", placeholder="hf_... (uses env var HF_TOKEN if empty)")
                owner_name_input = gr.Textbox(label="HF Owner Name", placeholder="e.g., your-username")
                space_name_input = gr.Textbox(label="HF Space Name", value="") # Default to empty
                load_space_button = gr.Button("🔄 Load Existing Space", variant="secondary")
                gr.Markdown("---") # Separator
                # Manual Space Actions (outside AI flow, but use same backend)
                # set_privacy_button = gr.Button("🔒 Toggle Space Privacy", variant="secondary") # Could be checkbox + button
                # delete_space_button = gr.Button("💥 Delete Entire Space", variant="stop") # Needs confirmation modal


            with gr.Accordion("🤖 AI Model Settings", open=True):
                # --- MODIFIED: Set up default provider and model logic on load ---
                available_providers = get_available_providers()
                default_provider = 'Groq'
                # Fallback if 'Groq' is not an option, or if list is smaller than 3
                if default_provider not in available_providers:
                    default_provider = available_providers[0] if available_providers else None
                elif len(available_providers) < 3:
                     default_provider = available_providers[0] if available_providers else None


                # Get initial models and the default model for the selected provider
                initial_models = get_models_for_provider(default_provider) if default_provider else []
                initial_model = get_default_model_for_provider(default_provider) if default_provider else None
                # Fallback for the model as well
                if initial_model not in initial_models:
                    initial_model = initial_models[0] if initial_models else None

                provider_select = gr.Dropdown(
                    label="AI Provider",
                    choices=available_providers,
                    value=default_provider,
                    allow_custom_value=False
                )
                model_select = gr.Dropdown(
                    label="AI Model",
                    choices=initial_models,
                    value=initial_model,
                    allow_custom_value=False
                )
                # --- END MODIFICATION ---
                provider_api_key_input = gr.Textbox(label="Model Provider API Key (Optional)", type="password", placeholder="sk_... (overrides backend settings)")
                system_prompt_input = gr.Textbox(label="System Prompt", lines=10, value=DEFAULT_SYSTEM_PROMPT, elem_id="system-prompt")

    with gr.Column(scale=2):
        gr.Markdown("## 💬 AI Assistant Chat")
        chatbot_display = gr.Chatbot(label="AI Chat", height=500, bubble_full_width=False, avatar_images=(None))
        with gr.Row():
            chat_message_input = gr.Textbox(show_label=False, placeholder="Your Message...", scale=7)
            send_chat_button = gr.Button("Send", variant="primary", scale=1)
        status_output = gr.Textbox(label="Last Action Status", interactive=False, value="Ready.")

    # Confirmation Accordion - Initially hidden
    with gr.Accordion("📝 Proposed Changes (Pending Confirmation)", open=False, visible=False) as confirm_accordion:
        changeset_display = gr.Markdown("*No changes proposed.*")
        with gr.Row():
            confirm_button = gr.Button("✅ Confirm & Apply Changes", variant="primary", visible=False)
            cancel_button = gr.Button("❌ Cancel", variant="stop", visible=False)

    with gr.Tabs():
        with gr.TabItem("📝 Generated Markdown & Build"):
            with gr.Row():
                with gr.Column(scale=2):
                    # This textbox shows the full markdown representation based on the *current state cache*
                    formatted_space_output_display = gr.Textbox(label="Current Space Definition (Generated Markdown)", lines=20, interactive=True, value="*Load or create a space to see its definition.*")
                    download_button = gr.DownloadButton(label="Download .md", interactive=False)
                with gr.Column(scale=1):
                    gr.Markdown("### Build Controls")
                    # Manual build controls
                    space_sdk_select = gr.Dropdown(label="Space SDK", choices=["gradio", "streamlit", "docker", "static"], value="gradio", interactive=True)
                    space_private_checkbox = gr.Checkbox(label="Make Space Private", value=False, interactive=True)
                    # Manual build button now builds from the content in `formatted_space_output_display`
                    build_space_button = gr.Button("🚀 Build / Update Space from Markdown", variant="primary")
                    build_status_display = gr.Textbox(label="Manual Build/Update Status", interactive=False)
                    gr.Markdown("---") # Separator
                    # Manual status check (uses build_logic)
                    refresh_status_button = gr.Button("🔄 Refresh Runtime Status")
                    space_runtime_status_display = gr.Markdown("*Runtime status will appear here.*")

        with gr.TabItem("🔍 Files Preview"):
             # This markdown shows the *latest* version of each file from the cache
            detected_files_preview = gr.Markdown(value="*A preview of the latest file versions will appear here.*")

        with gr.TabItem("✏️ Live File Editor & Preview"):
             with gr.Row():
                with gr.Column(scale=1):
                    gr.Markdown("### Live Editor")
                    # Dropdown lists files from the current state cache
                    file_browser_dropdown = gr.Dropdown(label="Select File in Space", choices=[], interactive=True)
                    # Editor for selected file content
                    file_content_editor = gr.Code(label="File Content Editor", language="python", lines=15, interactive=True, value="")
                    commit_message_input = gr.Textbox(label="Commit Message", placeholder="e.g., Updated app.py", interactive=True, value="")
                    with gr.Row():
                        # Manual file actions (use build_logic directly)
                        update_file_button = gr.Button("Commit Changes", variant="primary", interactive=True)
                        delete_file_button = gr.Button("🗑️ Delete Selected File", variant="stop", interactive=True)
                    edit_status_display = gr.Textbox(label="File Edit/Delete Status", interactive=False, value="")
                with gr.Column(scale=1):
                    gr.Markdown("### Live Space Preview")
                    # Iframe preview of the space (updates on load and successful build)
                    space_iframe_display = gr.HTML(value="", visible=True)

    # --- Event Listeners ---

    # Model dropdown update logic
    provider_select.change(update_models_dropdown, inputs=provider_select, outputs=model_select)

    # Chat submission logic
    chat_inputs = [
        chat_message_input, chatbot_display, hf_api_key_input,
        provider_api_key_input, provider_select, model_select, system_prompt_input,
        owner_name_input, space_name_input # Need current space info for context/actions
    ]
    chat_outputs = [
        chat_message_input, chatbot_display, status_output,
        detected_files_preview, formatted_space_output_display, download_button,
        changeset_state, changeset_display, confirm_accordion, confirm_button, cancel_button
    ]
    send_chat_button.click(handle_chat_submit, inputs=chat_inputs, outputs=chat_outputs)
    chat_message_input.submit(handle_chat_submit, inputs=chat_inputs, outputs=chat_outputs)

    # Confirmation Button Listeners for AI-proposed changes
    confirm_inputs = [hf_api_key_input, owner_name_input, space_name_input, changeset_state]
    confirm_outputs = [
        status_output, formatted_space_output_display, detected_files_preview, download_button,
        confirm_accordion, confirm_button, cancel_button, changeset_state, changeset_display # Also clear summary display
    ]
    confirm_button.click(handle_confirm_changes, inputs=confirm_inputs, outputs=confirm_outputs)

    cancel_outputs = [
        status_output, changeset_state, changeset_display,
        confirm_accordion, confirm_button, cancel_button
    ]
    cancel_button.click(handle_cancel_changes, inputs=None, outputs=cancel_outputs)

    # Load Existing Space Button logic
    load_space_outputs = [
        formatted_space_output_display, detected_files_preview, status_output,
        file_browser_dropdown, owner_name_input, space_name_input, # Update these inputs as well
        space_iframe_display, download_button, build_status_display,
        edit_status_display, space_runtime_status_display,
        chatbot_display, # Keep chat history
        changeset_state, changeset_display, confirm_accordion, confirm_button, cancel_button # Clear and hide confirm UI
    ]
    load_space_button.click(
        fn=handle_load_existing_space,
        inputs=[hf_api_key_input, owner_name_input, space_name_input],
        outputs=load_space_outputs
    )

    # Manual Build Button logic
    build_outputs = [
        build_status_display, space_iframe_display, file_browser_dropdown,
        owner_name_input, space_name_input, # Update inputs based on build result
        changeset_state, changeset_display, confirm_accordion, confirm_button, cancel_button, # Clear and hide confirm UI
        formatted_space_output_display, detected_files_preview, download_button # Update markdown displays
    ]
    build_inputs = [
        hf_api_key_input, space_name_input, owner_name_input, space_sdk_select,
        space_private_checkbox, formatted_space_output_display # Use content from this textbox
    ]
    build_space_button.click(fn=handle_build_space_button, inputs=build_inputs, outputs=build_outputs)


    # Manual File Editor Load logic
    file_edit_load_outputs = [file_content_editor, edit_status_display, commit_message_input, file_content_editor] # last one updates language
    file_browser_dropdown.change(fn=handle_load_file_for_editing, inputs=[hf_api_key_input, space_name_input, owner_name_input, file_browser_dropdown], outputs=file_edit_load_outputs)

    # Manual File Commit logic
    commit_file_outputs = [edit_status_display, file_browser_dropdown, formatted_space_output_display, detected_files_preview, download_button]
    update_file_button.click(fn=handle_commit_file_changes, inputs=[hf_api_key_input, space_name_input, owner_name_input, file_browser_dropdown, file_content_editor, commit_message_input], outputs=commit_file_outputs)

    # Manual File Delete logic
    delete_file_outputs = [
        edit_status_display, file_browser_dropdown,
        file_content_editor, commit_message_input, file_content_editor, # Clear editor fields
        formatted_space_output_display, detected_files_preview, download_button # Update markdown displays
    ]
    delete_file_button.click(fn=handle_delete_file, inputs=[hf_api_key_input, space_name_input, owner_name_input, file_browser_dropdown], outputs=delete_file_outputs)

    # Refresh Runtime Status logic
    refresh_status_button.click(fn=handle_refresh_space_status, inputs=[hf_api_key_input, owner_name_input, space_name_input], outputs=[space_runtime_status_display])

if __name__ == "__main__":
    demo.launch(debug=False, mcp_server=True)