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 (
        create_space as build_logic_create_space,
        _get_api_token as build_logic_get_api_token,
        whoami as build_logic_whoami,
        list_space_files_for_browsing,
        get_space_repository_info,
        get_space_file_content,
        update_space_file,
        parse_markdown as build_logic_parse_markdown,
        delete_space_file as build_logic_delete_space_file,
        get_space_runtime_status
    )
    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:
    print("Warning: Local modules (build_logic.py, model_logic.py) not found. Using dummy functions.")
    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."
    def build_logic_create_space(*args, **kwargs): return "Error: build_logic not found."
    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.")
    def get_space_repository_info(*args): return (None, [], "Error: build_logic not found.")
    def get_space_file_content(*args): return ("", "Error: build_logic not found.")
    def update_space_file(*args, **kwargs): return "Error: build_logic not found."
    def build_logic_parse_markdown(md): return {"files": []}
    def build_logic_delete_space_file(*args): return "Error: build_logic not found."
    def get_space_runtime_status(*args): return (None, "Error: build_logic not found.")
# --- END: Dummy functions ---


# --- New Feature Functions (can be moved to build_logic.py) ---
def build_logic_set_space_privacy(hf_api_key, repo_id, private: bool):
    """Sets the privacy of a Hugging Face Space."""
    print(f"[ACTION] Setting privacy for '{repo_id}' to {private}.")
    try:
        token, err = build_logic_get_api_token(hf_api_key)
        if err or not token: return f"Error getting token: {err or 'Token not found.'}"
        api = HfApi(token=token)
        api.update_repo_visibility(repo_id=repo_id, private=private, repo_type='space')
        return f"Successfully set privacy for {repo_id} to {private}."
    except Exception as e:
        print(f"Error setting privacy: {e}")
        return f"Error setting privacy: {e}"

def build_logic_delete_space(hf_api_key, owner, space_name):
    """Deletes an entire Hugging Face Space."""
    repo_id = f"{owner}/{space_name}"
    print(f"[ACTION] Deleting space '{repo_id}'. THIS IS A DESTRUCTIVE ACTION.")
    try:
        token, err = build_logic_get_api_token(hf_api_key)
        if err or not token: return f"Error getting token: {err or 'Token not found.'}"
        api = HfApi(token=token)
        api.delete_repo(repo_id=repo_id, repo_type='space')
        return f"Successfully deleted space {repo_id}."
    except Exception as e:
        print(f"Error deleting space: {e}")
        return f"Error deleting space: {e}"


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

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.

**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:
## 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.
- Filenames in the '### File:' line should be clean paths (e.g., 'src/app.py', 'Dockerfile') and should NOT include Markdown backticks.

**Hugging Face Space Actions:**
To perform direct actions on the Hugging Face Space, use the `HF_ACTION` command. This is a powerful tool to manage the repository programmatically.
The format is `### HF_ACTION: COMMAND arguments...` on a single line.

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.
- `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.

You can issue multiple actions. For example, to delete a file and then add a new one:
### HF_ACTION: DELETE_FILE old_app.py
### File: new_app.py
{bbb}python
# new code
{bbb}
Use these actions when the user's request explicitly calls for them (e.g., "delete the readme file", "make this space private", "create a new private space called my-test-app"). If no code is provided, assist the user with their tasks.
"""

# --- Helper Functions (largely unchanged) ---
def escape_html_for_markdown(text):
    if not isinstance(text, str): return ""
    return text.replace("&", "&").replace("<", "<").replace(">", ">")

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()
    text = re.sub(r'[`\*_]+', '', text)
    path_match = re.match(r'^([\w\-\.\s\/\\]+)', text)
    if path_match:
        parts = re.split(r'\s*\(', path_match.group(1).strip(), 1)
        return parts[0].strip() if parts else ""
    backtick_match = re.search(r'`([^`]+)`', text)
    if backtick_match:
        potential_fn = backtick_match.group(1).strip()
        parts = re.split(r'\s*\(|\s{2,}', potential_fn, 1)
        cleaned_fn = parts[0].strip() if parts else ""
        cleaned_fn = cleaned_fn.strip('`\'":;,')
        if cleaned_fn: return cleaned_fn
    parts = re.split(r'\s*\(|\s{2,}', text, 1)
    filename_candidate = parts[0].strip() if parts else text.strip()
    filename_candidate = filename_candidate.strip('`\'":;,')
    return filename_candidate if filename_candidate else text.strip()

