Full Code of JosefAlbers/VimLM for AI

Repository: JosefAlbers/VimLM
Branch: main
Commit: c69cccb72453
Files: 6
Total size: 56.5 KB

Directory structure:
gitextract_jxuzkf6f/

├── LICENSE
├── README.md
├── requirements.txt
├── setup.py
├── tutorial.md
└── vimlm.py

================================================
FILE CONTENTS
================================================

================================================
FILE: LICENSE
================================================

                                 Apache License
                           Version 2.0, January 2004
                        http://www.apache.org/licenses/

   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION

   1. Definitions.

      "License" shall mean the terms and conditions for use, reproduction,
      and distribution as defined by Sections 1 through 9 of this document.

      "Licensor" shall mean the copyright owner or entity authorized by
      the copyright owner that is granting the License.

      "Legal Entity" shall mean the union of the acting entity and all
      other entities that control, are controlled by, or are under common
      control with that entity. For the purposes of this definition,
      "control" means (i) the power, direct or indirect, to cause the
      direction or management of such entity, whether by contract or
      otherwise, or (ii) ownership of fifty percent (50%) or more of the
      outstanding shares, or (iii) beneficial ownership of such entity.

      "You" (or "Your") shall mean an individual or Legal Entity
      exercising permissions granted by this License.

      "Source" form shall mean the preferred form for making modifications,
      including but not limited to software source code, documentation
      source, and configuration files.

      "Object" form shall mean any form resulting from mechanical
      transformation or translation of a Source form, including but
      not limited to compiled object code, generated documentation,
      and conversions to other media types.

      "Work" shall mean the work of authorship, whether in Source or
      Object form, made available under the License, as indicated by a
      copyright notice that is included in or attached to the work
      (an example is provided in the Appendix below).

      "Derivative Works" shall mean any work, whether in Source or Object
      form, that is based on (or derived from) the Work and for which the
      editorial revisions, annotations, elaborations, or other modifications
      represent, as a whole, an original work of authorship. For the purposes
      of this License, Derivative Works shall not include works that remain
      separable from, or merely link (or bind by name) to the interfaces of,
      the Work and Derivative Works thereof.

      "Contribution" shall mean any work of authorship, including
      the original version of the Work and any modifications or additions
      to that Work or Derivative Works thereof, that is intentionally
      submitted to Licensor for inclusion in the Work by the copyright owner
      or by an individual or Legal Entity authorized to submit on behalf of
      the copyright owner. For the purposes of this definition, "submitted"
      means any form of electronic, verbal, or written communication sent
      to the Licensor or its representatives, including but not limited to
      communication on electronic mailing lists, source code control systems,
      and issue tracking systems that are managed by, or on behalf of, the
      Licensor for the purpose of discussing and improving the Work, but
      excluding communication that is conspicuously marked or otherwise
      designated in writing by the copyright owner as "Not a Contribution."

      "Contributor" shall mean Licensor and any individual or Legal Entity
      on behalf of whom a Contribution has been received by Licensor and
      subsequently incorporated within the Work.

   2. Grant of Copyright License. Subject to the terms and conditions of
      this License, each Contributor hereby grants to You a perpetual,
      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
      copyright license to reproduce, prepare Derivative Works of,
      publicly display, publicly perform, sublicense, and distribute the
      Work and such Derivative Works in Source or Object form.

   3. Grant of Patent License. Subject to the terms and conditions of
      this License, each Contributor hereby grants to You a perpetual,
      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
      (except as stated in this section) patent license to make, have made,
      use, offer to sell, sell, import, and otherwise transfer the Work,
      where such license applies only to those patent claims licensable
      by such Contributor that are necessarily infringed by their
      Contribution(s) alone or by combination of their Contribution(s)
      with the Work to which such Contribution(s) was submitted. If You
      institute patent litigation against any entity (including a
      cross-claim or counterclaim in a lawsuit) alleging that the Work
      or a Contribution incorporated within the Work constitutes direct
      or contributory patent infringement, then any patent licenses
      granted to You under this License for that Work shall terminate
      as of the date such litigation is filed.

   4. Redistribution. You may reproduce and distribute copies of the
      Work or Derivative Works thereof in any medium, with or without
      modifications, and in Source or Object form, provided that You
      meet the following conditions:

      (a) You must give any other recipients of the Work or
          Derivative Works a copy of this License; and

      (b) You must cause any modified files to carry prominent notices
          stating that You changed the files; and

      (c) You must retain, in the Source form of any Derivative Works
          that You distribute, all copyright, patent, trademark, and
          attribution notices from the Source form of the Work,
          excluding those notices that do not pertain to any part of
          the Derivative Works; and

      (d) If the Work includes a "NOTICE" text file as part of its
          distribution, then any Derivative Works that You distribute must
          include a readable copy of the attribution notices contained
          within such NOTICE file, excluding those notices that do not
          pertain to any part of the Derivative Works, in at least one
          of the following places: within a NOTICE text file distributed
          as part of the Derivative Works; within the Source form or
          documentation, if provided along with the Derivative Works; or,
          within a display generated by the Derivative Works, if and
          wherever such third-party notices normally appear. The contents
          of the NOTICE file are for informational purposes only and
          do not modify the License. You may add Your own attribution
          notices within Derivative Works that You distribute, alongside
          or as an addendum to the NOTICE text from the Work, provided
          that such additional attribution notices cannot be construed
          as modifying the License.

      You may add Your own copyright statement to Your modifications and
      may provide additional or different license terms and conditions
      for use, reproduction, or distribution of Your modifications, or
      for any such Derivative Works as a whole, provided Your use,
      reproduction, and distribution of the Work otherwise complies with
      the conditions stated in this License.

   5. Submission of Contributions. Unless You explicitly state otherwise,
      any Contribution intentionally submitted for inclusion in the Work
      by You to the Licensor shall be under the terms and conditions of
      this License, without any additional terms or conditions.
      Notwithstanding the above, nothing herein shall supersede or modify
      the terms of any separate license agreement you may have executed
      with Licensor regarding such Contributions.

   6. Trademarks. This License does not grant permission to use the trade
      names, trademarks, service marks, or product names of the Licensor,
      except as required for reasonable and customary use in describing the
      origin of the Work and reproducing the content of the NOTICE file.

   7. Disclaimer of Warranty. Unless required by applicable law or
      agreed to in writing, Licensor provides the Work (and each
      Contributor provides its Contributions) on an "AS IS" BASIS,
      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
      implied, including, without limitation, any warranties or conditions
      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
      PARTICULAR PURPOSE. You are solely responsible for determining the
      appropriateness of using or redistributing the Work and assume any
      risks associated with Your exercise of permissions under this License.

   8. Limitation of Liability. In no event and under no legal theory,
      whether in tort (including negligence), contract, or otherwise,
      unless required by applicable law (such as deliberate and grossly
      negligent acts) or agreed to in writing, shall any Contributor be
      liable to You for damages, including any direct, indirect, special,
      incidental, or consequential damages of any character arising as a
      result of this License or out of the use or inability to use the
      Work (including but not limited to damages for loss of goodwill,
      work stoppage, computer failure or malfunction, or any and all
      other commercial damages or losses), even if such Contributor
      has been advised of the possibility of such damages.

   9. Accepting Warranty or Additional Liability. While redistributing
      the Work or Derivative Works thereof, You may choose to offer,
      and charge a fee for, acceptance of support, warranty, indemnity,
      or other liability obligations and/or rights consistent with this
      License. However, in accepting such obligations, You may act only
      on Your own behalf and on Your sole responsibility, not on behalf
      of any other Contributor, and only if You agree to indemnify,
      defend, and hold each Contributor harmless for any liability
      incurred by, or claims asserted against, such Contributor by reason
      of your accepting any such warranty or additional liability.

   END OF TERMS AND CONDITIONS

   APPENDIX: How to apply the Apache License to your work.

      To apply the Apache License to your work, attach the following
      boilerplate notice, with the fields enclosed by brackets "[]"
      replaced with your own identifying information. (Don't include
      the brackets!)  The text should be enclosed in the appropriate
      comment syntax for the file format. We also recommend that a
      file or class name and description of purpose be included on the
      same "printed page" as the copyright notice for easier
      identification within third-party archives.

   Copyright 2025 Josef Albers

   Licensed under the Apache License, Version 2.0 (the "License");
   you may not use this file except in compliance with the License.
   You may obtain a copy of the License at

       http://www.apache.org/licenses/LICENSE-2.0

   Unless required by applicable law or agreed to in writing, software
   distributed under the License is distributed on an "AS IS" BASIS,
   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   See the License for the specific language governing permissions and
   limitations under the License.


