Local LM Visual Studio Code Helper (Ollama)

import os
import sys
import shutil
import subprocess
from datetime import datetime
import requests
import re
import threading
import time

# =============================================================================
# Configuration
# If not hosted locally you can simply run this file as python code_helper_v6.py <file_path> <prompt>
# =============================================================================
OLLAMA_API_URL = "http://localhost:11434/api/generate"
MODEL_NAME = "qwen2.5-coder:3b"

# Disclaimer to be injected into C# files when the user requests it via the prompt add disclaimer
DISCLAIMER_TEXT = """/*
 * -----------------------------------------------------------------------------
 * © 2026 Skyline Cloud LLC. All rights reserved.
 * NOTICE: This source code file contains structural modifications optimized via
 * an internally hosted, private Qwen LLM instance. All data remains subject 
 * to strict corporate security and data sovereignty policies.
 * SkylineCloud.AI is a proprietary tool and may not be redistributed or used outside of authorized environments.
 * -----------------------------------------------------------------------------
 */"""

# =============================================================================
# Visual Spinner & External Prompt Loading
# architect_rules.txt will be auto-generated if missing, ensuring consistent LLM behavior across environments.
# Updates to the rules file will be reflected in subsequent LLM interactions without requiring code changes.
# =============================================================================
def show_spinner(stop_event):
    spinner = ['⠋', '⠙', '⠹', '⠸', '⠼', '⠴', '⠦', '⠧', '⠇', '⠏']
    idx = 0
    while not stop_event.is_set():
        sys.stdout.write(f"\r[INFO] Skyline LLM is thinking... {spinner[idx]} ")
        sys.stdout.flush()
        idx = (idx + 1) % len(spinner)
        time.sleep(0.1)
    sys.stdout.write("\r" + " "*50 + "\r") # Clears the line cleanly

def get_system_prompt(file_path):
    # Determine the file extension to load language-specific rules
    _, ext = os.path.splitext(file_path.lower())
    script_dir = os.path.dirname(os.path.abspath(__file__))
    
    # Define mapping of file extensions to target rule files and their initial default prompts
    agent_configs = {
        ".cs": {
            "filename": "architect_rules_cs.txt",
            "default_prompt": (
                "You are an expert C# refactoring assistant. Your job is to modify the provided C# code "
                "according to the user's instructions. You MUST return ONLY the raw, updated C# source code. "
                "Do NOT include any introduction, explanations, markdown code blocks (like ```csharp), or trailing text. "
                "Preserve the existing structure, namespaces, and formatting of the class as much as possible."
            )
        },
        ".py": {
            "filename": "architect_rules_py.txt",
            "default_prompt": (
                "You are an expert Python programming assistant. Your job is to modify the provided Python code "
                "according to the user's instructions. You MUST return ONLY the raw, updated Python source code. "
                "Do NOT include any introduction, explanations, markdown code blocks (like ```python), or trailing text. "
                "Adhere to PEP 8 guidelines and preserve formatting as much as possible."
            )
        },
        ".sql": {
            "filename": "architect_rules_sql.txt",
            "default_prompt": (
                "You are an expert SQL database assistant. Your job is to modify or generate database scripts "
                "according to the user's instructions. You MUST return ONLY valid SQL commands. "
                "Do NOT include markdown syntax (like ```sql), conversational introductions, or explanations. "
                "Ensure syntax matches standard ANSI SQL or T-SQL conventions where applicable."
            )
        }
    }
    
    # Fallback for generic/unsupported file extensions
    fallback_config = {
        "filename": "architect_rules.txt",
        "default_prompt": (
            "You are an expert software development assistant. Your job is to modify or generate code "
            "according to the user's instructions. You MUST return ONLY the raw, updated source code. "
            "Do NOT include markdown syntax wrapping, introductions, or structural explanations."
        )
    }
    
    # Select configuration based on extension
    config = agent_configs.get(ext, fallback_config)
    prompt_file = os.path.join(script_dir, config["filename"])
    
    if os.path.exists(prompt_file):
        with open(prompt_file, 'r', encoding='utf-8') as f:
            return f.read().strip()
    else:
        # Auto-generate the appropriate configuration file if it's missing
        with open(prompt_file, 'w', encoding='utf-8') as f:
            f.write(config["default_prompt"])
        print(f"[INFO] Created default agent rule configuration at '{config['filename']}'")
        return config["default_prompt"]