def _parse_chat_stream_logic(latest_bot_message_content, existing_files_state=None):
    global parsed_code_blocks_state_cache
    latest_blocks_dict = {}
    if existing_files_state:
        for block in existing_files_state:
             if not block.get("is_structure_block"):
                 latest_blocks_dict[block["filename"]] = block.copy()

    results = {"parsed_code_blocks": [], "preview_md": "", "default_selected_filenames": [], "error_message": 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(?: - [^\]]+)?\]))")
    structure_pattern = re.compile(r"## File Structure\n```(?:(?P<struct_lang>[\w.-]*)\n)?(?P<structure_code>[\s\S]*?)\n```")

    structure_match = structure_pattern.search(content)
    if structure_match:
        latest_blocks_dict["File Structure (original)"] = {"filename": "File Structure (original)", "language": structure_match.group("struct_lang") or "plaintext", "code": structure_match.group("structure_code").strip(), "is_binary": False, "is_structure_block": True}
    else:
         existing_structure_block = next((b for b in parsed_code_blocks_state_cache if b.get("is_structure_block")), None)
         if existing_structure_block:
              latest_blocks_dict["File Structure (original)"] = existing_structure_block.copy()

    current_message_file_blocks = {}
    for match in file_pattern.finditer(content):
        filename = _clean_filename(match.group("filename_line"))
        if not filename: continue
        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"], item_data["language"] = code_block.strip(), (lang.strip().lower() if lang else _infer_lang_from_filename(filename))
        elif binary_msg is not None:
            item_data["code"], item_data["language"], item_data["is_binary"] = binary_msg.strip(), "binary", True
        else: continue
        current_message_file_blocks[filename] = item_data

    latest_blocks_dict.update(current_message_file_blocks)
    current_parsed_blocks = list(latest_blocks_dict.values())
    current_parsed_blocks.sort(key=lambda b: (0, b["filename"]) if b.get("is_structure_block") else (1, b["filename"]))
    results["parsed_code_blocks"] = current_parsed_blocks
    results["default_selected_filenames"] = [b["filename"] for b in current_parsed_blocks if not b.get("is_structure_block")]
    return results

def _export_selected_logic(selected_filenames, space_line_name_for_md, parsed_blocks_for_export):
    results = {"output_str": "", "error_message": None, "download_filepath": None}
    exportable_blocks_content = [b for b in parsed_blocks_for_export if not b.get("is_structure_block") and not b.get("is_binary") and not (b.get("code", "").startswith(("[Error loading content:", "[Binary or Skipped file]")))]
    binary_blocks_content = [b for b in parsed_blocks_for_export if b.get("is_binary") or b.get("code", "").startswith("[Binary or Skipped file]")]
    all_filenames_in_state = sorted(list(set(b["filename"] for b in parsed_blocks_for_export if not b.get("is_structure_block"))))

    if not all_filenames_in_state and not any(b.get("is_structure_block") for b in parsed_blocks_for_export):
        results["output_str"] = f"# Space: {space_line_name_for_md}\n## File Structure\n{bbb}\n📁 Root\n{bbb}\n\n*No files to list in structure or export.*"
        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 export: {e}")
        return results

    output_lines = [f"# Space: {space_line_name_for_md}"]
    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:
        output_lines.extend(["## File Structure", bbb, "📁 Root"])
        if all_filenames_in_state:
            for fname in all_filenames_in_state: output_lines.append(f"  📄 {fname}")
        output_lines.extend([bbb, ""])

    output_lines.append("Below are the contents of all files in the space:\n")
    files_to_export_content = [b for b in exportable_blocks_content if not selected_filenames or b["filename"] in selected_filenames]
    binary_error_blocks_to_export = [b for b in binary_blocks_content if not selected_filenames or b["filename"] in selected_filenames]
    all_blocks_to_export_content = sorted(files_to_export_content + binary_error_blocks_to_export, key=lambda b: b["filename"])

    exported_content = False
    for block in all_blocks_to_export_content:
        output_lines.append(f"### File: {block['filename']}")
        if block.get('is_binary') or block.get("code", "").startswith(("[Binary file", "[Error loading content:", "[Binary or Skipped file]")):
            output_lines.append(block.get('code','[Binary or Skipped file]'))
        else:
            output_lines.extend([f"{bbb}{block.get('language', 'plaintext') or 'plaintext'}", block.get('code',''), bbb])
        output_lines.append(""); exported_content = True

    if not exported_content and not all_filenames_in_state: output_lines.append("*No files in state.*")
    elif not exported_content: output_lines.append("*No files with editable content are in the state or selected.*")

    final_output_str = "\n".join(output_lines)
    results["output_str"] = final_output_str
    try:
        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: {e}")
        results["error_message"] = "Could not prepare file for download."
    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})
    if current_user_message: messages.append({"role": "user", "content": current_user_message})
    return messages

