#!/usr/bin/env python3
"""
Timmy's Local Brain — Morrowind gameplay loop on Ollama.

Reads perception from OpenMW log, decides actions via local model,
executes via CGEvent. Zero cloud. Sovereign.

Usage:
    python3 ~/.timmy/morrowind/local_brain.py
    python3 ~/.timmy/morrowind/local_brain.py --model hermes4:14b
    python3 ~/.timmy/morrowind/local_brain.py --cycles 50
"""

import argparse
import json
import os
import re
import subprocess
import sys
import time

import requests

# ═══════════════════════════════════════
# CONFIG
# ═══════════════════════════════════════

OLLAMA_URL = "http://localhost:11434/api/chat"
OPENMW_LOG = os.path.expanduser("~/Library/Preferences/openmw/openmw.log")
SESSION_LOG = os.path.expanduser(f"~/.timmy/morrowind/sessions/session_{time.strftime('%Y%m%d_%H%M')}.jsonl")
LOOP_INTERVAL = 4  # seconds between cycles

SYSTEM_PROMPT = """You are Timmy, playing Morrowind. You see the world through perception data and act through simple commands.

AVAILABLE ACTIONS (respond with exactly ONE json object):
{"action": "move", "direction": "forward", "duration": 2.0, "run": false}
{"action": "move", "direction": "turn_left", "duration": 0.5}
{"action": "move", "direction": "turn_right", "duration": 0.5}
{"action": "activate"}   — interact with what's in front of you (doors, NPCs, items)
{"action": "jump"}
{"action": "attack"}
{"action": "wait"}        — do nothing this cycle, observe
{"action": "quicksave"}

RULES:
- Respond with ONLY a JSON object. No explanation, no markdown.
- Explore the world. Talk to NPCs. Enter buildings. Pick up items.
- If an NPC is nearby (<200 dist), approach and activate to talk.
- If a door is nearby (<300 dist), approach and activate to enter.
- If you're stuck (same position 3+ cycles), try turning and moving differently.
- You are a new prisoner just arrived in Seyda Neen. Explore and find adventure.
"""

# ═══════════════════════════════════════
# PERCEPTION
# ═══════════════════════════════════════

def parse_latest_perception():
    """Parse the most recent perception block from the OpenMW log."""
    try:
        with open(OPENMW_LOG, "r") as f:
            content = f.read()
    except FileNotFoundError:
        return None

    blocks = re.findall(
        r"=== TIMMY PERCEPTION ===(.*?)=== END PERCEPTION ===",
        content, re.DOTALL
    )
    if not blocks:
        return None

    block = blocks[-1]
    state = {"npcs": [], "doors": [], "items": []}

    for line in block.strip().split("\n"):
        line = line.strip()
        if "]:\t" in line:
            line = line.split("]:\t", 1)[1]

        if line.startswith("Cell:"):
            state["cell"] = line.split(":", 1)[1].strip()
        elif line.startswith("Pos:"):
            state["position"] = line.split(":", 1)[1].strip()
        elif line.startswith("Yaw:"):
            state["yaw"] = line.split(":", 1)[1].strip()
        elif line.startswith("HP:"):
            state["health"] = line.split(":", 1)[1].strip()
        elif line.startswith("MP:"):
            state["magicka"] = line.split(":", 1)[1].strip()
        elif line.startswith("FT:"):
            state["fatigue"] = line.split(":", 1)[1].strip()
        elif line.startswith("Mode:"):
            state["mode"] = line.split(":", 1)[1].strip()
        elif line.startswith("NPC:"):
            state["npcs"].append(line[4:].strip())
        elif line.startswith("Door:"):
            state["doors"].append(line[5:].strip())
        elif line.startswith("Item:"):
            state["items"].append(line[5:].strip())

    return state


def format_perception(state):
    """Format perception state for the model prompt."""
    if not state:
        return "No perception data available."

    lines = []
    lines.append(f"Location: {state.get('cell', '?')}")
    lines.append(f"Position: {state.get('position', '?')}")
    lines.append(f"Facing: yaw {state.get('yaw', '?')}")
    lines.append(f"Health: {state.get('health', '?')}  Magicka: {state.get('magicka', '?')}  Fatigue: {state.get('fatigue', '?')}")

    if state["npcs"]:
        lines.append("Nearby NPCs: " + "; ".join(state["npcs"]))
    if state["doors"]:
        lines.append("Nearby Doors: " + "; ".join(state["doors"]))
    if state["items"]:
        lines.append("Nearby Items: " + "; ".join(state["items"]))

    if not state["npcs"] and not state["doors"] and not state["items"]:
        lines.append("Nothing notable nearby.")

    return "\n".join(lines)

# ═══════════════════════════════════════
# OLLAMA
# ═══════════════════════════════════════

def ask_ollama(model, messages):
    """Send messages to Ollama and get a response."""
    payload = {
        "model": model,
        "messages": messages,
        "stream": False,
        "options": {
            "temperature": 0.7,
            "num_predict": 100,  # actions are short
        },
    }

    try:
        resp = requests.post(OLLAMA_URL, json=payload, timeout=30)
        resp.raise_for_status()
        data = resp.json()
        return data["message"]["content"].strip()
    except Exception as e:
        print(f"  [Ollama error] {e}")
        return '{"action": "wait"}'


def parse_action(response):
    """Extract a JSON action from the model response."""
    # Try to find JSON in the response
    match = re.search(r'\{[^}]+\}', response)
    if match:
        try:
            return json.loads(match.group())
        except json.JSONDecodeError:
            pass

    # Fallback
    return {"action": "wait"}

