#!/usr/bin/env python3
"""Append one PR under a bead claim, keeping bd's TTL alive with 120 s heartbeats."""
from __future__ import annotations

import argparse
import copy
from datetime import datetime, timezone
import json
import os
from pathlib import Path
import re
import subprocess
import sys
import tempfile
import time

from stack_top import select_stack_top


PULL_FILTER = '.[] | {number, base: .base.ref, head: .head.ref, sha: .head.sha, draft}'
STACK_FILTER = '.[] | {number, base: .base.ref, prs: [.pull_requests[].number]}'
COMMAND_TIMEOUT = 120
HEARTBEAT_INTERVAL = 120


class StackError(Exception):
    pass


class Parser(argparse.ArgumentParser):
    def error(self, message):
        raise StackError(message)


def rows(output):
    try:
        value = json.loads(output)
        return value if isinstance(value, list) else [value]
    except json.JSONDecodeError:
        try:
            return [json.loads(line) for line in output.splitlines() if line.strip()]
        except json.JSONDecodeError as error:
            raise StackError(f"invalid JSON from command: {error}") from error


def timestamp(value):
    if not isinstance(value, str):
        return None
    try:
        parsed = datetime.fromisoformat(value.replace('Z', '+00:00'))
        return parsed.timestamp() if parsed.tzinfo is not None else None
    except ValueError:
        return None