def _generate_ui_outputs_from_cache(owner, space_name):
    global parsed_code_blocks_state_cache
    preview_md_val = "*No files in cache to display.*"
    formatted_md_val = f"# Space: {owner}/{space_name}\n## File Structure\n{bbb}\n📁 Root\n{bbb}\n\n*No files in cache.*" if owner and space_name else "*Load or define a Space to see its Markdown structure.*"
    download_file = None

    if parsed_code_blocks_state_cache:
        preview_md_lines = ["## Detected/Updated Files & Content (Latest Versions):"]
        for block in parsed_code_blocks_state_cache:
            preview_md_lines.append(f"\n----\n**File:** `{escape_html_for_markdown(block['filename'])}`")
            if block.get('is_structure_block'): preview_md_lines.append(f" (Original File Structure from AI)\n")
            elif block.get('is_binary'): preview_md_lines.append(f" (Binary File)\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)}`\n")
            else:
                 preview_md_lines.append(f"\n{bbb}{block.get('language', 'plaintext') or 'plaintext'}\n{content}\n{bbb}\n")
        preview_md_val = "\n".join(preview_md_lines)
        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"]

    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, and generates a structured changeset.
    Returns the changeset and a markdown summary for display.
    """
    changeset = []
    current_files_dict = {f["filename"]: f for f in current_files_state if not f.get("is_structure_block")}

    # 1. Parse proposed files from AI response
    parsing_result = _parse_chat_stream_logic(ai_response_content, existing_files_state=current_files_state)
    proposed_files = parsing_result.get("parsed_code_blocks", [])

    # 2. Parse HF_ACTION commands from AI response
    action_pattern = re.compile(r"### HF_ACTION:\s*(?P<command_line>[^\n]+)")
    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 == "DELETE_FILE" and args:
            changeset.append({"type": "DELETE_FILE", "path": args[0]})
        elif command == "SET_PRIVATE" and args:
            changeset.append({"type": "SET_PRIVACY", "private": args[0].lower() == 'true', "repo_id": f"{hf_owner_name}/{hf_repo_name}"})
        elif command == "DELETE_SPACE":
            changeset.append({"type": "DELETE_SPACE", "owner": hf_owner_name, "space_name": hf_repo_name})
        elif command == "CREATE_SPACE" and args:
            repo_id = args[0]
            sdk = "gradio"  # default
            private = False # default
            if '--sdk' in args: sdk = args[args.index('--sdk') + 1]
            if '--private' in args: private = args[args.index('--private') + 1].lower() == 'true'
            changeset.append({"type": "CREATE_SPACE", "repo_id": repo_id, "sdk": sdk, "private": private})

    # 3. Compare proposed files with current files to determine CREATE/UPDATE
    for file_block in proposed_files:
        if file_block.get("is_structure_block"): continue

        filename = file_block["filename"]
        if filename not in current_files_dict:
            changeset.append({"type": "CREATE_FILE", "path": filename, "content": file_block["code"], "lang": file_block["language"]})
        elif file_block["code"] != current_files_dict[filename]["code"]:
            changeset.append({"type": "UPDATE_FILE", "path": filename, "content": file_block["code"], "lang": file_block["language"]})

    # 4. Format the changeset into a human-readable Markdown string
    if not changeset:
        return [], "The AI did not propose any specific changes to files or the space.", parsing_result

    md_summary = ["### 📋 Proposed Changes Plan\n"]
    md_summary.append("The AI has proposed the following changes. Please review and confirm.")

    for change in changeset:
        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']}`")
        elif change["type"] == "CREATE_SPACE":
            md_summary.append(f"- **🚀 Create New Space:** `{change['repo_id']}` (SDK: {change['sdk']}, Private: {change['private']})")
        elif change["type"] == "SET_PRIVACY":
            md_summary.append(f"- **🔒 Set Privacy:** Set `{change['repo_id']}` to `private={change['private']}`")
        elif change["type"] == "DELETE_SPACE":
            md_summary.append(f"- **💥 DELETE ENTIRE SPACE:** `{change['owner']}/{change['space_name']}` **(DESTRUCTIVE ACTION)**")

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

# --- 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)

    # UI updates for streaming
    yield (
        _chat_msg_in, _chat_hist, "Initializing...",
        gr.update(), gr.update(), gr.update(interactive=False),
        [], gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
    )

    if not user_message.strip():
        yield (
            _chat_msg_in, _chat_hist, "Cannot send an empty message.",
            gr.update(), gr.update(), gr.update(),
            [], gr.update(visible=False), 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(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
    )

    # Prepare context for the AI
    current_sys_prompt = system_prompt.strip() or DEFAULT_SYSTEM_PROMPT
    export_result = _export_selected_logic(None, f"{hf_owner_name}/{hf_repo_name}", parsed_code_blocks_state_cache)
    current_files_context = f"\n\n## Current Space Context: {hf_owner_name}/{hf_repo_name}\n{export_result['output_str']}"
    user_message_with_context = user_message.strip() + "\n" + 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) 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(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
            )

        if "Error:" in full_bot_response_content:
            _status = full_bot_response_content
            yield (_chat_msg_in, _chat_hist, _status, gr.update(), gr.update(), gr.update(), [], gr.update(), gr.update(), gr.update(), gr.update())
            return

        # Instead of applying, generate and stage changes
        _status = "Stream complete. Generating change plan..."
        yield (_chat_msg_in, _chat_hist, _status, gr.update(), gr.update(), gr.update(), [], gr.update(), gr.update(), gr.update(), gr.update())

        staged_changeset, summary_md, parsing_res = generate_and_stage_changes(full_bot_response_content, parsed_code_blocks_state_cache, hf_owner_name, hf_repo_name)

        if parsing_res["error_message"]:
            _status = f"Parsing Error: {parsing_res['error_message']}"
            yield (_chat_msg_in, _chat_hist, _status, gr.update(), gr.update(), gr.update(), [], gr.update(), gr.update(), gr.update(), gr.update())
            return

        if not staged_changeset:
            _status = summary_md # "No changes proposed" message
            # Still update the cache with the AI's *view* of the world, even if no changes.
            parsed_code_blocks_state_cache = parsing_res["parsed_code_blocks"]
            _formatted, _detected, _download = _generate_ui_outputs_from_cache(hf_owner_name, hf_repo_name)
            yield (_chat_msg_in, _chat_hist, _status, _detected, _formatted, _download, [], gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False))
        else:
            _status = "Change plan generated. Please review and confirm below."
            yield (
                _chat_msg_in, _chat_hist, _status,
                gr.update(), gr.update(), gr.update(),
                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}")
        if _chat_hist: _chat_hist[-1] = (user_message, error_msg)
        yield (
            _chat_msg_in, _chat_hist, error_msg,
            gr.update(), gr.update(), gr.update(),
            [], gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
        )

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
    if not changeset:
        return "No changes to apply.", gr.update(), gr.update(), gr.update(), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)

    status_messages = []

    # Handle space creation first, as other ops might depend on it
    create_space_op = next((c for c in changeset if c['type'] == 'CREATE_SPACE'), None)
    if create_space_op:
        repo_parts = create_space_op['repo_id'].split('/')
        if len(repo_parts) == 2:
            owner, repo = repo_parts
            # We need to pass the full markdown for creation. Let's build it from the plan.
            # This is a simplification; a more robust solution would pass the planned files directly.
            # For now, we assume the AI provides file content for the new space.

            planned_files_md = [f"# Space: {create_space_op['repo_id']}"]
            for change in changeset:
                if change['type'] in ['CREATE_FILE', 'UPDATE_FILE']:
                     planned_files_md.append(f"### File: {change['path']}\n{bbb}{change.get('lang', 'plaintext')}\n{change['content']}\n{bbb}")

            markdown_for_creation = "\n\n".join(planned_files_md)

            result = build_logic_create_space(
                ui_api_token_from_textbox=hf_api_key,
                space_name_ui=repo,
                owner_ui=owner,
                sdk_ui=create_space_op['sdk'],
                private=create_space_op['private'],
                markdown_input=markdown_for_creation
            )
            status_messages.append(f"CREATE_SPACE: {result}")
            if "Error" in result:
                # Stop if space creation failed
                final_status = " | ".join(status_messages)
                return final_status, gr.update(), gr.update(), gr.update(), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), []

    # Apply all other changes
    for change in changeset:
        try:
            if change['type'] == 'UPDATE_FILE':
                msg = update_space_file(hf_api_key, space_name, owner_name, change['path'], change['content'], f"AI: Update {change['path']}")
                status_messages.append(f"UPDATE '{change['path']}': {msg}")
                if "Success" in msg:
                    # Update cache on success
                    for block in parsed_code_blocks_state_cache:
                        if block['filename'] == change['path']:
                            block['code'] = change['content']
                            break
            elif change['type'] == 'CREATE_FILE' and not create_space_op: # Don't re-create if handled by CREATE_SPACE
                msg = update_space_file(hf_api_key, space_name, owner_name, change['path'], change['content'], f"AI: Create {change['path']}")
                status_messages.append(f"CREATE '{change['path']}': {msg}")
                if "Success" in msg:
                    parsed_code_blocks_state_cache.append({'filename': change['path'], 'code': change['content'], 'language': change['lang'], 'is_binary': False})
            elif change['type'] == 'DELETE_FILE':
                msg = build_logic_delete_space_file(hf_api_key, space_name, owner_name, change['path'])
                status_messages.append(f"DELETE '{change['path']}': {msg}")
                if "Success" in msg:
                    parsed_code_blocks_state_cache = [b for b in parsed_code_blocks_state_cache if b["filename"] != change['path']]
            elif change['type'] == 'SET_PRIVACY':
                msg = build_logic_set_space_privacy(hf_api_key, change['repo_id'], change['private'])
                status_messages.append(f"SET_PRIVACY: {msg}")
            elif change['type'] == 'DELETE_SPACE':
                msg = build_logic_delete_space(hf_api_key, change['owner'], change['space_name'])
                status_messages.append(f"DELETE_SPACE: {msg}")
                if "Success" in msg:
                    parsed_code_blocks_state_cache = [] # Clear everything
        except Exception as e:
            status_messages.append(f"Error applying {change['type']} for {change.get('path', '')}: {e}")

    final_status = " | ".join(status_messages)
    _formatted, _detected, _download = _generate_ui_outputs_from_cache(owner_name, space_name)

    # Hide the confirmation UI and clear the state
    return final_status, _formatted, _detected, _download, gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), []

def handle_cancel_changes():
    """Clears the staged changeset and hides the confirmation UI."""
    return "Changes cancelled.", [], gr.update(value="*No changes proposed.*"), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)

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 = []
    outputs = [_formatted_md_val, _detected_preview_val, _status_val, _file_browser_update, gr.update(value=ui_owner_name), gr.update(value=ui_space_name), _iframe_html_update, _download_btn_update, _build_status_clear, _edit_status_clear, _runtime_status_clear, _chat_history_clear]
    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)
            _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)
    if err and not file_list:
        _status_val = f"File List Error: {err}"; parsed_code_blocks_state_cache = []
        _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
        yield outputs; return

    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)

    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
        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
    _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}."
    outputs[0], outputs[1], outputs[2], outputs[7] = _formatted, _detected, _status_val, _download
    outputs[3] = gr.update(visible=True, choices=sorted(file_list or []), value=None)
    yield outputs

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 space build process...", gr.update(value=None, visible=False), gr.update(visible=False, choices=[], value=None)
    yield _build_status, _iframe_html, _file_browser_update, gr.update(value=ui_owner_name_part), gr.update(value=ui_space_name_part)
    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(); return

    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"Build Process: {result_message}"

    if "Successfully" in result_message:
        sub_owner = re.sub(r'[^a-z0-9\-]+', '-', ui_owner_name_part.lower()).strip('-')
        sub_repo = re.sub(r'[^a-z0-9\-]+', '-', ui_space_name_part.lower()).strip('-')
        iframe_url = f"https://{sub_owner}-{sub_repo}{'.static.hf.space' if space_sdk_ui == '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)
        file_list, err = list_space_files_for_browsing(hf_api_key_ui, ui_space_name_part, ui_owner_name_part)
        _file_browser_update = gr.update(visible=True, choices=sorted(file_list or []), value=None)
    yield _build_status, _iframe_html, _file_browser_update, gr.update(value=ui_owner_name_part), gr.update(value=ui_space_name_part)

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:
        yield f"Error: {err}", 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):
    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)
    global parsed_code_blocks_state_cache
    if "Successfully" in status_msg:
        # Update cache
        for block in parsed_code_blocks_state_cache:
            if block["filename"] == file_to_edit_path:
                block["code"] = edited_content
                break
    _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)
    global parsed_code_blocks_state_cache
    if "Successfully" in status_msg:
        parsed_code_blocks_state_cache = [b for b in parsed_code_blocks_state_cache if b["filename"] != file_to_delete_path]
    _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), "", "", "plaintext", _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, err = get_space_runtime_status(hf_api_key_ui, ui_space_name, ui_owner_name)
    if err: return f"**Error:** {err}"
    if not status: return "Could not retrieve status."

    md = f"### Status for {ui_owner_name}/{ui_space_name}\n"
    for key, val in status.items():
        md += f"- **{key.replace('_', ' ').title()}:** `{val}`\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; }
.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; }
"""