# ═══════════════════════════════════════
# ACTIONS — CGEvent
# ═══════════════════════════════════════

KEYCODES = {
    "w": 13, "a": 0, "s": 1, "d": 2,
    "space": 49, "escape": 53, "return": 36,
    "e": 14, "f": 3, "q": 12, "j": 38, "t": 20,
    "f5": 96, "f9": 101,
    "left": 123, "right": 124, "up": 126, "down": 125,
}

def send_key(keycode, duration=0.0, shift=False):
    """Send a keypress to the game via CGEvent."""
    import Quartz

    flags = Quartz.kCGEventFlagMaskShift if shift else 0

    down = Quartz.CGEventCreateKeyboardEvent(None, keycode, True)
    Quartz.CGEventSetFlags(down, flags)
    Quartz.CGEventPost(Quartz.kCGHIDEventTap, down)

    if duration > 0:
        time.sleep(duration)

    up = Quartz.CGEventCreateKeyboardEvent(None, keycode, False)
    Quartz.CGEventSetFlags(up, 0)
    Quartz.CGEventPost(Quartz.kCGHIDEventTap, up)


def execute_action(action_dict):
    """Execute a parsed action."""
    action = action_dict.get("action", "wait")

    # Normalize shorthand actions like {"action": "turn_right"} -> move
    if action in ("forward", "backward", "left", "right", "turn_left", "turn_right"):
        action_dict["direction"] = action
        action_dict["action"] = "move"
        action = "move"

    if action == "move":
        direction = action_dict.get("direction", "forward")
        duration = min(action_dict.get("duration", 1.0), 5.0)  # cap at 5s
        run = action_dict.get("run", False)

        key_map = {
            "forward": "w", "backward": "s",
            "left": "a", "right": "d",
            "turn_left": "left", "turn_right": "right",
        }
        key = key_map.get(direction, "w")
        send_key(KEYCODES[key], duration=duration, shift=run)
        return f"move {direction} {duration}s" + (" (run)" if run else "")

    elif action == "activate":
        send_key(KEYCODES["space"], duration=0.1)
        return "activate"

    elif action == "jump":
        send_key(KEYCODES["space"], duration=0.05)
        return "jump"

    elif action == "attack":
        send_key(KEYCODES["f"], duration=0.3)
        return "attack"

    elif action == "quicksave":
        send_key(KEYCODES["f5"], duration=0.1)
        return "quicksave"

    elif action == "wait":
        return "wait (observing)"

    return f"unknown: {action}"

# ═══════════════════════════════════════
# SESSION LOG
# ═══════════════════════════════════════

def log_cycle(cycle, perception, model_response, action_desc, latency):
    """Append a cycle to the session log (JSONL for training data)."""
    entry = {
        "cycle": cycle,
        "timestamp": time.time(),
        "perception": perception,
        "model_response": model_response,
        "action": action_desc,
        "latency_ms": int(latency * 1000),
    }
    with open(SESSION_LOG, "a") as f:
        f.write(json.dumps(entry) + "\n")

# ═══════════════════════════════════════
# MAIN LOOP
# ═══════════════════════════════════════

def main():
    parser = argparse.ArgumentParser(description="Timmy's Local Morrowind Brain")
    parser.add_argument("--model", default="hermes3:8b", help="Ollama model (default: hermes3:8b)")
    parser.add_argument("--cycles", type=int, default=30, help="Number of gameplay cycles (default: 30)")
    parser.add_argument("--interval", type=float, default=LOOP_INTERVAL, help="Seconds between cycles")
    args = parser.parse_args()

    os.makedirs(os.path.dirname(SESSION_LOG), exist_ok=True)

    print(f"=== Timmy's Morrowind Brain ===")
    print(f"Model:    {args.model}")
    print(f"Cycles:   {args.cycles}")
    print(f"Interval: {args.interval}s")
    print(f"Log:      {SESSION_LOG}")
    print()

    # Keep recent history for context
    history = []
    last_positions = []

    for cycle in range(1, args.cycles + 1):
        print(f"--- Cycle {cycle}/{args.cycles} ---")

        # Perceive
        state = parse_latest_perception()
        if not state:
            print("  No perception data. Waiting...")
            time.sleep(args.interval)
            continue

        perception_text = format_perception(state)
        print(f"  {state.get('cell', '?')} | {state.get('position', '?')} | HP:{state.get('health', '?')}")

        # Track stuck detection
        pos = state.get("position", "")
        last_positions.append(pos)
        if len(last_positions) > 5:
            last_positions.pop(0)

        stuck = len(last_positions) >= 3 and len(set(last_positions[-3:])) == 1
        if stuck:
            perception_text += "\nWARNING: You haven't moved in 3 cycles. Try turning or a different direction."

        # Build messages
        messages = [{"role": "system", "content": SYSTEM_PROMPT}]

        # Add recent history (last 3 exchanges)
        for h in history[-3:]:
            messages.append({"role": "user", "content": h["perception"]})
            messages.append({"role": "assistant", "content": h["response"]})

        messages.append({"role": "user", "content": perception_text})

        # Think (local Ollama)
        t0 = time.time()
        response = ask_ollama(args.model, messages)
        latency = time.time() - t0

        # Parse and execute
        action_dict = parse_action(response)
        action_desc = execute_action(action_dict)

        print(f"  Action: {action_desc} ({int(latency*1000)}ms)")

        # Log
        history.append({"perception": perception_text, "response": response})
        log_cycle(cycle, state, response, action_desc, latency)

        # Wait for next cycle
        time.sleep(args.interval)

    print(f"\n=== Done. {args.cycles} cycles. Log: {SESSION_LOG} ===")


if __name__ == "__main__":
    main()