import json, re, glob, random, pathlib

RAW_DIR   = pathlib.Path.home() / ".devstyle/raw"
OUT_DIR   = pathlib.Path("./data")
MAX_CHARS = 8_000          # ≈2k tokens; keeps training memory sane
SYSTEM    = ("You are a senior engineer who writes code in the exact "
             "style, structure, and conventions of this user's codebase.")

SECRET_RE = re.compile(
    r'(?i)(api[_-]?key|secret|token|password|authorization|aws_[a-z_]+)'
    r'\s*[:=]\s*["\047]?[A-Za-z0-9+/_\-\.]{8,}["\047]?')
IP_RE     = re.compile(r'\b(?:10|172\.(?:1[6-9]|2\d|3[01])|192\.168)(?:\.\d{1,3}){2,3}\b')
EMAIL_RE  = re.compile(r'[\w\.-]+@[\w\.-]+\.\w+')

def scrub(text):
    text = SECRET_RE.sub(r'\1=[REDACTED_SECRET]', text)
    text = IP_RE.sub('[REDACTED_IP]', text)
    return EMAIL_RE.sub('[REDACTED_EMAIL]', text)

def usable(ev):
    diff = ev.get("diff_output", "")
    if ev.get("status") != "success" or not diff.strip():
        return False
    if len(diff) > MAX_CHARS * 3:        # generated / vendored monster
        return False
    return True

def to_chat(ev):
    intent = scrub(ev.get("commit_message", "Update functionality").strip())
    code   = scrub(ev["diff_output"])[:MAX_CHARS]
    return {"messages": [
        {"role": "system",    "content": SYSTEM},
        {"role": "user",      "content": f"Implement the following change:\n{intent}"},
        {"role": "assistant", "content": code},
    ]}

events = [json.load(open(p)) for p in glob.glob(str(RAW_DIR / "*.json"))]
rows   = [to_chat(e) for e in events if usable(e)]

# De-duplicate on assistant content, then shuffle and split 90/10
seen, unique = set(), []
for r in rows:
    key = r["messages"][2]["content"][:400]
    if key not in seen:
        seen.add(key); unique.append(r)

random.seed(13); random.shuffle(unique)
cut = max(1, int(len(unique) * 0.1))
OUT_DIR.mkdir(exist_ok=True)
for name, chunk in [("valid", unique[:cut]), ("train", unique[cut:])]:
    with open(OUT_DIR / f"{name}.jsonl", "w") as f:
        for row in chunk:
            f.write(json.dumps(row) + "\n")

print(f"{len(unique)} examples → train/valid written to {OUT_DIR}/")

# Clean virtual environment
python3 -m venv .mlx_env
source .mlx_env/bin/activate

# mlx-lm with training dependencies, plus the HF hub for model downloads
pip install -U "mlx-lm[train]" huggingface_hub

# Optional but useful
pip install gitleaks-py wandb   # experiment logging: add --report-to wandb

# Base: 4-bit quantized ⇒ training runs as QLoRA automatically
model: "mlx-community/Qwen2.5-Coder-7B-Instruct-4bit"
train: true
data: "./data"              # contains train.jsonl (+ valid.jsonl)
adapter_path: "./devstyle-adapters"

fine_tune_type: lora        # lora | dora | full
num_layers: 16              # layers to adapt; lower ⇒ less memory
batch_size: 1
grad_accumulation_steps: 4  # effective batch of 4 at batch-1 memory
iters: 800
learning_rate: 1e-5
mask_prompt: true           # loss on assistant tokens only
steps_per_eval: 100         # validation-loss cadence
save_every: 200             # checkpoint cadence

mlx_lm.lora --config devstyle.yaml

# Interrupted? Resume from the last checkpoint:
mlx_lm.lora --config devstyle.yaml \
  --resume-adapter-file ./devstyle-adapters/adapters.safetensors

mlx_lm.lora --model mlx-community/Qwen2.5-Coder-7B-Instruct-4bit \
  --adapter-path ./devstyle-adapters --data ./data --test

# With your adapter
mlx_lm.generate --model mlx-community/Qwen2.5-Coder-7B-Instruct-4bit \
  --adapter-path ./devstyle-adapters \
  --prompt "Implement the following change:\nAdd pagination to the /users endpoint"

# Same prompt, no adapter — the control
mlx_lm.generate --model mlx-community/Qwen2.5-Coder-7B-Instruct-4bit \
  --prompt "Implement the following change:\nAdd pagination to the /users endpoint"

mlx_lm.fuse --model mlx-community/Qwen2.5-Coder-7B-Instruct-4bit \
  --adapter-path ./devstyle-adapters \
  --save-path ./devstyle-model

name: DevStyle Local
version: 1.0.0
schema: v1
models:
  # Route A/B — mlx_lm.server or LM Studio (OpenAI-compatible)
  - name: DevStyle 7B
    provider: openai
    model: devstyle-model
    apiBase: http://localhost:8080/v1
    roles: [chat, edit, apply]

  # Route C — Ollama
  - name: DevStyle 7B (Ollama)
    provider: ollama
    model: devstyle
    roles: [chat, edit]

  # Keep a small fast model for tab-autocomplete
  - name: Autocomplete
    provider: ollama
    model: qwen2.5-coder:1.5b
    roles: [autocomplete]