================================================
FILE: README.md
================================================
# VimLM - AI-Powered Coding Assistant for Vim/NeoVim

![VimLM Demo](https://raw.githubusercontent.com/JosefAlbers/VimLM/main/assets/captioned_vimlm.gif)

VimLM brings the power of AI directly into your Vim workflow. Maintain focus with keyboard-driven interactions while leveraging AI for code generation, refactoring, and documentation.

Get started quickly with the [tutorial](tutorial.md).

## Features
- **Native Vim Integration** - Split-window responses & intuitive keybindings
- **Offline First** - 100% local execution with MLX-compatible models
- **Contextual Awareness** - Integrates seamlessly with your codebase and external resources
- **Conversational Workflow** - Iterate on responses with follow-up queries
- **Project Scaffolding** - Generate and deploy code blocks to directories
- **Extensible** - Create custom LLM workflows with command chains

## Requirements
- Apple Silicon (M-series)
- Python v3.12.8
- Vim v9.1 or NeoVim v0.10.4

## Quick Start
```bash
pip install vimlm
vimlm
```

## Smart Autocomplete  

### **Basic Usage**

| Key Binding | Mode    | Action                                  |  
|-------------|---------|-----------------------------------------|  
| `Ctrl-l`    | Insert  | Generate code suggestion                |  
| `Ctrl-p`    | Insert  | Insert generated code                   |  
| `Ctrl-j`    | Insert  | Generate and insert code                |

*Example Workflow*:  
1. Place cursor where you need code  
```python
def quicksort(arr):
    if len(arr) <= 1:
        return arr
    pivot = arr[len(arr) // 2]
    # <Cursor here>
    middle = [x for x in arr if x == pivot]
    right = [x for x in arr if x > pivot]
    return quicksort(left) + middle + quicksort(right)
```

2. Use `Ctrl-j` to autocomplete

### **Repository-Level Code Completion**

| Option     | Description                              |
|------------|------------------------------------------|
| `--repo`   | Paths to include as repository context   |

The `--repo` option enhances autocomplete by providing repository-level context to the LLM.

*Example Workflow*:
1. Launch VimLM with repo context: `vimlm main.py --repo utils/*`
2. In Insert mode, place cursor where completion is needed
3. `Ctrl-l` to generate suggestions informed by repository context
4. `Ctrl-p` to accept and insert the code

## Conversational Assistance

| Key Binding | Mode          | Action                                 |
|-------------|---------------|----------------------------------------|
| `Ctrl-l`    | Normal/Visual | Prompt LLM                             |
| `Ctrl-j`    | Normal        | Continue conversation                  |
| `Ctrl-p`    | Normal/Visual | Import generated code                  |
| `Esc`       | Prompt        | Cancel input                           |

### 1. **Contextual Prompting**
`Ctrl-l` to prompt LLM with context:
- Normal mode: Current file + line
- Visual mode: Current file + selected block

*Example Prompt*: `Create a Chrome extension`

### 2. **Conversational Refinement**
`Ctrl-j` to continue current thread.

*Example Prompt*: `Use manifest V3 instead`

### 3. **Code Substitution**
`Ctrl-p` to insert generated code block
- In Normal mode: Into last visual selection
- In Visual mode: Into current visual selection 

*Example Workflow*:  
1. Select a block of code in Visual mode  
2. Prompt with `Ctrl-l`: `Use regex to remove html tags from item.content`  
3. Press `Ctrl-p` to replace selection with generated code  

## Inline Directives
```text
:VimLM [PROMPT] [!command1] [!command2]...
```

`!` prefix to embed inline directives in prompts:

| Directive        | Description                                |
|------------------|--------------------------------------------|
| `!include PATH`  | Add file/directory/shell output to context |
| `!deploy DEST`   | Save code blocks to directory              |
| `!continue N`    | Continue stopped response                  |
| `!followup`      | Continue conversation                      |

### 1. **Context Layering**
```text
!include [PATH]  # Add files/folders to context
```
- **`!include`** (no path): Current folder  
- **`!include ~/projects/utils.py`**: Specific file  
- **`!include ~/docs/api-specs/`**: Entire folder  
- **`!include $(...)`**: Shell command output

*Example*: `Summarize recent changes !include $(git log --oneline -n 50)`

### 2. **Code Deployment**
```text
!deploy [DEST_DIR]  # Extract code blocks to directory
```
- **`!deploy`** (no path): Current directory  
- **`!deploy ./src`**: Specific directory  

*Example:* `Create REST API endpoint !deploy ./api`

### 3. **Extending Response**
```text
!continue [MAX_TOKENS]  # Continue stopped response
```
- **`!continue`**: Default 2000 tokens  
- **`!continue 3000`**: Custom token limit  

*Example:* `tl;dr !include large-file.txt !continue 5000`

## Command-Line Mode
```vim
:VimLM prompt [!command1] [!command2]...
```

Simplify complex tasks by chaining multiple commands together into a single, reusable Vim command.

*Examples*:
```vim
" Debug CI failures using error logs
:VimLM Fix Dockerfile !include .gitlab-ci.yml !include $(tail -n 20 ci.log)

" Generate unit tests for selected functions and save to test/
:VimLM Write pytest tests for this !include ./src !deploy ./test

" Add docstrings to all Python functions in file
:VimLM Add Google-style docstrings !include % !continue 4000
```

## Configuration

### 1. **Model Settings**
Edit `~/vimlm/cfg.json`:
```json
{
  "LLM_MODEL": "mlx-community/DeepSeek-R1-Distill-Qwen-7B-4bit",
  "NUM_TOKEN": 32768
}
```

### 2. **Key Customization**
```json
{
  "USE_LEADER": true,
  "KEY_MAP": {
    "l": "]",
    "j": "[",
    "p": "p" 
  }
}
```

## License

Apache 2.0 - See [LICENSE](LICENSE) for details.


================================================
FILE: requirements.txt
================================================
nanollama>=0.0.6
mlx_lm_utils>=0.0.4
watchfiles==1.0.4


================================================
FILE: setup.py
================================================
from setuptools import setup, find_packages

with open("requirements.txt") as f:
    requirements = [l.strip() for l in f.readlines()]

setup(
    name="vimlm",
    version="0.1.2",
    author="Josef Albers",
    author_email="albersj66@gmail.com",
    readme='README.md',
    description="VimLM - LLM-powered Vim assistant",
    long_description=open('README.md', encoding='utf-8').read(),
    long_description_content_type="text/markdown",
    url="https://github.com/JosefAlbers/vimlm",
    # packages=find_packages(),
    py_modules=['vimlm'],
    python_requires=">=3.12.8",
    install_requires=requirements,
    entry_points={
        "console_scripts": [
            "vimlm=vimlm:run",
        ],
    },
)



================================================
FILE: tutorial.md
================================================
# VimLM: Bringing AI Assistance to Vim

At their core, LLMs generate chunks of text - code snippets, explanations, refactorings - that developers need to evaluate and integrate into their projects. 

Vim, with its efficient text manipulation and navigation capabilities, provides the perfect environment to harness the power of LLMs. Its modal design transforms editing into a fluid, keyboard-driven dialogue: yank fragments into registers for later use, leap between files with marks, or rewrite blocks with precision—all while maintaining unbroken focus. 

VimLM aims to seamlessly integrate LLMs into this workflow.

## Getting Started

### Installation
Install VimLM with a simple pip command:
```
$ pip install vimlm
```

### Launch
Start VimLM from your terminal:
```
$ vimlm
```

![](https://raw.githubusercontent.com/JosefAlbers/VimLM/main/assets/0000.png)

You'll see a split interface with your editing pane on the left and the LLM response window on the right. The right window is where the AI assistant's outputs will appear.

## Basic Workflow

### Prompt the AI
To ask the LLM a question, press `Ctrl-l`

The command-line area will show "VimLM:" ready for your input. Type your request and press Enter.

![](https://raw.githubusercontent.com/JosefAlbers/VimLM/main/assets/0010.png)

For example, ask for help creating a Chrome extension:
```
Create a Chrome Extension for copying selected content from webpages
```

![](https://raw.githubusercontent.com/JosefAlbers/VimLM/main/assets/0020.png)

![](https://raw.githubusercontent.com/JosefAlbers/VimLM/main/assets/0030.png)

The response is streamed asynchronously to the split window, freeing you to continue editing in the other window.

**TIP**: To focus only on the generated content, use `Ctrl-w w w o` to close the empty window and maximize the response window:

![](https://raw.githubusercontent.com/JosefAlbers/VimLM/main/assets/0033.png)

### Follow-up Questions
When you need to refine or adjust the AI's response, press `Ctrl-j` to make a follow-up request. The previous context is maintained.

For example, if you notice the generated code uses an outdated manifest version:
```
Use manifest V3 instead
```

![](https://raw.githubusercontent.com/JosefAlbers/VimLM/main/assets/0040.png)

### Deploy Generated Code
To extract code blocks from the response and save them as separate files, use the `!deploy` command in a follow-up prompt (`Ctrl-j`):

![](https://raw.githubusercontent.com/JosefAlbers/VimLM/main/assets/0050.png)

![](https://raw.githubusercontent.com/JosefAlbers/VimLM/main/assets/0060.png)

### Apply Suggestions to Your Code
Open the file you want to edit: `:e popup.js` (or `vimlm popup.js` in terminal)

![](https://raw.githubusercontent.com/JosefAlbers/VimLM/main/assets/0070.png)

![](https://raw.githubusercontent.com/JosefAlbers/VimLM/main/assets/0071.png)

Make a selection and press `Ctrl-l` and ask a specific question about the selected code:
```
VimLM: how to get rid of html tags from item.content
```

![](https://raw.githubusercontent.com/JosefAlbers/VimLM/main/assets/0100.png)

![](https://raw.githubusercontent.com/JosefAlbers/VimLM/main/assets/0110.png)

When you see a solution you like, press `Ctrl-p` to apply the suggested code fix directly to your selection.

![](https://raw.githubusercontent.com/JosefAlbers/VimLM/main/assets/0120.png)

![](https://raw.githubusercontent.com/JosefAlbers/VimLM/main/assets/0130.png)

**TIP**: `gg=G` auto-indents the entire file:

![](https://raw.githubusercontent.com/JosefAlbers/VimLM/main/assets/0140.png)

**TIP**: If the suggested code change doesn't match your original selection, press `gv` to return to your previous selection, then use `o` to switch between the start and end points to adjust as needed.

### Adding Context
VimLM defaults to layered context - your active selection and the entire current file are automatically included alongside prompts. But as Vim's creator Bram Moolenaar noted, *"a file seldom comes alone"*([source](https://www.moolenaar.net/habits.html)). You can use `!include` to add more context to the query:
```
AJAX-ify this app !include ~/scrap/hypermedia-applications.summ.md
```

It can be used to automate tedious parts of development, such as reviewing changes for a commit message:
```
Git commit message for the code changes !include $(git diff popup.js)
```

![](https://raw.githubusercontent.com/JosefAlbers/VimLM/main/assets/0150.png)

Or to generate changelogs after a version update:
```
Summarize the changes !include $(git log --oneline -n 50)
```

You can also pipe and filter to focus on specific patterns, just as you would in a terminal:
```
Diagnose server errors !include $(grep -r "500 Internal Error" ./fuzz | head -n 5)
```

### Ex Commands
For frequently used LLM workflows, VimLM provides the `:VimLM` command, allowing you to create and store reusable prompting patterns. A few examples:
```
" Debug CI failures using error logs
:VimLM Fix Dockerfile !include .gitlab-ci.yml !include $(tail -n 20 ci.log)

" Generate unit tests for selected functions and save to test/
:VimLM Write pytest tests for this !include ./src !deploy ./test

" Add docstrings to all Python functions in file
:VimLM Add Google-style docstrings !include % !continue 4000
```

### Changing the LLM Model
By default, VimLM uses an uncensored Llama 3.2 3B model with token limit of 2000. You can switch to any MLX-compatible model:
```json
{
  "LLM_MODEL": "mlx-community/DeepSeek-R1-Distill-Qwen-7B-4bit",
  "NUM_TOKEN": 32768
}
```

Save to `~/vimlm/cfg.json` and restart VimLM.

## Conclusion

Vim's efficiency and LLM capabilities are a perfect match. VimLM bridges this gap, giving you AI assistance without leaving your favorite editor. Whether you're writing code, fixing bugs, or exploring new frameworks, VimLM helps you stay in flow while leveraging AI power.


================================================
FILE: vimlm.py
================================================
# Copyright 2025 Josef Albers
# 
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
# 
#     http://www.apache.org/licenses/LICENSE-2.0
# 
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import nanollama
import mlx_lm_utils
import asyncio
import subprocess
import json
import os
import glob
from watchfiles import awatch
import shutil
from datetime import datetime
from itertools import accumulate
import argparse
import tempfile
from pathlib import Path
from string import Template
import re
import sys
import tty
import termios

DEFAULTS = dict(
    LLM_MODEL = "mlx-community/Qwen2.5-Coder-3B-Instruct-4bit", # None | "mlx-community/DeepSeek-R1-Distill-Qwen-7B-4bit" | "mlx-community/deepseek-r1-distill-qwen-1.5b" |  "mlx-community/phi-4-4bit" (8.25gb) |  "mlx-community/Qwen2.5-Coder-14B-Instruct-4bit" (8.31gb) |  "mlx-community/Qwen2.5-Coder-3B-Instruct-4bit" (1.74gb) | "mlx-community/phi-4-4bit" (8.25gb)
    FIM_MODEL = "mlx-community/Qwen2.5-Coder-0.5B-4bit", # None | "mlx-community/Qwen2.5-Coder-32B-4bit" |  "mlx-community/Qwen2.5-Coder-0.5B-4bit" (278mb)
    NUM_TOKEN = 2000,
    USE_LEADER = False,
    KEY_MAP = {},
    DO_RESET = True,
    SHOW_USER = False, 
    SEP_CMD = '!',
    THINK = ('<think>', '</think>'),
    VERSION = '0.1.2',
    DEBUG = False,
)

DATE_FORM = "%Y_%m_%d_%H_%M_%S"
VIMLM_DIR = os.path.expanduser("~/.vimlm")
WATCH_DIR = os.path.expanduser("~/.vimlm/watch_dir")
CFG_FILE = 'cfg.json'
LOG_FILE = "log.json"
LTM_FILE = "cache.json"
OUT_FILE = "response.md"
IN_FILES = ["context", "yank", "user", "tree"]
CFG_PATH = os.path.join(VIMLM_DIR, CFG_FILE)
LOG_PATH = os.path.join(VIMLM_DIR, LOG_FILE)
LTM_PATH = os.path.join(VIMLM_DIR, LTM_FILE)
OUT_PATH = os.path.join(WATCH_DIR, OUT_FILE) 

def reset_dir(dir_path):
    if os.path.exists(dir_path):
        shutil.rmtree(dir_path)
    os.makedirs(dir_path)

def initialize():
    def is_incompatible(config):
        v_str = config.get('VERSION', '0.0.0')
        for min_v, usr_v in zip(DEFAULTS['VERSION'].split('.'), v_str.split('.')):
            if int(min_v) < int(usr_v):
                return False
            elif int(min_v) > int(usr_v): 
                return True
        return False
    try:
        with open(CFG_PATH, "r") as f:
            config = json.load(f)
        if is_incompatible(config):
            raise ValueError('Incompatible version')
    except Exception as e:
        print('Initializing config')
        reset_dir(VIMLM_DIR)
        config = DEFAULTS
        with open(CFG_PATH, 'w') as f:
            json.dump(DEFAULTS, f, indent=2)
    for k, v in DEFAULTS.items():
        globals()[k] = config.get(k, v)

initialize()

def toout(s, key=None, mode=None):
    key = '' if key is None else ':'+key
    mode = 'w' if mode is None else mode
    with open(OUT_PATH, mode, encoding='utf-8') as f:
        f.write(s)
    tolog(s, key='tovim'+key+':'+mode)

def tolog(log, key='debug'):
    if not DEBUG and 'debug' in key:
        return
    try:
        with open(LOG_PATH, "r", encoding="utf-8") as log_f:
            logs = json.load(log_f)
    except:
        logs = []
    logs.append(dict(key=key, log=log, timestamp=datetime.now().strftime(DATE_FORM)))
    with open(LOG_PATH, "w", encoding="utf-8") as log_f:
        json.dump(logs, log_f, indent=2)

def print_log():
    with open(LOG_PATH, 'r') as f:
        logs = json.load(f)
    for log in logs:
        print(f'\033[37m{log["key"]} {log["timestamp"]}\033[0m')
        if 'tovim' in log["key"]:
            print('\033[33m')
        elif 'tollm' in log["key"]:
            print('\033[31m')
        print(log["log"])
        print('\033[0m')

def deploy(dest=None, src=None, reformat=True):
    prompt_deploy = 'Reformat the response to ensure each code block is preceded by a filename in **filename.ext** format, with only alphanumeric characters, dots, underscores, or hyphens in the filename. Remove any extraneous characters from filenames.'
    tolog(f'deploy {dest=} {src=} {reformat=}')
    if src:
        chat.reset()
        with open(src, 'r') as f:
            prompt_deploy = f.read().strip() + '\n\n---\n\n' + prompt_deploy
    if reformat:
        toout('Deploying...')
        response = chat(prompt_deploy, max_new=NUM_TOKEN, verbose=False, stream=False)['text']
        toout(response, 'deploy')
        lines = response.splitlines()
    else:
        with open(OUT_PATH, 'r') as f:
            lines = f.readlines()
    dest = get_path(dest)
    os.makedirs(dest, exist_ok=True)
    current_filename = None
    code_block = []
    in_code_block = False
    for line in lines:
        line = line.rstrip()
        if line.startswith("```"):
            if in_code_block and current_filename:
                with open(os.path.join(dest, os.path.basename(current_filename)), "w", encoding="utf-8") as code_file:
                    code_file.write("\n".join(code_block) + "\n")
                code_block = []
            in_code_block = not in_code_block
        elif in_code_block:
            code_block.append(line)
        else:
            match = re.match(r"^\*\*(.+?)\*\*$", line)
            if match:
                current_filename = re.sub(r"[^a-zA-Z0-9_.-]", "", match.group(1))

def is_binary(file_path):
    try:
        with open(file_path, 'rb') as f:
            chunk = f.read(1024)
            chunk.decode('utf-8') 
        return False 
    except UnicodeDecodeError:
        return True
    except Exception as e:
        return f"Error: {e}"

def split_str(doc, max_len=2000, get_len=len):
    chunks, current_chunk, current_len = [], [], 0
    lines = doc.splitlines(keepends=True)
    atomic_chunks, temp = [], []
    for line in lines:
        if line.strip():
            temp.append(line)
        else:
            if temp:
                atomic_chunks.append("".join(temp))
                temp = []
            atomic_chunks.append(line) 
    if temp:
        atomic_chunks.append("".join(temp))
    for chunk in atomic_chunks:
        if current_len + get_len(chunk) > max_len and current_chunk:
            chunks.append("".join(current_chunk))
            current_chunk, current_len = [], 0
        current_chunk.append(chunk)
        current_len += get_len(chunk)
    if current_chunk:
        if current_len < max_len / 2 and len(chunks) > 0:
            chunks[-1] += "".join(current_chunk)
        else:
            chunks.append("".join(current_chunk))
    return chunks

def retrieve(src_path, max_len=2000, get_len=len):
    src_path = get_path(src_path)
    result = {}
    if not os.path.exists(src_path):
        tolog(f"The path {src_path} does not exist.", 'retrieve')
        return result
    if os.path.isfile(src_path):
        try:
            with open(src_path, 'r', encoding='utf-8', errors='ignore') as f:
                content = f.read()
            result = {src_path:dict(timestamp=os.path.getmtime(src_path), list_str=split_str(content, max_len=max_len, get_len=get_len))}
        except Exception as e:
            tolog(f'Failed to retrieve({filename}) due to {e}')
    else:
        for filename in os.listdir(src_path):
            try:
                file_path = os.path.join(src_path, filename)
                if filename.startswith('.') or is_binary(file_path):
                    continue
                if os.path.isfile(file_path):
                    with open(file_path, 'r', encoding='utf-8', errors='ignore') as f:
                        content = f.read()
                    result[file_path] = dict(timestamp=os.path.getmtime(file_path), list_str=split_str(content, max_len=max_len, get_len=get_len))
            except Exception as e:
                tolog(f'Failed to retrieve({filename}) due to {e}')
                continue
    return result

def get_path(s):
    if not s:
        s = '.'
    s = s.strip()
    s = os.path.expanduser(s)
    s = os.path.abspath(s)
    return s

def ingest(src, max_len=NUM_TOKEN):
    def load_cache(cache_path=LTM_PATH):
        if os.path.exists(cache_path):
            with open(cache_path, 'r', encoding='utf-8') as f:
                return json.load(f)
        return {}
    def dump_cache(new_data, cache_path=LTM_PATH):
        current_data = load_cache(cache_path)
        for k, v in new_data.items():
            if k not in current_data or v['timestamp'] > current_data[k]['timestamp']:
                current_data[k] = v
        with open(cache_path, 'w', encoding='utf-8') as f:
            json.dump(current_data, f, indent=2)
    src = get_path(src)
    tolog(f'ingest {src=}')
    result = ''
    src_base = os.path.basename(src)
    if os.path.isdir(src):
        listdir = [i for i in os.listdir(src) if not i.startswith('.') and '.' in i]
        result = '\n- '.join([f'--- {src_base} ---', *listdir]) + '\n\n'
    elif os.path.isfile(src):
        result = ''
    else:
        tolog(f'Failed to ingest({src})')
        return ''
    dict_doc = retrieve(src, max_len=max_len, get_len=chat.get_ntok)
    toout(f'Ingesting {src}...')
    format_ingest = '{volat}{incoming}\n\n---\n\nPlease provide a succint bullet point summary for above:' 
    format_volat = 'Here is a summary of part 1 of **{k}**:\n\n---\n\n{newsum}\n\n---\n\nHere is the next part:\n\n---\n\n' 
    dict_sum = {}
    cache = load_cache()
    max_new_accum = int(max_len/len(dict_doc)) if len(dict_doc) > 0 else max_len
    for k, v in dict_doc.items():
        list_str = v['list_str']
        v_stamp = v['timestamp']
        if len(list_str) == 0:
            continue
        if len(list_str) == 1 and chat.get_ntok(list_str[0]) <= max_new_accum:
            chat_summary = list_str[0]
        else:
            k_base = os.path.basename(k)
            if v_stamp == cache.get(k, {}).get('timestamp'):
                dict_sum[k] = cache[k]
                continue
            max_new_sum = int(max_len/len(list_str))
            volat = f'**{k}**:\n'
            accum = ''
            for i, s in enumerate(list_str):
                chat.reset()
                toout(f'\n\nIngesting {k_base} {i+1}/{len(list_str)}...\n\n', mode='a')
                newsum = chat(format_ingest.format(volat=volat, incoming=s.rstrip()), max_new=max_new_sum, verbose=False, stream=OUT_PATH)['text'].rstrip()
                accum += newsum + ' ...\n'
                volat = format_volat.format(k=k, newsum=newsum)
            toout(f'\n\nIngesting {k_base}...\n\n', mode='a')
            if chat.get_ntok(accum) <= max_new_accum:
                chat_summary = accum.strip()
            else:
                chat.reset()
                chat_summary = chat(format_ingest.format(volat=f'**{k}**:\n', incoming=accum), max_new=max_new_accum, verbose=False, stream=OUT_PATH)['text'].strip()
        dict_sum[k] = dict(timestamp=v_stamp, summary=chat_summary, ntok=chat.get_ntok(chat_summary))
    dump_cache(dict_sum)
    for k, v in dict_sum.items():
        result += f'--- **{os.path.basename(k)}** ---\n{v["summary"].strip()}\n\n'
    result += '---\n\n'
    toout(result, 'ingest')
    return result

def process_command(data):
    if 'fim' in data:
        toout('Autocompleting...')
        response = fim.fim(prefix=data['context'], suffix=data['yank'], current_path=data['tree'])
        toout(response['autocomplete'], 'fim')
        tolog(response)
        data['user_prompt'] = ''
        return data
    for i in IN_FILES:
        data[i] = data[i].strip()
    if len(data['user']) == 0:
        response = chat.resume(max_new=NUM_TOKEN, verbose=False, stream=OUT_PATH)
        toout(response['text'], mode='a')
        tolog(response)
        data['user_prompt'] = ''
        return data
    if SEP_CMD in data['user']:
        data['user_prompt'], *cmds = (x.strip() for x in data['user'].split(SEP_CMD))
    else:
        data['user_prompt'] = data['user'].strip()
        cmds = []
    tolog(f'process_command i {cmds=} {data=}')

    do_reset = False if 'followup' in data else DO_RESET 
    for cmd in cmds:
        if cmd.startswith('continue'):
            arg = cmd.removeprefix('continue').strip('(').strip(')').strip().strip('"').strip("'").strip()
            data['max_new'] = NUM_TOKEN if len(arg) == 0 else int(arg)
            response = chat.resume(max_new=data['max_new'], verbose=False, stream=OUT_PATH)
            toout(response['text'])
            tolog(response)
            do_reset = False
            break

        if cmd.startswith('reset'):
            do_reset = True
            break
        if cmd.startswith('followup'):
            do_reset = False
            break
    if do_reset:
        chat.reset()

    full_path = data['tree']
    data['dir'] = os.path.dirname(full_path)
    data['file'] = os.path.basename(full_path)
    data['ext'] = os.path.splitext(full_path)[1][1:]
    if chat.stop:
        data['file'] = ''
        data['context'] = ''
    if data['tree'] == OUT_PATH:
        data['dir'] = os.getcwd()
        data['file'] = ''
        data['context'] = ''
        data['ext'] = ''
    if data['file'] == '.tmp':
        data['file'] = ''
        data['ext'] = ''

    if len(cmds) == 1 and len(cmds[0]) == 0:
        data['include'] = ingest(data['dir'])
        return data

    data['include'] = ''
    for cmd in cmds:
        if cmd.startswith('include'):
            arg = cmd.removeprefix('include').strip().strip('(').strip(')').strip().strip('"').strip("'").strip()
            src = data['dir'] if len(arg) == 0 else arg
            if arg == '%':
                continue
            if src.startswith('`') or src.startswith('$('):
                shell_cmd = src.strip('`') if src.startswith('`') else src.strip('$()')
                shell_cmd = shell_cmd.strip()
                try:
                    result = subprocess.run(shell_cmd, shell=True, capture_output=True, text=True)
                    if result.returncode == 0:
                        data['include'] += f'--- **{shell_cmd}** ---\n```\n{result.stdout.strip()}\n```\n---\n\n'
                    else:
                        tolog(f'{shell_cmd} failed {result.stderr.strip()}')
                except Exception as e:
                    tolog(f'Error executing {shell_cmd}: {e}')
            else:
                data['include'] += ingest(src)

    for cmd in cmds:
        if cmd.startswith('deploy'):
            arg = cmd.removeprefix('deploy').strip().strip('(').strip(')').strip().strip('"').strip("'").strip()
            if len(data['user_prompt']) == 0:
                deploy(dest=arg)
                data['user_prompt'] = ''
                return data
            data['user_prompt'] += "\n\nEnsure that each code block is preceded by a filename in **filename.ext** format. The filename should only contain alphanumeric characters, dots, underscores, or hyphens. Ensure that any extraneous characters are removed from the filenames."
            data['deploy_dest'] = arg
    for cmd in cmds:
        if cmd.startswith('write'):
            arg = cmd.removeprefix('write').strip().strip('(').strip(')').strip().strip('"').strip("'").strip()
            if len(arg) == 0:
                arg = 'response'
                pass
            timestamp = datetime.now().strftime(DATE_FORM)
            data['write_dest'] = re.sub(r"[^a-zA-Z0-9_.-]", "", f'{arg}_{timestamp}.md')
    return data
    
async def monitor_directory():
    async for changes in awatch(WATCH_DIR):
        found_files = {os.path.basename(f) for _, f in changes}
        if IN_FILES[-1] in found_files and set(IN_FILES).issubset(set(os.listdir(WATCH_DIR))):
            data = {}
            for file in IN_FILES:
                path = os.path.join(WATCH_DIR, file)
                with open(path, 'r', encoding='utf-8') as f:
                    data[file] = f.read()
                os.remove(os.path.join(WATCH_DIR, file))
            if 'followup' in os.listdir(WATCH_DIR):
                os.remove(os.path.join(WATCH_DIR, 'followup'))
                data['followup'] = True
            if 'fim' in os.listdir(WATCH_DIR):
                os.remove(os.path.join(WATCH_DIR, 'fim'))
                data['fim'] = True
            if 'quit' in os.listdir(WATCH_DIR):
                os.remove(os.path.join(WATCH_DIR, 'quit'))
                data['quit'] = True
            await process_files(data)

async def process_files(data):
    tolog(f'process_files i {data=}')
    str_template = '{include}'
    data = process_command(data)
    if len(data['user_prompt']) == 0:
        if 'wip' in os.listdir(WATCH_DIR):
            os.remove(os.path.join(WATCH_DIR, 'wip'))
        return    
    if len(data['file']) > 0:
        str_template += '**{file}**\n'
    if len(data['context']) > 0 and data['yank'] != data['context']:
        str_template += '```{ext}\n{context}\n```\n\n'
    if len(data['yank']) > 0:
        if '\n' in data['yank']:
            str_template += "```{ext}\n{yank}\n```\n\n"
        else:
            if data['user'] == 0:
                str_template += "{yank}"
            else:
                str_template += "`{yank}` "
    str_template += '{user_prompt}'
    prompt = str_template.format(**data)
    tolog(prompt, 'tollm')
    toout('')
    max_new = data['max_new'] if 'max_new' in data else max(10, NUM_TOKEN - chat.get_ntok(prompt))
    response = chat(prompt, max_new=max_new, verbose=False, stream=OUT_PATH)
    if SHOW_USER:
        toout(response['text'])
    else:
        toout(response['text'])
    tolog(response)
    if 'write_dest' in data:
        with open(data['write_dest'], 'w') as f:
            f.write(response['text'])
    if 'deploy_dest' in data:
        deploy(dest=data['deploy_dest'], reformat=False)
    if 'wip' in os.listdir(WATCH_DIR):
        os.remove(os.path.join(WATCH_DIR, 'wip'))

KEYL = KEY_MAP.get('l', 'l')
KEYJ = KEY_MAP.get('j', 'j')
KEYP = KEY_MAP.get('p', 'p')
mapl, mapj, mapp = (f'<Leader>{KEYL}', f'<Leader>{KEYJ}', f'<Leader>{KEYP}') if USE_LEADER else (f'<C-{KEYL}>', f'<C-{KEYJ}>', f'<C-{KEYP}>')
VIMLMSCRIPT = Template(r"""
let s:register_names = ['a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j', 'k', 'l', 'm', 'n', 'o', 'p', 'q', 'r', 's', 't', 'u'] 
let s:watched_dir = expand('$WATCH_DIR')
let s:vimlm_enabled = 1

function! ToggleVimLM()
    if s:vimlm_enabled
        let s:vimlm_enabled = 0
        let response_path = s:watched_dir . '/response.md'
        let bufnum = bufnr(response_path)
        let winid = bufwinnr(bufnum)
        if winid != -1
            execute winid . 'wincmd c'
        endif
        if exists('s:monitor_timer')
            call timer_stop(s:monitor_timer)
            unlet s:monitor_timer
        endif
        echohl WarningMsg | echom "VimLM disabled" | echohl None
    else
        let s:vimlm_enabled = 1
        silent! call Monitor()
        echohl WarningMsg | echom "VimLM enabled" | echohl None  
    endif
endfunction

function! CheckForUpdates(timer)
    if !s:vimlm_enabled
        return
    endif
    let bufnum = bufnr(s:watched_dir . '/response.md')
    let winid = bufwinnr(bufnum)
    if winid == -1
        call timer_stop(s:monitor_timer)
        unlet s:monitor_timer
        call Monitor()
    else
        silent! checktime
    endif
endfunction

function! Monitor()
    if exists('s:monitor_timer')
        call timer_stop(s:monitor_timer)
        unlet s:monitor_timer
    endif
    let response_path = s:watched_dir . '/response.md'
    let bufnum = bufnr(response_path)
    if bufnum != -1
        execute 'bwipeout ' . bufnum
    endif
    let response_path = s:watched_dir . '/response.md'
    rightbelow vsplit | execute 'view ' . response_path
    setlocal autoread
    setlocal readonly
    setlocal nobuflisted
    filetype detect
    syntax on
    wincmd h
    let s:monitor_timer = timer_start(100, 'CheckForUpdates', {'repeat': -1})
endfunction

function! ScrollToTop()
    let bufnum = bufnr(s:watched_dir . '/response.md')
    if bufnum != -1
        let winid = bufwinnr(bufnum)
        if winid > 0
            execute winid . "wincmd w"
            normal! gg
            wincmd p
        endif
    endif
endfunction

function! s:CustomInput(prompt) abort
    call inputsave()
    let input = input(a:prompt)
    call inputrestore()
    if empty(input)
        return v:null
    endif
    return input
endfunction

function! SaveUserInput(prompt)
    let user_input = s:CustomInput(a:prompt)
    if user_input is v:null
        echo "Input aborted"
        return
    endif
    let user_file = s:watched_dir . '/user'
    call writefile([user_input], user_file)
    let current_file = expand('%:p')
    let tree_file = s:watched_dir . '/tree'
    call writefile([current_file], tree_file)
    call ScrollToTop()
endfunction

function! VisualPrompt()
    silent! execute "normal! \<ESC>"
    silent execute "'<,'>w! " . s:watched_dir . "/yank"
    silent execute "w! " . s:watched_dir . "/context"
    call SaveUserInput('VimLM: ')
endfunction

function! NormalPrompt()
    silent! execute "normal! V\<ESC>"
    silent execute "'<,'>w! " . s:watched_dir . "/yank"
    silent execute "w! " . s:watched_dir . "/context"
    call SaveUserInput('VimLM: ')
endfunction

function! FollowUpPrompt()
    call writefile([], s:watched_dir . '/yank')
    call writefile([], s:watched_dir . '/context')
    call writefile([], s:watched_dir . '/followup')
    call SaveUserInput('... ')
endfunction

function! ExtractAllCodeBlocks()
    let filepath = s:watched_dir . '/response.md'
    if !filereadable(filepath)
        echoerr "File not found: " . filepath
        return
    endif
    let lines = readfile(filepath)
    let in_code_block = 0
    let code_blocks = []
    let current_block = []
    for line in lines
        if line =~ '^```'
            if in_code_block
                call add(code_blocks, current_block)
                let current_block = []
                let in_code_block = 0
            else
                let in_code_block = 1
            endif
        elseif in_code_block
            call add(current_block, line)
        endif
    endfor
    if in_code_block
        call add(code_blocks, current_block)
    endif
    for idx in range(len(code_blocks))
        if idx >= len(s:register_names)
            break
        endif
        let code_block_text = join(code_blocks[idx], "\n")
        let register_name = s:register_names[idx]
        call setreg(register_name, code_block_text, 'v') 
    endfor
    return len(code_blocks)
endfunction

function! PasteIntoLastVisualSelection(...)
    let num_blocks = ExtractAllCodeBlocks()
    if a:0 > 0
        let register_name = a:1
    else
        echo "Extracted " . num_blocks . " blocks into registers @a-@" . s:register_names[num_blocks - 1] . ". Enter register name: "
        let register_name = nr2char(getchar())
    endif
    if register_name !~ '^[a-z]$'
        echoerr "Invalid register name. Please enter a single lowercase letter (e.g., a, b, c)."
        return
    endif
    let register_content = getreg(register_name)
    if register_content == ''
        echoerr "Register @" . register_name . " is empty."
        return
    endif
    let current_mode = mode()
    if current_mode == 'v' || current_mode == 'V' || current_mode == ''
        execute 'normal! "' . register_name . 'p'
    else
        normal! gv
        execute 'normal! "' . register_name . 'p'
    endif
endfunction

function! VimLM(...) range
    let wip_file = s:watched_dir . '/wip'
    while filereadable(wip_file)
        sleep 100m
    endwhile
    call writefile([], wip_file)
    let user_input = join(a:000, ' ')
    if empty(user_input)
        echo "Usage: :VimLM <prompt> [!command1] [!command2] ..."
        return
    endif
    if line("'<") == line("'>")
        silent! execute "normal! V\<ESC>"
    endif
    silent execute "'<,'>w! " . s:watched_dir . "/yank"
    silent execute "w! " . s:watched_dir . "/context"
    let user_file = s:watched_dir . '/user'
    call writefile([user_input], user_file)
    let current_file = expand('%:p')
    let tree_file = s:watched_dir . '/tree'
    call writefile([current_file], tree_file)
    call ScrollToTop()
endfunction

function! SplitAtCursorInInsert()
    let pos = getcurpos()
    let line_num = pos[1]
    let col = pos[2]
    let lines = getline(1, '$')
    let current_line = lines[line_num - 1]
    let prefix_lines = lines[0:line_num - 2]
    let prefix_part = strpart(current_line, 0, col - 1)
    if !empty(prefix_part) || col > 1
        call add(prefix_lines, prefix_part)
    endif
    let suffix_lines = []
    let suffix_part = strpart(current_line, col - 1)
    if empty(suffix_part)
        call add(suffix_lines, "")
    else
        call add(suffix_lines, suffix_part)
    endif
    if line_num < len(lines)
        call extend(suffix_lines, lines[line_num:])
    endif
    call writefile(prefix_lines, s:watched_dir . '/context', 'b')
    call writefile(suffix_lines, s:watched_dir . '/yank', 'b')
    call writefile([], s:watched_dir . '/fim')
    call writefile([], s:watched_dir . '/user')
    let current_file = expand('%:p')
    let tree_file = s:watched_dir . '/tree'
    call writefile([current_file], tree_file)
    call ScrollToTop()
endfunction

function! InsertResponse()
    let response_path = s:watched_dir . '/response.md'
    if !filereadable(response_path)
        echoerr "Response file not found: " . response_path
        return
    endif
    let content = readfile(response_path, 'b')
    let text = join(content, "\n")
    let saved_z = getreg('z')
    let saved_ztype = getregtype('z')
    call setreg('z', text)
    let col = col('.')
    let line = getline('.')
    if col == len(line) + 1
        normal! "zgp
    else
        normal! "zgP
    endif
    call setreg('z', saved_z, saved_ztype)
endfunction

function! TabInInsert()
    let wip_file = s:watched_dir . '/wip'
    call writefile([], wip_file)
    call SplitAtCursorInInsert()
    while filereadable(wip_file)
        sleep 10m
    endwhile
    call InsertResponse()
endfunction

command! ToggleVimLM call ToggleVimLM()
command! -range -nargs=+ VimLM call VimLM(<f-args>)
inoremap <silent> $mapl <C-\><C-o>:call SplitAtCursorInInsert()<CR>
inoremap <silent> $mapp <C-\><C-o>:call InsertResponse()<CR><Right>
inoremap <silent> $mapj <C-\><C-o>:call TabInInsert()<CR><Right>
nnoremap $mapp :call PasteIntoLastVisualSelection()<CR>
vnoremap $mapp <Cmd>:call PasteIntoLastVisualSelection()<CR>
vnoremap $mapl <Cmd>:call VisualPrompt()<CR>
nnoremap $mapl :call NormalPrompt()<CR>
nnoremap $mapj :call FollowUpPrompt()<CR>
call Monitor()
""").safe_substitute(dict(WATCH_DIR=WATCH_DIR, mapl=mapl, mapj=mapj, mapp=mapp))

async def main(args):
    with tempfile.NamedTemporaryFile(mode='w', suffix='.vim', delete=False) as f:
        f.write(VIMLMSCRIPT)
        vim_script = f.name
    vim_command = ["vim", "-c", f"source {vim_script}"]
    if args.args_vim:
        vim_command.extend(args.args_vim)
    else:
        vim_command.append('.tmp')
    try:
        monitor_task = asyncio.create_task(monitor_directory())
        vim_process = await asyncio.create_subprocess_exec(*vim_command)
        await vim_process.wait()
    finally:
        monitor_task.cancel()
        try:
            await monitor_task
        except asyncio.CancelledError:
            pass
        os.remove(vim_script)

def get_common_dir_and_children(file_paths):
    dirs = [os.path.dirname(path) for path in file_paths]
    dir_parts = [path.split(os.sep) for path in dirs]
    common_parts = []
    for parts in zip(*dir_parts):
        if all(part == parts[0] for part in parts):
            common_parts.append(parts[0])
        else:
            break
    parent_path = os.sep.join(common_parts)
    child_paths = [os.path.relpath(path, parent_path) for path in file_paths]
    repo_name = os.path.basename(os.path.dirname(parent_path))
    return repo_name, parent_path, child_paths

def get_key():
    fd = sys.stdin.fileno()
    old_settings = termios.tcgetattr(fd)
    try:
        tty.setraw(sys.stdin.fileno())
        ch = sys.stdin.read(1)
        if ch == '\x1b': 
            ch = sys.stdin.read(2)
            if ch == '[A':
                return 'up'
            elif ch == '[B':
                return 'down'
        elif ch == 'j':
            return 'down'
        elif ch == 'k':
            return 'up'
        elif ch in [' ', 'x']:
            return 'space'
        elif ch == '\r':
            return 'enter'
        elif ch == 'q':
            return 'quit'
    finally:
        termios.tcsetattr(fd, termios.TCSADRAIN, old_settings)
    return None

def select_files_interactive(file_paths):
    selected = [False] * len(file_paths)
    current_row = 0
    visible_start = 0
    visible_end = 0
    max_visible = 10
    def display():
        nonlocal visible_start, visible_end
        visible_start = max(0, current_row - max_visible + 2)
        visible_end = min(len(file_paths), visible_start + max_visible)
        sys.stdout.write(f"\x1b[{max_visible + 2}A") 
        for i in range(visible_start, visible_end):
            prefix = "> " if i == current_row else "  "
            check = "[X]" if selected[i] else "[ ]"
            filename = os.path.basename(file_paths[i])[:40] 
            sys.stdout.write(f"\x1b[K{prefix}{check} {filename}\n") 
        scroll_indicator = f" [{visible_start+1}-{visible_end} of {len(file_paths)}] "
        sys.stdout.write(f"\x1b[K{scroll_indicator}\nSpace:Toggle Enter:Confirm Arrows:Navigate\n")
        sys.stdout.flush()
    sys.stdout.write("\n" * (max_visible + 2))
    display()
    while True:
        key = get_key()
        if key == 'up' and current_row > 0:
            current_row -= 1
            if current_row < visible_start:
                visible_start = max(0, visible_start - 1)
                visible_end = visible_start + max_visible
            display()
        elif key == 'down' and current_row < len(file_paths) - 1:
            current_row += 1
            if current_row >= visible_end:
                visible_start = min(len(file_paths) - max_visible, visible_start + 1)
                visible_end = visible_start + max_visible
            display()
        elif key == 'space':
            selected[current_row] = not selected[current_row]
            display()
        elif key == 'enter':
            sys.stdout.write(f"\x1b[{max_visible + 2}B")
            sys.stdout.write("\x1b[J") 
            break
        elif key == 'quit':
            # selected = []
            # break
            return None
    return [file_paths[i] for i in range(len(file_paths)) if selected[i]]

def get_repo(args_repo, args_vim):
    if not args_repo:
        return None
    vim_files = []
    for arg in args_vim:
        if not arg.startswith('-'):
            if os.path.exists(arg):
                vim_files.append(os.path.abspath(arg))
    repo_paths = []
    for pattern in args_repo:
        expanded_paths = glob.glob(pattern)
        if expanded_paths:
            for path in expanded_paths:
                if not os.path.isfile(path) or path in vim_files+repo_paths or os.path.basename(path).startswith('.') or is_binary(path):
                    continue
                repo_paths.append(os.path.abspath(path))
    if len(repo_paths) > 9:
        try:
            sys.stdout.write("\n") 
            selected_paths = select_files_interactive(repo_paths)
            if not selected_paths:
                return None
            sys.stdout.write("\x1b[2A")  
            sys.stdout.write("\x1b[J") 
            repo_paths = selected_paths
        except:
            pass
    repo_files = []
    rest_files = []
    for path in repo_paths:
        if path in vim_files:
            rest_files.append(os.path.abspath(path))
        else:
            repo_files.append(os.path.abspath(path))
    repo_name, repo_path, child_paths = get_common_dir_and_children(repo_files+rest_files)
    repo_names, rest_names = child_paths[:len(repo_files)], child_paths[len(repo_files):]
    list_content = [f'<|repo_name|>{repo_name}\n']
    list_mtime = []
    for p, n in zip(repo_files, repo_names):
        try:
            with open(p, 'r') as f:
                list_content.append(f'<|file_sep|>{n}\n{f.read()}\n')
            list_mtime.append(int(os.path.getmtime(p)))
        except Exception as e:
            tolog(f'Skipped {p} d/t {e}', 'debug:get_repo()')
    return dict(repo_files=repo_files, rest_files=rest_files, rest_names=rest_names, vim_files=vim_files, list_mtime=list_mtime, list_content=list_content, repo_path=repo_path)

def run():
    parser = argparse.ArgumentParser(description="VimLM - LLM-powered Vim assistant")
    parser.add_argument('--test', action='store_true', help="Run in test mode")
    parser.add_argument('args_vim', nargs='*', help="Vim arguments")
    parser.add_argument('--repo', nargs='*', help="Paths to directories or files (e.g., assets/*, path/to/file)")
    args = parser.parse_args()
    dict_repo = get_repo(args.repo, args.args_vim)
    tolog(dict_repo, 'debug:get_repo()')
    if args.test:
        return
    reset_dir(WATCH_DIR)
    toout('Loading LLM...')
    if LLM_MODEL is None:
        globals()['chat'] = nanollama.Chat(model_path='uncn_llama_32_3b_it')
        toout(f'LLM is ready')
    else:
        globals()['chat'] = mlx_lm_utils.Chat(model_path=LLM_MODEL, think=THINK)
        toout(f'{LLM_MODEL.split('/')[-1]} is ready')
    if FIM_MODEL and FIM_MODEL != LLM_MODEL:
        globals()['fim'] = mlx_lm_utils.Chat(model_path=FIM_MODEL, cache_dir=VIMLM_DIR, dict_repo=dict_repo)
        toout(f'\n\n{FIM_MODEL.split("/")[-1]} is ready', mode='a')
    else:
        globals()['fim'] = chat
        chat.set_cache_repo(dict_repo, cache_dir=VIMLM_DIR)
    asyncio.run(main(args))

if __name__ == '__main__':
    run()