class Linker:
    def __init__(self, values, runner, clock, sleep, actor):
        self.values = values
        self.runner = runner
        self.clock = clock
        self.sleep = sleep
        self.actor = actor
        self.held = False
        self.heartbeat_at = 0
        self.reclaimed = None
        self.mode = None
        self.lean_reason = None

    def command(self, argv, *, check=True, allow_transport_failure=False):
        if self.held and argv[:2] != ['bd', 'heartbeat'] and self.clock() - self.heartbeat_at >= HEARTBEAT_INTERVAL:
            self.heartbeat()
        try:
            result = self.runner(argv, capture_output=True, text=True, timeout=COMMAND_TIMEOUT)
        except (OSError, subprocess.TimeoutExpired) as error:
            if allow_transport_failure:
                return subprocess.CompletedProcess(argv, 1, '', str(error))
            raise StackError(f"{' '.join(argv)}: {error}") from error
        if check and result.returncode:
            detail = (result.stderr or '').strip().splitlines()
            raise StackError(f"{' '.join(argv)} [{result.returncode}] {detail[0] if detail else 'command failed'}")
        return result

    def heartbeat(self):
        self.command(['bd', 'heartbeat', self.values['stack_bead'], '--actor', self.actor])
        self.heartbeat_at = self.clock()

    def bead(self):
        result = self.command(['bd', 'show', self.values['stack_bead'], '--json'])
        found = rows(result.stdout)
        if len(found) != 1 or not isinstance(found[0], dict):
            raise StackError('bd show returned an invalid stack bead')
        return found[0]

    def acquire(self):
        deadline = self.clock() + self.values['timeout_s']
        while True:
            result = self.command(['bd', 'update', self.values['stack_bead'], '--claim', '--actor', self.actor], check=False)
            if result.returncode == 0:
                self.held = True
                self.heartbeat_at = self.clock()
                return None
            bead = self.bead()
            if not bead.get('assignee') or bead.get('assignee') == self.actor:
                detail = (result.stderr or '').strip().splitlines()
                raise StackError(f"bd update --claim [{result.returncode}] {detail[0] if detail else 'claim failed'}")
            expiry = timestamp(bead.get('lease_expires_at'))
            if bead.get('assignee') and expiry is not None and self.clock() - expiry > 1800:
                reclaimed = {'holder': bead['assignee'], 'claimed_at': bead.get('started_at')}
                reclaimed_result = self.command(['bd', 'reclaim', '--id', self.values['stack_bead'], '--older-than', '1800s', '--actor', self.actor], check=False)
                if reclaimed_result.returncode == 0 and not self.bead().get('assignee'):
                    self.reclaimed = reclaimed
                    continue
            remaining = deadline - self.clock()
            if remaining <= 0:
                holder = bead.get('assignee')
                try:
                    task = self.command(['atm', 'task', 'list', '--all', '--json'])
                    data = json.loads(task.stdout)
                    tasks = data.get('tasks', []) if isinstance(data, dict) else data
                    active = next((item.get('task_id') or item.get('id') for item in tasks
                                   if isinstance(item, dict) and item.get('assignee') == holder
                                   and item.get('state', item.get('status')) in ('active', 'in_progress', 'started')), None)
                except (StackError, ValueError, TypeError):
                    active = None
                return {'result': 'lock_timeout', 'holder': holder, 'claimed_at': bead.get('started_at'), 'holder_task': active}
            self.sleep(min(5, remaining))

    def release(self):
        if self.held:
            self.held = False
            self.command(['bd', 'update', self.values['stack_bead'], '--if-assignee', self.actor,
                          '--status', 'open', '--assignee', '', '--actor', self.actor])

    def pulls(self):
        result = self.command(['gh', 'api', 'repos/{owner}/{repo}/pulls?state=open&per_page=100',
                               '--paginate', '--jq', PULL_FILTER])
        found = rows(result.stdout) if result.stdout.strip() else []
        if any(not isinstance(item, dict) or type(item.get('number')) is not int or
               not all(isinstance(item.get(key), str) and item[key] for key in ('base', 'head', 'sha')) or
               type(item.get('draft')) is not bool for item in found):
            raise StackError('filtered pull list has invalid fields')
        if len({item['number'] for item in found}) != len(found):
            raise StackError('filtered pull list contains duplicate pull requests')
        return found

    def lean_mode(self):
        if self.values['budget'] == 'lean':
            self.lean_reason = 'budget_lean'
        else:
            quota = self.command(['gh', 'api', 'rate_limit', '--jq', '.resources.graphql'], check=False, allow_transport_failure=True)
            if quota.returncode:
                self.lean_reason = 'rate_limit_call_failed'
            else:
                try:
                    self.lean_reason = None if json.loads(quota.stdout)['remaining'] > 0 else 'graphql_remaining_0'
                except (ValueError, KeyError, TypeError):
                    self.lean_reason = 'rate_limit_response_invalid'
        self.mode = 'lean' if self.lean_reason else 'full'
        if self.lean_reason:
            print(f'sanity-stack-link: lean mode: {self.lean_reason}', file=sys.stderr)
        return self.mode == 'lean'

    def stacks(self):
        result = self.command(['gh', 'api', 'repos/{owner}/{repo}/stacks', '--paginate', '--jq', STACK_FILTER])
        found = rows(result.stdout) if result.stdout.strip() else []
        if any(not isinstance(item, dict) or type(item.get('number')) is not int or
               not isinstance(item.get('base'), str) or not isinstance(item.get('prs'), list) or
               any(type(number) is not int for number in item['prs']) for item in found):
            raise StackError('filtered hosted stack list has invalid fields')
        return found

    def reconcile(self, metadata, pulls, stacks):
        layers = metadata['layers']
        known = {entry['pr'] for entry in layers}
        matching = [stack for stack in stacks if stack['number'] == metadata['stack_number'] or
                    (stack['base'] == metadata['trunk'] and known.intersection(stack['prs']))]
        if len(matching) > 1:
            raise StackError('ambiguous hosted stacks overlap the ledger')
        if not matching:
            if metadata['stack_number'] is not None or any(layer['linked'] for layer in layers):
                raise StackError('hosted stack missing for linked ledger')
            return
        stack = matching[0]
        if stack['base'] != metadata['trunk'] or len(set(stack['prs'])) != len(stack['prs']):
            raise StackError('hosted stack has conflicting ancestry')
        by_number = {pull['number']: pull for pull in pulls}
        previous = metadata['trunk']
        for index, number in enumerate(stack['prs']):
            pull = by_number.get(number)
            if not pull or pull['draft'] or pull['base'] != previous:
                raise StackError('hosted stack pull chain is missing or ambiguous')
            ref = self.command(['git', 'rev-parse', 'refs/remotes/origin/' + pull['head']])
            if ref.stdout.strip() != pull['sha']:
                raise StackError('hosted stack head differs from fetched branch')
            ancestry = self.command(['git', 'merge-base', '--is-ancestor', 'origin/' + previous, pull['sha']], check=False)
            if ancestry.returncode:
                raise StackError('hosted stack has conflicting git ancestry')
            if index < len(layers):
                layer = layers[index]
                if (layer['pr'], layer['branch'], layer['head'], layer['base']) != (number, pull['head'], pull['sha'], previous):
                    raise StackError('hosted stack conflicts with ledger ancestry')
                layer['linked'] = True
            else:
                if (sum(other['head'] == pull['head'] for other in pulls) != 1 or
                        any(layer['branch'] == pull['head'] or layer['pr'] == number for layer in layers)):
                    raise StackError('hosted-only stack layer is not uniquely attributable')
                layers.append({'branch': pull['head'], 'pr': number, 'head': pull['sha'], 'base': previous, 'linked': True})
            previous = pull['head']
        if any(layer['linked'] for layer in layers[len(stack['prs']):]):
            raise StackError('linked ledger layer is absent from hosted stack')
        metadata['stack_number'] = stack['number']

    def write(self, metadata):
        with tempfile.NamedTemporaryFile(mode='w', suffix='.json', encoding='utf-8') as stream:
            json.dump(metadata, stream)
            stream.flush()
            self.command(['bd', 'update', self.values['stack_bead'], '--if-assignee', self.actor,
                          '--metadata', '@' + stream.name, '--actor', self.actor])

    def execute(self):
        timeout = self.acquire()
        if timeout:
            return timeout
        bead = self.bead()
        metadata = copy.deepcopy(bead.get('metadata'))
        if (not isinstance(metadata, dict) or not isinstance(metadata.get('trunk'), str) or not metadata['trunk'] or
                not isinstance(metadata.get('layers'), list) or
                (metadata.get('stack_number') is not None and (type(metadata['stack_number']) is not int or metadata['stack_number'] <= 0))):
            raise StackError('stack bead metadata is invalid')
        layers = metadata['layers']
        previous = metadata['trunk']
        branches, numbers = set(), set()
        for layer in layers:
            if (not isinstance(layer, dict) or not isinstance(layer.get('branch'), str) or not layer['branch'] or
                    type(layer.get('pr')) is not int or layer['pr'] <= 0 or not isinstance(layer.get('head'), str) or
                    not re.fullmatch(r'[0-9a-fA-F]{7,64}', layer['head']) or type(layer.get('linked')) is not bool or
                    layer.get('base') != previous or layer['branch'] in branches or layer['pr'] in numbers):
                raise StackError('stack bead has an invalid layer')
            previous = layer['branch']
            branches.add(layer['branch'])
            numbers.add(layer['pr'])
        original = copy.deepcopy(metadata)
        lean = self.lean_mode()
        pulls = self.pulls()
        if not lean and (layers or self.reclaimed):
            self.command(['git', 'fetch', 'origin'])
            self.reconcile(metadata, pulls, self.stacks())
        candidate = next((pull for pull in pulls if pull['number'] == self.values['pr_number']), None)
        if not candidate or candidate['head'] != self.values['branch'] or candidate['sha'] != self.values['commit'] or candidate['draft']:
            raise StackError('checked pull request is missing, draft, or differs from the pinned branch and commit')
        refreshed = self.bead()
        if refreshed.get('assignee') != self.actor or refreshed.get('metadata') != original:
            raise StackError('stack claim or ledger changed during reconciliation')
        self.command(['git', 'fetch', 'origin'])

        def resolve_head(branch):
            return self.command(['git', 'rev-parse', 'refs/remotes/origin/' + branch]).stdout.strip()

        base, recorded_head = select_stack_top({'metadata': metadata}, metadata['trunk'], resolve_head)
        if layers and layers[-1]['pr'] == self.values['pr_number']:
            top = layers[-1]
            actual = resolve_head(base)
            if (base != self.values['branch'] or recorded_head != self.values['commit'] or
                    actual != recorded_head or candidate['base'] != top['base']):
                raise StackError('checked pull request conflicts with the recorded stack top')
            if metadata != original:
                self.write(metadata)
            return {'result': 'stacked', 'top': top, 'stack_number': metadata.get('stack_number'), 'reclaimed': self.reclaimed}
        if any(layer['pr'] == self.values['pr_number'] or layer['branch'] == self.values['branch'] for layer in layers):
            raise StackError('checked layer is already below the stack top')
        ancestor = self.command(['git', 'merge-base', '--is-ancestor', 'origin/' + base, self.values['commit']], check=False)
        actual = resolve_head(base) if layers else recorded_head
        if actual != recorded_head:
            raise StackError('fetched stack top differs from its recorded head')
        if ancestor.returncode == 1:
            return {'result': 'rebase_required', 'top': {'branch': base, 'head': actual},
                    'evidence': f"origin/{base} is not an ancestor of {self.values['commit']}"}
        if ancestor.returncode:
            raise StackError(f"git merge-base --is-ancestor [{ancestor.returncode}] {(ancestor.stderr or '').strip()}")
        if (lean or not layers) and candidate['base'] != base:
            raise StackError(f"pull request base {candidate['base']} differs from stack top {base}")
        new = {'branch': self.values['branch'], 'pr': self.values['pr_number'], 'head': self.values['commit'], 'base': base, 'linked': False}
        if not lean and layers:
            pending = [entry for entry in layers if not entry['linked']] + [new]
            stack_number = metadata.get('stack_number')
            if stack_number is None:
                self.command(['gh', 'stack', 'link', '--base', metadata['trunk']] + [str(entry['pr']) for entry in pending])
            else:
                for entry in pending:
                    self.command(['gh', 'stack', 'link', str(stack_number), str(entry['pr'])])
            hosted = self.stacks()
            expected = [entry['pr'] for entry in layers] + [new['pr']]
            matches = [stack for stack in hosted if stack['base'] == metadata['trunk'] and stack['prs'] == expected
                       and (stack_number is None or stack['number'] == stack_number)]
            if len(matches) != 1:
                raise StackError('linked hosted stack identity or order is missing or ambiguous')
            metadata['stack_number'] = matches[0]['number']
            for entry in pending:
                entry['linked'] = True
        layers.append(new)
        self.write(metadata)
        return {'result': 'stacked', 'top': new, 'stack_number': metadata.get('stack_number'), 'reclaimed': self.reclaimed}