# =============================================================================
# CORE EXECUTION
# Pass an existing file path and a prompt to modify it, or a new file path to scaffold a new C# file.
# =============================================================================
def modify_code(file_path, prompt):
    file_exists = os.path.exists(file_path)
    original_code = ""

    # FEATURE 1: Handle File Generation vs Refactoring
    if file_exists:
        with open(file_path, 'r', encoding='utf-8') as f:
            original_code = f.read()

        try:
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            backup_path = f"{file_path}.{timestamp}.bak"
            shutil.copy2(file_path, backup_path)
            print(f"[INFO] Backup created: '{os.path.basename(backup_path)}'")
        except Exception as e:
            print(f"[WARN] Failed to create backup. Error: {e}")
            return
            
        # CONDITIONAL BYPASS: Local File Mutation (Only works if file exists)
        if "add disclaimer" in prompt.lower():
            print("[INFO] Shortcut matched: Injecting local corporate disclaimer...")
            using_pattern = re.compile(r'((?:^using\s+[^;]+;\s*\r?\n)+)', re.MULTILINE)
            match = using_pattern.search(original_code)
            
            if match:
                end_of_usings = match.end()
                updated_code = original_code[:end_of_usings] + "\n" + DISCLAIMER_TEXT + "\n" + original_code[end_of_usings:]
            else:
                updated_code = DISCLAIMER_TEXT + "\n\n" + original_code

            with open(file_path, 'w', encoding='utf-8') as f:
                f.write(updated_code)
            print(f"[SUCCESS] Disclaimer applied to {os.path.basename(file_path)} without LLM dependency.")
            return

    else:
        print(f"[INFO] File does not exist. Initiating scaffolding for '{os.path.basename(file_path)}'...")
        # Create an empty file so VS Code Diff viewer doesn't crash on the left side
        os.makedirs(os.path.dirname(file_path) or '.', exist_ok=True)
        open(file_path, 'a').close()

    print(f"[NETWORK] Sending '{os.path.basename(file_path)}' to Skyline Cloud LLM...")

    # Build Context-Aware Prompt (Now routing rules based on extension)
    system_instruction = get_system_prompt(file_path)
    
    if file_exists:
        full_prompt = f"{system_instruction}\n\nOriginal Code:\n{original_code}\n\nInstruction:\n{prompt}\n\nUpdated Code:"
    else:
        full_prompt = f"{system_instruction}\n\nInstruction: Scaffold a completely new file based on the following request.\n\nRequest:\n{prompt}\n\nGenerated Code:"

    # FEATURE 2: Expand Context Window (num_ctx)
    payload = {
        "model": MODEL_NAME,
        "prompt": full_prompt,
        "stream": False,
        "options": {
            "temperature": 0.2,
            "num_ctx": 8192 
        }
    }

    # Execute Network Call with Threaded Spinner
    stop_spinner = threading.Event()
    spinner_thread = threading.Thread(target=show_spinner, args=(stop_spinner,))
    spinner_thread.start()

    try:
        # FEATURE 3: Timeout Fail-Safe
        response = requests.post(OLLAMA_API_URL, json=payload, timeout=120)
        response.raise_for_status()
    except requests.exceptions.Timeout:
        stop_spinner.set()
        print(f"\n[FATAL] The LLM timed out after 120 seconds. Aborting.")
        return
    except requests.exceptions.RequestException as e:
        stop_spinner.set()
        print(f"\n[FATAL] Error communicating with Ollama: {e}")
        return
    finally:
        stop_spinner.set() # Guarantee the spinner stops no matter what happens
        spinner_thread.join()

    # Extract the updated code
    updated_code = response.json().get("response", "").strip()
    
    # Strip markdown block wrappers cleanly
    if updated_code.startswith("```"):
        lines = updated_code.splitlines()
        if lines and lines[0].startswith("```"):
            lines = lines[1:]
        if lines and lines[-1].startswith("```"):
            lines = lines[:-1]
        updated_code = "\n".join(lines).strip()

    if not updated_code:
        print("[ERROR] Received empty response from the model.")
        return

    # Create temporary proposal
    proposal_path = f"{file_path}.proposal"
    with open(proposal_path, 'w', encoding='utf-8') as f:
        f.write(updated_code)

    # Launch Visual Diff
    print(f"[INFO] Launching visual diff in VS Code...")
    subprocess.run(["code", "--diff", file_path, proposal_path], shell=True)

    # Await User Confirmation
    print("\n" + "="*50)
    user_decision = input(f"Accept LLM changes to {os.path.basename(file_path)}? (y/n): ").strip().lower()
    print("="*50)

    if user_decision == 'y':
        shutil.move(proposal_path, file_path)
        print(f"[SUCCESS] Changes committed to '{os.path.basename(file_path)}'.")
    else:
        os.remove(proposal_path)
        # If we scaffolded an empty file and rejected it, clean up the empty file
        if not file_exists:
            os.remove(file_path)
        print(f"[INFO] Changes discarded. Codebase remains untouched.")


if __name__ == "__main__":
    if len(sys.argv) < 3:
        print("Usage: ai <file_path> <prompt...>")
    else:
        target_file = sys.argv[1]
        user_prompt = " ".join(sys.argv[2:])
        modify_code(target_file, user_prompt)