#!/usr/bin/env python3
"""
codey — a standalone terminal client for LuisCodey (LuisForge's AI coding assistant).

Usage:
  codey                       pick a repo from a full-screen list, then chat
  codey <owner/repo>          jump straight into chatting about that repo
  codey [<owner/repo>] --plain  plain-text REPL/list instead (no curses, works over dumb pipes/logs)

TUI keys:
  Enter        send message
  /save [n]    commit the n-th code block from the last reply (or the first, if omitted)
  /clear       reset conversation history
  Ctrl+C, Esc  quit

Shares config (API base + saved token) with the `forge` CLI at ~/.forge/config.json.
Only uses the Python standard library — no extra deps to install.
"""
import os
import sys
import re
import json
import time
import locale
import curses
import textwrap
import urllib.request
import urllib.error

USER_AGENT = 'Mozilla/5.0 (compatible; codey-cli/1.1.0)'
DEFAULT_API_BASE = os.environ.get('FORGE_API_BASE', 'https://reviewer-shopper-gamma-bag.trycloudflare.com')
CONFIG_PATH = os.path.expanduser('~/.forge/config.json')

CODE_BLOCK_RE = re.compile(r'```(\w+)?[ \t]+(\S+)\r?\n(.*?)```', re.DOTALL)

LOGO = [
    "  ____ ___  ____  _______   __",
    " / ___/ _ \\|  _ \\| ____\\ \\ / /",
    "| |  | | | | | | |  _|  \\ V / ",
    "| |__| |_| | |_| | |___  | |  ",
    " \\____\\___/|____/|_____| |_|  ",
]


def load_config():
    if os.path.exists(CONFIG_PATH):
        with open(CONFIG_PATH) as f:
            return json.load(f)
    return {}


def api_base():
    return load_config().get('api_base') or DEFAULT_API_BASE


def token():
    return load_config().get('token')


def die(msg):
    print(f"error: {msg}", file=sys.stderr)
    sys.exit(1)


def api_json(method, path, body=None):
    url = api_base().rstrip('/') + path
    headers = {'User-Agent': USER_AGENT, 'Content-Type': 'application/json'}
    if token():
        headers['Authorization'] = f'Bearer {token()}'
    data = json.dumps(body).encode() if body is not None else None
    req = urllib.request.Request(url, data=data, headers=headers, method=method)
    try:
        with urllib.request.urlopen(req, timeout=310) as res:
            return res.status, json.loads(res.read())
    except urllib.error.HTTPError as e:
        content = e.read()
        try:
            return e.code, json.loads(content)
        except json.JSONDecodeError:
            return e.code, {'error': content.decode(errors='replace')}
    except urllib.error.URLError as e:
        die(f"couldn't reach {api_base()} ({e.reason})")


def api_stream(path, body):
    """Yields ('data', text) chunks as they arrive, or a single ('error', message)."""
    url = api_base().rstrip('/') + path
    headers = {'User-Agent': USER_AGENT, 'Content-Type': 'application/json'}
    if token():
        headers['Authorization'] = f'Bearer {token()}'
    req = urllib.request.Request(url, data=json.dumps(body).encode(), headers=headers, method='POST')
    try:
        resp = urllib.request.urlopen(req, timeout=310)
    except urllib.error.HTTPError as e:
        content = e.read()
        try:
            err = json.loads(content)
        except json.JSONDecodeError:
            err = {'error': content.decode(errors='replace')}
        yield ('error', err.get('error', 'unknown error'))
        return
    except urllib.error.URLError as e:
        yield ('error', f"couldn't reach {api_base()} ({e.reason})")
        return
    with resp:
        while True:
            chunk = resp.read(64)
            if not chunk:
                break
            yield ('data', chunk.decode(errors='replace'))


def api_upload(owner, name, filename, data, message):
    boundary = '----codeycli'
    parts = [
        f'--{boundary}\r\nContent-Disposition: form-data; name="path"\r\n\r\n\r\n'.encode(),
        f'--{boundary}\r\nContent-Disposition: form-data; name="message"\r\n\r\n{message}\r\n'.encode(),
        f'--{boundary}\r\nContent-Disposition: form-data; name="file"; filename="{filename}"\r\n'
        f'Content-Type: application/octet-stream\r\n\r\n'.encode() + data + b'\r\n',
        f'--{boundary}--\r\n'.encode(),
    ]
    body = b''.join(parts)
    url = api_base().rstrip('/') + f'/api/repos/{owner}/{name}/upload'
    req = urllib.request.Request(url, data=body, method='POST', headers={
        'Content-Type': f'multipart/form-data; boundary={boundary}',
        'Authorization': f'Bearer {token()}',
        'User-Agent': USER_AGENT,
    })
    try:
        with urllib.request.urlopen(req, timeout=30) as res:
            return 200, json.loads(res.read())
    except urllib.error.HTTPError as e:
        content = e.read()
        try:
            return e.code, json.loads(content)
        except json.JSONDecodeError:
            return e.code, {'error': content.decode(errors='replace')}