def main(argv=None, runner=subprocess.run, clock=time.time, sleep=time.sleep):
    parser = Parser(description=__doc__)
    parser.add_argument('--input', required=True, type=Path)
    linker = None
    try:
        args = parser.parse_args(argv)
        values = json.loads(args.input.read_text(encoding='utf-8'))
        if not isinstance(values, dict) or any(not isinstance(values.get(key), str) or not values[key] for key in ('checked_bead', 'branch', 'commit', 'stack_bead')):
            raise StackError('input requires checked_bead, branch, commit and stack_bead strings')
        if type(values.get('pr_number')) is not int or values['pr_number'] <= 0 or values.get('budget') not in ('full', 'lean'):
            raise StackError('input requires a positive pr_number and budget full or lean')
        if not re.fullmatch(r'[0-9a-fA-F]{7,64}', values['commit']):
            raise StackError('commit must be a hexadecimal Git object ID')
        values.setdefault('timeout_s', 60)
        if type(values['timeout_s']) not in (int, float) or not 0 <= values['timeout_s'] <= 3600:
            raise StackError('timeout_s must be between 0 and 3600')
        actor = os.environ.get('ATM_IDENTITY')
        if not actor:
            raise StackError('ATM_IDENTITY is required')
        linker = Linker(values, runner, clock, sleep, actor)
        result = None
        try:
            try:
                result = linker.execute()
            except (StackError, OSError, ValueError, TypeError) as error:
                result = {'result': 'stack_failed', 'evidence': str(error)}
        finally:
            try:
                linker.release()
            except StackError as error:
                if result is None:
                    raise
                if result['result'] == 'stacked':
                    result['release_warning'] = str(error)
                elif result['result'] == 'stack_failed':
                    result['evidence'] += f'; release failed: {error}'
                else:
                    result = {'result': 'stack_failed', 'evidence': f'release failed after {result["result"]}: {error}'}
    except (StackError, OSError, ValueError, TypeError) as error:
        result = {'result': 'stack_failed', 'evidence': str(error)}
    if linker is not None and linker.mode is not None:
        result.update(mode=linker.mode, lean_reason=linker.lean_reason)
    print(json.dumps(result))
    return int(result['result'] == 'stack_failed')


if __name__ == '__main__':
    sys.exit(main())