# --- Gradio UI Definition ---
with gr.Blocks(theme=custom_theme, css=custom_css) as demo:
    # State to hold the plan
    changeset_state = gr.State([])

    gr.Markdown("# 🤖 AI-Powered Hugging Face Space Builder")
    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")

    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="my-ai-space")
                load_space_button = gr.Button("🔄 Load Existing Space", variant="secondary")

            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
                if default_provider not in available_providers:
                    default_provider = available_providers[2] 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
                )
                model_select = gr.Dropdown(
                    label="AI Model",
                    choices=initial_models,
                    value=initial_model
                )
                # --- 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
    with gr.Accordion("📝 Proposed Changes (Pending Confirmation)", 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):
                    formatted_space_output_display = gr.Textbox(label="Current Space Definition (Editable)", 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")
                    space_sdk_select = gr.Dropdown(label="Space SDK", choices=["gradio", "streamlit", "docker", "static"], value="gradio")
                    space_private_checkbox = gr.Checkbox(label="Make Space Private", value=False)
                    build_space_button = gr.Button("🚀 Build / Update Space from Manual Edit", variant="primary")
                    build_status_display = gr.Textbox(label="Build Operation Status", interactive=False)
                    refresh_status_button = gr.Button("🔄 Refresh Runtime Status")
                    space_runtime_status_display = gr.Markdown("*Runtime status will appear here.*")

        with gr.TabItem("🔍 Files Preview"):
            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")
                    file_browser_dropdown = gr.Dropdown(label="Select File in Space", choices=[], interactive=True)
                    file_content_editor = gr.Code(label="File Content Editor", language="python", lines=15, interactive=True)
                    commit_message_input = gr.Textbox(label="Commit Message", placeholder="e.g., Updated app.py")
                    with gr.Row():
                        update_file_button = gr.Button("Commit Changes", variant="primary")
                        delete_file_button = gr.Button("🗑️ Delete Selected File", variant="stop")
                    edit_status_display = gr.Textbox(label="File Edit/Delete Status", interactive=False)
                with gr.Column(scale=1):
                    gr.Markdown("### Live Space Preview")
                    space_iframe_display = gr.HTML(value="", visible=True)

    # --- Event Listeners ---
    provider_select.change(update_models_dropdown, inputs=provider_select, outputs=model_select)

    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]
    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
    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
    ]
    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_space_outputs = [formatted_space_output_display, detected_files_preview, status_output, file_browser_dropdown, owner_name_input, space_name_input, space_iframe_display, download_button, build_status_display, edit_status_display, space_runtime_status_display, chatbot_display]
    load_space_button.click(fn=handle_load_existing_space, inputs=[hf_api_key_input, owner_name_input, space_name_input], outputs=load_space_outputs)

    build_outputs = [build_status_display, space_iframe_display, file_browser_dropdown, owner_name_input, space_name_input]
    build_inputs = [hf_api_key_input, space_name_input, owner_name_input, space_sdk_select, space_private_checkbox, formatted_space_output_display]
    build_space_button.click(fn=handle_build_space_button, inputs=build_inputs, outputs=build_outputs)

    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)

    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)

    delete_file_outputs = [edit_status_display, file_browser_dropdown, file_content_editor, commit_message_input, file_content_editor, formatted_space_output_display, detected_files_preview, download_button]
    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_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)