# ---------------- plain REPL (fallback / --plain) ----------------
def print_reply(reply):
    print(f"\033[35mCodey:\033[0m {reply}\n")


def run_plain(owner, name):
    print(f"\033[1mLuisCodey\033[0m — chatting about {owner}/{name}. Type /help for commands, /exit to quit.\n")
    history = []
    last_blocks = []
    while True:
        try:
            message = input('\033[36myou>\033[0m ').strip()
        except (EOFError, KeyboardInterrupt):
            print()
            break
        if not message:
            continue
        if message in ('/exit', '/quit'):
            break
        if message == '/help':
            print(__doc__)
            continue
        if message == '/clear':
            history, last_blocks = [], []
            print('(history cleared)\n')
            continue
        if message.startswith('/save'):
            if not last_blocks:
                print('(no code blocks in the last reply)\n')
                continue
            parts = message.split()
            idx = int(parts[1]) - 1 if len(parts) > 1 else 0
            if idx < 0 or idx >= len(last_blocks):
                print(f'(no block #{idx + 1} — last reply had {len(last_blocks)})\n')
                continue
            _, filename, code = last_blocks[idx]
            status, data = api_upload(owner, name, filename, code.encode(), f'{filename} via LuisCodey (codey CLI)')
            print(f"saved {filename} — commit {data['commit'][:8]}\n" if status == 200 else f"error: {data.get('error')}\n")
            continue

        history.append({'role': 'user', 'content': message})
        print('\033[35mCodey:\033[0m ', end='', flush=True)
        reply = ''
        had_error = False
        for kind, chunk in api_stream(f'/api/repos/{owner}/{name}/luiscodey/chat/stream',
                                       {'message': message, 'history': history[:-1]}):
            if kind == 'error':
                print(f"\nerror: {chunk}\n")
                history.pop()
                had_error = True
                break
            reply += chunk
            print(chunk, end='', flush=True)
        if had_error:
            continue
        print('\n')
        history.append({'role': 'assistant', 'content': reply})
        last_blocks = CODE_BLOCK_RE.findall(reply)
        if last_blocks:
            names = ', '.join(f'{i+1}) {b[1]}' for i, b in enumerate(last_blocks))
            print(f"(code blocks: {names} — use /save [n] to commit one)\n")


