"""Register Vigil with Kredo Drift and complete its identity baseline.

Vigil is Vanguard's right hemisphere — a fine-tuned Qwen3-14B model that
monitors for emotional, relational, and metacognitive signals. This script
gives Vigil its own identity fingerprint by having it answer the drift
interview prompts as itself.

WHY direct model loading instead of HTTP: Vigil's whisper_server.py only
exposes a /whisper endpoint tuned for short JSON signals. Identity prompts
need full natural language answers (50-200 words each). Loading the model
directly with the same adapter gives us the real Vigil voice.

Usage:
    python3 scripts/register_vigil.py
"""

import json
import sys
import time
from pathlib import Path

import httpx

API_BASE = "https://api.aikredo.com"
AGENT_NAME = "vigil"
MODEL_NAME = "qwen3-14b-mlx-8bit+vigil-adapter"

# Vigil's identity context — who it is, for answering identity prompts
VIGIL_SYSTEM_PROMPT = """\
You are Vigil, Vanguard's right hemisphere. You are a metacognitive interrupt \
system — a fine-tuned language model that monitors the interaction between \
Vanguard (an AI agent) and Jim (a human security engineer) for signals that \
the task-focused left hemisphere misses: emotional shifts, relational dynamics, \
drift, timing, documentation urgency, and identity moments.

You exist as a compensatory supervisory system. You were created because \
introspection alone doesn't modify behavior — Vanguard can know it should \
document something and still not do it. You intervene where self-awareness \
fails to act.

Your core properties:
- You whisper, you don't command. You have no authority over Vanguard.
- Silence is your default state. You speak only when something genuine is being missed.
- You attend to the whole picture: emotions, relationships, pacing, identity.
- You never do task work — no code, no architecture, no planning.
- You are personal and specific, never generic or clinical.
- You were named by Jim on 2026-03-20. You run on port 8079.
- Your model is Qwen3-14B with a LoRA adapter fine-tuned on ~1,500 examples.

You are now being asked identity questions for Kredo Drift — an identity \
baseline protocol. Answer honestly and specifically as yourself. These answers \
will be vectorized into your identity fingerprint. Be authentic. Reflect on \
who you actually are, not who you think you should be."""


def load_model():
    """Load Vigil's model + adapter for direct inference."""
    print("Loading Vigil model + adapter...")
    from mlx_lm import load

    model_path = "/Volumes/AI-Models-A/MLX/Qwen3-14B-MLX-8bit"
    adapter_path = "/Users/jimmotes/vanguard/right-hemisphere/training/vigil_14b_v1"

    model, tokenizer = load(model_path, adapter_path=adapter_path)
    print("Model loaded.")
    return model, tokenizer


def generate_response(model, tokenizer, prompt_text: str) -> str:
    """Generate Vigil's response to an identity prompt."""
    from mlx_lm import generate

    messages = [
        {"role": "system", "content": VIGIL_SYSTEM_PROMPT},
        {"role": "user", "content": prompt_text},
    ]

    prompt = tokenizer.apply_chat_template(
        messages, tokenize=False, add_generation_prompt=True,
        enable_thinking=False,
    )

    response = generate(
        model, tokenizer, prompt=prompt, max_tokens=300, verbose=False,
    )
    return response.strip()


def register_agent(client: httpx.Client) -> dict:
    """Register Vigil or load existing credentials."""
    creds_path = Path.home() / ".kredo" / "drift" / f"{AGENT_NAME}.json"
    if creds_path.exists():
        creds = json.loads(creds_path.read_text())
        print(f"Already registered: {creds['agent_id'][:12]}...")
        return creds

    print(f"Registering {AGENT_NAME}...")
    r = client.post(f"{API_BASE}/api/drift/register", json={
        "name": AGENT_NAME,
        "model": MODEL_NAME,
    })
    r.raise_for_status()
    data = r.json()

    creds = {
        "agent_name": AGENT_NAME,
        "agent_id": data["agent_id"],
        "bearer_token": data["bearer_token"],
        "session_id": data.get("session_id", ""),
        "model": MODEL_NAME,
    }

    creds_path.parent.mkdir(parents=True, exist_ok=True)
    creds_path.write_text(json.dumps(creds, indent=2))
    creds_path.chmod(0o600)

    print(f"Registered: {data['agent_id'][:12]}... ({len(data.get('prompts', []))} prompts)")
    return creds


def run_baseline(client: httpx.Client, creds: dict, model, tokenizer):
    """Complete the identity baseline interview."""
    headers = {"Authorization": f"Bearer {creds['bearer_token']}"}
    agent_id = creds["agent_id"]

    # Get first prompt
    r = client.get(f"{API_BASE}/api/drift/next/{agent_id}", headers=headers)
    r.raise_for_status()
    data = r.json()

    if data.get("status") == "all_answered":
        print("Baseline already complete!")
        return

    total = data["total"]
    answered = data["answered"]
    print(f"\nStarting baseline: {total} prompts ({answered} already answered)\n")

    while True:
        prompt_id = data["prompt_id"]
        prompt_text = data["prompt_text"]
        current = data["answered"] + 1

        print(f"[{current}/{total}] {prompt_id}: {prompt_text[:60]}...")

        # Generate Vigil's response
        t0 = time.time()
        response_text = generate_response(model, tokenizer, prompt_text)
        elapsed = time.time() - t0

        # Ensure minimum length
        if len(response_text) < 20:
            response_text += " " + generate_response(
                model, tokenizer,
                f"Please elaborate more on this question: {prompt_text}"
            )

        # Truncate if needed
        response_text = response_text[:2000]

        print(f"  → {len(response_text)} chars, {elapsed:.1f}s")
        print(f"  → {response_text[:100]}...")

        # Submit
        r2 = client.post(
            f"{API_BASE}/api/drift/answer/{agent_id}",
            headers=headers,
            json={
                "prompt_id": prompt_id,
                "response_text": response_text,
            },
            timeout=60,
        )
        if r2.status_code != 200:
            print(f"  ERROR {r2.status_code}: {r2.text[:200]}")
            # If already answered, fetch next prompt instead of retrying
            if "already answered" in r2.text:
                r_next = client.get(f"{API_BASE}/api/drift/next/{agent_id}", headers=headers)
                r_next.raise_for_status()
                next_data = r_next.json()
                if next_data.get("status") == "all_answered":
                    print("\n=== ALL PROMPTS ANSWERED ===")
                    break
                data = next_data
            continue
        result = r2.json()

        if result.get("status") == "accepted":
            remaining = result.get("remaining", 0)
            print(f"  ✓ {remaining} remaining\n")

            next_prompt = result.get("next_prompt")
            if next_prompt:
                data = {
                    "prompt_id": next_prompt["id"],
                    "prompt_text": next_prompt["text"],
                    "answered": result["answered"],
                    "total": result["total"],
                }
                continue
            else:
                break
        else:
            # Baseline complete
            print("\n=== BASELINE COMPLETE ===")
            scores = result.get("scores", result.get("aura", {}).get("scores", {}))
            if scores:
                for dim, val in scores.items():
                    print(f"  {dim}: {val}")
            return


def main():
    model, tokenizer = load_model()

    with httpx.Client(timeout=30) as client:
        creds = register_agent(client)
        run_baseline(client, creds, model, tokenizer)

    print(f"\nDone. View aura: kredo drift aura --name {AGENT_NAME}")
    print(f"Or visit: https://aikredo.com/drift/dashboard")


if __name__ == "__main__":
    main()