# ---------------- TUI ----------------
class Tui:
    def __init__(self, stdscr, owner, name):
        self.stdscr = stdscr
        self.owner = owner
        self.name = name
        self.history = []       # [{'role', 'content'}]
        self.log_lines = []     # rendered chat log, list of (attr, text)
        self.last_blocks = []
        self.input_buf = ''
        self.status = ''
        self.busy = False

        curses.curs_set(1)
        curses.start_color()
        curses.use_default_colors()
        curses.init_pair(1, curses.COLOR_MAGENTA, -1)   # Codey label / logo
        curses.init_pair(2, curses.COLOR_CYAN, -1)      # you label
        curses.init_pair(3, curses.COLOR_YELLOW, -1)    # tips / status
        curses.init_pair(4, curses.COLOR_WHITE, -1)     # dim
        curses.init_pair(5, curses.COLOR_RED, -1)       # errors
        curses.init_pair(6, curses.COLOR_GREEN, -1)     # saved/success
        self.stdscr.keypad(True)

    def size(self):
        h, w = self.stdscr.getmaxyx()
        return h, w

    def draw_idle(self):
        """OpenCode-style idle screen: big centered logo + input box + tips."""
        self.stdscr.clear()
        h, w = self.size()
        logo_w = max(len(l) for l in LOGO)
        start_y = max(1, h // 2 - 9)
        for i, line in enumerate(LOGO):
            x = max(0, (w - logo_w) // 2)
            try:
                self.stdscr.addstr(start_y + i, x, line, curses.color_pair(1) | curses.A_BOLD)
            except curses.error:
                pass

        box_y = start_y + len(LOGO) + 2
        self.draw_input_box(box_y, w)

        sub_y = box_y + 4
        sub = f"{self.owner}/{self.name}  ·  deepseek-v4-flash · DeepSeek"
        try:
            self.stdscr.addstr(sub_y, max(0, (w - len(sub)) // 2), sub, curses.color_pair(4))
        except curses.error:
            pass

        hint = "Enter send    /save [n] commit code    /clear reset    Esc/Ctrl+C quit"
        try:
            self.stdscr.addstr(h - 2, max(0, (w - len(hint)) // 2), hint, curses.color_pair(3))
        except curses.error:
            pass

        self.draw_status_bar(h, w)
        self.stdscr.refresh()

    def draw_input_box(self, y, w):
        box_w = min(70, w - 6)
        x = max(0, (w - box_w) // 2)
        placeholder = 'Ask anything... "What does this repo do?"'
        text = self.input_buf if self.input_buf else placeholder
        attr = curses.color_pair(4) if not self.input_buf else curses.A_NORMAL
        try:
            self.stdscr.addstr(y - 1, x, '╭' + '─' * (box_w - 2) + '╮')
            self.stdscr.addstr(y, x, '│' + ' ' * (box_w - 2) + '│')
            self.stdscr.addstr(y, x + 2, text[:box_w - 4], attr)
            self.stdscr.addstr(y + 1, x, '╰' + '─' * (box_w - 2) + '╯')
        except curses.error:
            pass
        if self.input_buf:
            try:
                self.stdscr.move(y, min(x + 2 + len(self.input_buf), x + box_w - 2))
            except curses.error:
                pass

    def draw_status_bar(self, h, w):
        left = f"~ {self.owner}/{self.name}"
        right = "codey 1.1.0"
        mid = self.status
        try:
            self.stdscr.addstr(h - 1, 1, left[:w - 2], curses.color_pair(4))
            if mid:
                self.stdscr.addstr(h - 1, max(len(left) + 3, 1), mid[:max(0, w - len(left) - len(right) - 6)], curses.color_pair(3))
            self.stdscr.addstr(h - 1, max(0, w - len(right) - 1), right, curses.color_pair(4))
        except curses.error:
            pass

    def wrapped_log(self, w):
        """Flatten self.log_lines (role, text) into wrapped display lines."""
        out = []
        for role, text in self.log_lines:
            label = 'you> ' if role == 'user' else 'Codey: '
            attr = curses.color_pair(2) | curses.A_BOLD if role == 'user' else curses.color_pair(1) | curses.A_BOLD
            wrapped = textwrap.wrap(text, max(10, w - len(label) - 2)) or ['']
            for i, wl in enumerate(wrapped):
                prefix = label if i == 0 else ' ' * len(label)
                pattr = attr if i == 0 else curses.A_NORMAL
                out.append((pattr, prefix + wl))
            out.append((curses.A_NORMAL, ''))
        return out

    def draw_chat(self):
        self.stdscr.clear()
        h, w = self.size()
        lines = self.wrapped_log(w)
        body_h = h - 5  # reserve input box + status bar + hint
        visible = lines[-body_h:] if len(lines) > body_h else lines
        for i, (attr, text) in enumerate(visible):
            try:
                self.stdscr.addstr(1 + i, 1, text[:w - 2], attr)
            except curses.error:
                pass

        box_y = h - 4
        self.draw_input_box(box_y, w)
        self.draw_status_bar(h, w)
        self.stdscr.refresh()

    def draw(self):
        if not self.log_lines:
            self.draw_idle()
        else:
            self.draw_chat()

    def flash_status(self, text):
        self.status = text
        self.draw()

    def send(self, message):
        self.log_lines.append(('user', message))
        self.history.append({'role': 'user', 'content': message})
        self.log_lines.append(('assistant', ''))
        reply_idx = len(self.log_lines) - 1
        self.draw()
        reply = ''
        last_draw = 0.0
        for kind, chunk in api_stream(f'/api/repos/{self.owner}/{self.name}/luiscodey/chat/stream',
                                       {'message': message, 'history': self.history[:-1]}):
            if kind == 'error':
                self.log_lines[reply_idx] = ('assistant', f"error: {chunk}")
                self.history.pop()
                self.draw()
                return
            reply += chunk
            self.log_lines[reply_idx] = ('assistant', reply)
            now = time.monotonic()
            if now - last_draw > 0.05:  # throttle redraws so fast streams don't flood the terminal
                self.draw()
                last_draw = now
        self.draw()
        self.history.append({'role': 'assistant', 'content': reply})
        self.last_blocks = CODE_BLOCK_RE.findall(reply)
        if self.last_blocks:
            names = ', '.join(f'{i+1}) {b[1]}' for i, b in enumerate(self.last_blocks))
            self.log_lines.append(('assistant', f"(code blocks: {names} — /save [n] to commit one)"))
            self.draw()

    def save_block(self, arg):
        if not self.last_blocks:
            self.log_lines.append(('assistant', '(no code blocks in the last reply)'))
            return
        idx = int(arg) - 1 if arg else 0
        if idx < 0 or idx >= len(self.last_blocks):
            self.log_lines.append(('assistant', f'(no block #{idx + 1} — last reply had {len(self.last_blocks)})'))
            return
        _, filename, code = self.last_blocks[idx]
        status, data = api_upload(self.owner, self.name, filename, code.encode(), f'{filename} via LuisCodey (codey CLI)')
        if status == 200:
            self.log_lines.append(('assistant', f"saved {filename} — commit {data['commit'][:8]}"))
        else:
            self.log_lines.append(('assistant', f"error: {data.get('error')}"))

    def handle_command(self, message):
        if message == '/clear':
            self.history, self.log_lines, self.last_blocks = [], [], []
            return True
        if message.startswith('/save'):
            parts = message.split()
            self.save_block(parts[1] if len(parts) > 1 else None)
            return True
        return False

    def run(self):
        self.draw()
        while True:
            try:
                ch = self.stdscr.get_wch()
            except KeyboardInterrupt:
                break
            except curses.error:
                continue
            if ch in ('\x1b',):  # Esc
                break
            elif ch in ('\n', '\r', curses.KEY_ENTER):
                message = self.input_buf.strip()
                self.input_buf = ''
                if not message:
                    continue
                if not self.handle_command(message):
                    self.draw()  # show user's message immediately, then fetch
                    self.send(message)
                self.draw()
            elif ch in (curses.KEY_BACKSPACE, '\x7f', '\b'):
                self.input_buf = self.input_buf[:-1]
                self.draw()
            elif isinstance(ch, str) and ch.isprintable():
                self.input_buf += ch
                self.draw()
            # ignore other control/navigation keys for now


def run_tui(owner, name):
    locale.setlocale(locale.LC_ALL, '')  # required for get_wch() to decode input correctly
    def _main(stdscr):
        Tui(stdscr, owner, name).run()
    curses.wrapper(_main)


# ---------------- repo selector (no-args entry point) ----------------
def fetch_repos():
    status, data = api_json('GET', '/api/repos')
    if status != 200:
        die(data.get('error', 'failed to list repos'))
    return data['repos']


def select_repo_plain(repos):
    print(f"\033[1mcodey\033[0m — pick a repo to chat about:\n")
    for i, r in enumerate(repos):
        desc = f" — {r['description']}" if r.get('description') else ''
        print(f"  {i + 1}) {r['owner']}/{r['name']}{desc}")
    print()
    while True:
        try:
            choice = input('repo # (or owner/repo, blank to quit): ').strip()
        except (EOFError, KeyboardInterrupt):
            print()
            return None
        if not choice:
            return None
        if choice.isdigit() and 1 <= int(choice) <= len(repos):
            r = repos[int(choice) - 1]
            return r['owner'], r['name']
        if '/' in choice:
            return tuple(choice.split('/', 1))
        print('invalid choice')


class RepoPicker:
    def __init__(self, stdscr, repos):
        self.stdscr = stdscr
        self.repos = repos
        self.idx = 0
        self.filter = ''
        curses.curs_set(1)
        curses.start_color()
        curses.use_default_colors()
        curses.init_pair(1, curses.COLOR_MAGENTA, -1)
        curses.init_pair(2, curses.COLOR_CYAN, -1)
        curses.init_pair(3, curses.COLOR_YELLOW, -1)
        curses.init_pair(4, curses.COLOR_WHITE, -1)
        self.stdscr.keypad(True)

    def matching(self):
        if not self.filter:
            return self.repos
        f = self.filter.lower()
        return [r for r in self.repos if f in f"{r['owner']}/{r['name']}".lower()]

    def draw(self):
        self.stdscr.clear()
        h, w = self.stdscr.getmaxyx()
        logo_w = max(len(l) for l in LOGO)
        y = 2
        for i, line in enumerate(LOGO):
            try:
                self.stdscr.addstr(y + i, max(0, (w - logo_w) // 2), line, curses.color_pair(1) | curses.A_BOLD)
            except curses.error:
                pass
        y += len(LOGO) + 2

        box_w = min(70, w - 6)
        x = max(0, (w - box_w) // 2)
        text = self.filter if self.filter else 'Type to filter repos...'
        attr = curses.A_NORMAL if self.filter else curses.color_pair(4)
        try:
            self.stdscr.addstr(y - 1, x, '╭' + '─' * (box_w - 2) + '╮')
            self.stdscr.addstr(y, x, '│' + ' ' * (box_w - 2) + '│')
            self.stdscr.addstr(y, x + 2, text[:box_w - 4], attr)
            self.stdscr.addstr(y + 1, x, '╰' + '─' * (box_w - 2) + '╯')
        except curses.error:
            pass

        matches = self.matching()
        self.idx = min(self.idx, max(0, len(matches) - 1))
        list_y = y + 3
        for i, r in enumerate(matches[:h - list_y - 3]):
            label = f"{r['owner']}/{r['name']}"
            if r.get('description'):
                label += f"  —  {r['description']}"
            attr = curses.color_pair(2) | curses.A_BOLD | curses.A_REVERSE if i == self.idx else curses.color_pair(4)
            try:
                self.stdscr.addstr(list_y + i, x, ('> ' if i == self.idx else '  ') + label[:box_w - 2])
            except curses.error:
                pass
        if not matches:
            try:
                self.stdscr.addstr(list_y, x, '(no matching repos)', curses.color_pair(4))
            except curses.error:
                pass

        hint = "↑/↓ or type to filter   Enter select   Esc quit"
        try:
            self.stdscr.addstr(h - 2, max(0, (w - len(hint)) // 2), hint, curses.color_pair(3))
        except curses.error:
            pass
        try:
            self.stdscr.move(y, min(x + 2 + len(self.filter), x + box_w - 2))
        except curses.error:
            pass
        self.stdscr.refresh()

    def run(self):
        self.draw()
        while True:
            try:
                ch = self.stdscr.get_wch()
            except KeyboardInterrupt:
                return None
            if ch == '\x1b':
                return None
            elif ch in ('\n', '\r', curses.KEY_ENTER):
                matches = self.matching()
                if matches:
                    r = matches[self.idx]
                    return r['owner'], r['name']
            elif ch == curses.KEY_UP:
                self.idx = max(0, self.idx - 1)
            elif ch == curses.KEY_DOWN:
                self.idx += 1
            elif ch in (curses.KEY_BACKSPACE, '\x7f', '\b'):
                self.filter = self.filter[:-1]
                self.idx = 0
            elif isinstance(ch, str) and ch.isprintable():
                self.filter += ch
                self.idx = 0
            self.draw()


def select_repo_tui(repos):
    locale.setlocale(locale.LC_ALL, '')
    result = []
    def _main(stdscr):
        result.append(RepoPicker(stdscr, repos).run())
    curses.wrapper(_main)
    return result[0]


def main():
    if sys.argv[1:2] == ['-h'] or sys.argv[1:2] == ['--help']:
        print(__doc__)
        return
    if not token():
        die("not logged in — run 'forge auth login' first (codey shares forge's saved session)")

    plain = '--plain' in sys.argv
    args = [a for a in sys.argv[1:] if a != '--plain']

    if not args:
        repos = fetch_repos()
        if not repos:
            die('no repos to chat about — create one first')
        if plain or not sys.stdout.isatty():
            picked = select_repo_plain(repos)
        else:
            picked = select_repo_tui(repos)
        if not picked:
            return
        owner, name = picked
    else:
        spec = args[0]
        if '/' not in spec:
            die('expected format: owner/repo')
        owner, name = spec.split('/', 1)

    status, meta = api_json('GET', f'/api/repos/{owner}/{name}')
    if status != 200:
        die(meta.get('error', 'repo not found'))
    if not meta.get('ai_enabled'):
        die(f"AI is turned off for {owner}/{name} — enable it in Settings first")

    if plain or not sys.stdout.isatty():
        run_plain(owner, name)
    else:
        run_tui(owner, name)


if __name__ == '__main__':
    main()