#!/usr/bin/env python3
import argparse
import datetime
import json
import os
import pathlib
import re
import sys
import tempfile
import time
from typing import Dict, List, NoReturn, Optional, Sequence, Set, Tuple

from github_client import GitHubClient, GitHubError
from issue_lookup import lookup_issue_candidates, snapshots
from issues import occurrence_markers, plan_issue
from junit import ArtifactError, classify_reports, read_junit_zip
from model import (
    GovernanceState,
    JsonValue,
    Occurrence,
    PendingRecurrence,
    RejectedArtifact,
    Report,
    Summary,
    WorkflowRun,
    decode_json,
)


WORKFLOWS = ((".github/workflows/unit-test.yml", "unit-test-junit-attempt-"), (".github/workflows/integrations.yml", "integration-test-junit-attempt-"))
STATE_VERSION = 1


def log_progress(message: str) -> None:
    timestamp = datetime.datetime.now(datetime.timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")
    print("[flaky-governance] %s %s" % (timestamp, message), file=sys.stderr, flush=True)


def utc(value: str) -> datetime.datetime:
    return datetime.datetime.strptime(value, "%Y-%m-%dT%H:%M:%SZ").replace(tzinfo=datetime.timezone.utc)


def iso(value: datetime.datetime) -> str:
    return value.astimezone(datetime.timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")


class Arguments(argparse.Namespace):
    run_id: Optional[int]
    since: Optional[datetime.datetime]
    dry_run: bool

    def __init__(self) -> None:
        super().__init__()
        self.run_id = None
        self.since = None
        self.dry_run = False


def parse_args(argv: Sequence[str]) -> Arguments:
    parser = argparse.ArgumentParser(description="Reconcile flaky-test GitHub issues")
    setattr(parser, "error", argument_error)
    selectors = parser.add_mutually_exclusive_group()
    _ = selectors.add_argument("--run-id", type=int)
    _ = selectors.add_argument("--since", type=utc)
    _ = parser.add_argument("--dry-run", action="store_true")
    return parser.parse_args(argv, namespace=Arguments())


class ArgumentError(Exception):
    pass


def argument_error(message: str) -> NoReturn:
    raise ArgumentError(message)


def empty_state() -> GovernanceState:
    return GovernanceState(STATE_VERSION, None, [], [])


def _mapping(value: JsonValue) -> Dict[str, JsonValue]:
    if not isinstance(value, dict):
        raise RuntimeError("invalid state shape")
    return value


def _sequence(value: JsonValue) -> List[JsonValue]:
    if not isinstance(value, list):
        raise RuntimeError("invalid state shape")
    return value


def _integer(value: JsonValue) -> int:
    if isinstance(value, bool) or not isinstance(value, int):
        raise RuntimeError("invalid state integer")
    return value


def _string(value: JsonValue) -> str:
    if not isinstance(value, str):
        raise RuntimeError("invalid state string")
    return value


def _optional_string(value: JsonValue) -> Optional[str]:
    if value is None:
        return None
    return _string(value)


def load_state(path: pathlib.Path) -> GovernanceState:
    if not path.exists():
        return empty_state()
    try:
        raw = decode_json(path.read_text(encoding="utf-8"))
    except (OSError, json.JSONDecodeError) as error:
        raise RuntimeError("unable to read state") from error
    state = _mapping(raw)
    if _integer(state.get("schema_version")) != STATE_VERSION:
        raise RuntimeError("unsupported state schema")
    pending: List[PendingRecurrence] = []
    for value in _sequence(state.get("pending_recurrences", [])):
        record = _mapping(value)
        pending.append(PendingRecurrence(
            _integer(record.get("issue_number")),
            _optional_string(record.get("previous_closed_at")),
            [_string(item) for item in _sequence(record.get("fingerprints", []))],
            [_string(item) for item in _sequence(record.get("chunks", []))],
        ))
    rejected: List[RejectedArtifact] = []
    for value in _sequence(state.get("rejected_artifacts", [])):
        record = _mapping(value)
        rejected.append(RejectedArtifact(
            _integer(record.get("id")),
            _string(record.get("rejected_at")),
            _string(record.get("kind", "whole")),
        ))
    return GovernanceState(STATE_VERSION, _optional_string(state.get("successful_scan_started_at")), pending, rejected)


def save_state(path: pathlib.Path, state: GovernanceState) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    descriptor, temporary = tempfile.mkstemp(prefix="state.", suffix=".tmp", dir=str(path.parent))
    try:
        with os.fdopen(descriptor, "w", encoding="utf-8") as stream:
            payload = {
                "schema_version": state.schema_version,
                "successful_scan_started_at": state.successful_scan_started_at,
                "pending_recurrences": [record.__dict__ for record in state.pending_recurrences],
                "rejected_artifacts": [record.__dict__ for record in state.rejected_artifacts],
            }
            json.dump(payload, stream, sort_keys=True)
            stream.flush()
            os.fsync(stream.fileno())
        os.replace(temporary, str(path))
    finally:
        if os.path.exists(temporary):
            os.unlink(temporary)


def discover_partition(client: GitHubClient, workflow_id: int, start: datetime.datetime, end: datetime.datetime) -> List[WorkflowRun]:
    response = client.workflow_runs(workflow_id, iso(start), iso(end - datetime.timedelta(seconds=1)))
    if response.total_count < 1000:
        return list(response.runs)
    seconds = int((end - start).total_seconds())
    if seconds <= 1:
        raise GitHubError("workflow run result cap reached at one second")
    midpoint = start + datetime.timedelta(seconds=seconds // 2)
    combined = discover_partition(client, workflow_id, start, midpoint) + discover_partition(client, workflow_id, midpoint, end)
    return list({item.id: item for item in combined}.values())


class TransientBranchError(Exception):
    pass


def branch_for(run: WorkflowRun) -> str:
    if run.event == "pull_request":
        branches = {item.base_ref for item in run.pull_requests if item.base_ref is not None}
        if len(branches) == 1:
            return next(iter(branches))
        if not branches:
            # GitHub intermittently reports an empty pull_requests array for
            # pull_request runs (forks, older runs). A later scan may see the
            # association, so the run is skipped rather than permanently
            # rejected.
            raise TransientBranchError("pull request run %d has no associated base branch" % run.id)
        raise ArtifactError("ambiguous pull request base branch")
    if run.event == "workflow_dispatch" and run.head_branch:
        return run.head_branch
    raise ArtifactError("ambiguous workflow branch")


def discover(client: GitHubClient, args: Arguments, state: GovernanceState, now: datetime.datetime, retention_days: int) -> List[WorkflowRun]:
    if args.run_id is not None:
        run = client.run(args.run_id)
        if run.path not in dict(WORKFLOWS):
            raise GitHubError("run does not belong to a governed workflow")
        if run.status != "completed":
            raise GitHubError("workflow run is not completed")
        return [run]
    if args.since is not None:
        threshold = args.since
    elif state.successful_scan_started_at:
        threshold = utc(state.successful_scan_started_at) - datetime.timedelta(hours=1)
    else:
        threshold = now - datetime.timedelta(hours=int(os.environ.get("FLAKY_GOVERNANCE_INITIAL_LOOKBACK_HOURS", "24")))
    horizon = now - datetime.timedelta(days=max(30, retention_days))
    selected: Dict[int, WorkflowRun] = {}
    for path, _ in WORKFLOWS:
        log_progress("discovering workflow runs: %s" % path)
        workflow = client.workflow(path)
        for run in discover_partition(client, workflow.id, horizon, now):
            if run.status == "completed" and utc(run.updated_at) >= threshold:
                selected[run.id] = WorkflowRun(
                    run.id, run.status, run.run_attempt, run.event, run.head_branch,
                    run.pull_requests, run.updated_at, run.run_started_at, run.html_url,
                    run.name, run.path or path,
                )
    runs = list(selected.values())
    log_progress("workflow discovery found %d completed run(s)" % len(runs))
    return runs


def collect_occurrences(client: GitHubClient, run: WorkflowRun, summary: Summary, state: GovernanceState, now: datetime.datetime, record_rejections: bool) -> List[Occurrence]:
    started = time.monotonic()
    log_progress("collecting artifacts for run %d" % run.id)
    path = run.path
    if path is None:
        raise ArtifactError("workflow path is missing")
    prefix = dict(WORKFLOWS)[path]
    reports: List[Report] = []
    matched_artifacts = 0
    rejected_ids = {record.id for record in state.rejected_artifacts if record.kind == "whole"}
    for artifact in client.artifacts(run.id).artifacts:
        match = re.fullmatch(re.escape(prefix) + r"([1-9][0-9]*)", artifact.name)
        if not match:
            continue
        matched_artifacts += 1
        if artifact.expired:
            summary.rejected_artifacts += 1
            if record_rejections and artifact.id not in rejected_ids:
                state.rejected_artifacts.append(RejectedArtifact(artifact.id, iso(now)))
            continue
        if artifact.id in rejected_ids:
            summary.rejected_artifacts += 1
            continue
        try:
            report = read_junit_zip(client.download_artifact(artifact.id), int(match.group(1)))
            reports.append(report)
            summary.downloaded_artifacts += 1
            if report.ambiguous:
                summary.rejected_artifacts += 1
                if record_rejections and not any(record.id == artifact.id for record in state.rejected_artifacts):
                    state.rejected_artifacts.append(RejectedArtifact(artifact.id, iso(now), "partial"))
        except ArtifactError:
            summary.rejected_artifacts += 1
            if record_rejections and artifact.id not in rejected_ids:
                state.rejected_artifacts.append(RejectedArtifact(artifact.id, iso(now)))
    if not matched_artifacts:
        print("run %s has no matching artifact" % run.id, file=sys.stderr)
    attempt_times: Dict[int, str] = {}
    for report in reports:
        if report.attempt == run.run_attempt:
            attempt_times[report.attempt] = run.run_started_at
        else:
            attempt_times[report.attempt] = client.attempt_started_at(run.id, report.attempt)
    occurrences = classify_reports(client.repository, path, run.id, run.run_attempt, reports, run.name, run.html_url, "", attempt_times)
    log_progress("collected run %d: %d matching artifact(s), %d report(s), %d occurrence(s) in %.1fs" % (run.id, matched_artifacts, len(reports), len(occurrences), time.monotonic() - started))
    return occurrences


def apply_pending(client: GitHubClient, state: GovernanceState, summary: Summary) -> None:
    pending = list(state.pending_recurrences)
    log_progress("applying %d pending recurrence(s)" % len(pending))
    if not pending:
        return
    remotes = {item.number: item for item in snapshots(client, (record.issue_number for record in pending))}
    for record in pending:
        remote = remotes.get(record.issue_number)
        markers: Set[str] = set()
        if remote:
            markers.update(occurrence_markers(remote.body))
            for comment in remote.comments:
                markers.update(occurrence_markers(comment))
            if remote.state == "closed":
                client.reopen(record.issue_number)
                summary.reopens += 1
        for chunk in record.chunks:
            fingerprints = occurrence_markers(chunk)
            if not all(fingerprint in markers for fingerprint in fingerprints):
                client.comment(record.issue_number, chunk)
                summary.comments += 1
                markers.update(fingerprints)
        state.pending_recurrences.remove(record)
        if summary.state_path is None:
            raise RuntimeError("state path is missing")
        save_state(pathlib.Path(summary.state_path), state)


def new_summary() -> Summary:
    return Summary()


def run(argv: Sequence[str], summary: Optional[Summary] = None) -> Summary:
    started = time.monotonic()
    summary = summary if summary is not None else new_summary()
    args = parse_args(argv)
    mode = "dry-run" if args.dry_run else "live"
    selector = "run %d" % args.run_id if args.run_id is not None else "since %s" % iso(args.since) if args.since is not None else "saved cursor"
    log_progress("pipeline started (%s, %s)" % (mode, selector))
    repository = os.environ.get("GITHUB_REPOSITORY", "")
    labels = tuple(filter(None, (part.strip() for part in os.environ.get("FLAKY_GOVERNANCE_LABELS", "").split(","))))
    if not repository or not labels:
        raise RuntimeError("GITHUB_REPOSITORY and FLAKY_GOVERNANCE_LABELS are required")
    state_dir = pathlib.Path(os.environ.get("FLAKY_GOVERNANCE_STATE_DIR") or pathlib.Path(__file__).parent / ".state")
    state_path = state_dir / "state.json"
    state = load_state(state_path)
    log_progress("loaded state from %s" % state_path)
    now = utc(os.environ.get("FLAKY_GOVERNANCE_NOW", iso(datetime.datetime.now(datetime.timezone.utc))))
    summary.state_path = str(state_path)
    client = GitHubClient(repository)
    log_progress("verifying configured labels")
    client.verify_labels(labels)
    log_progress("configured labels verified")
    if not args.dry_run:
        apply_pending(client, state, summary)
    else:
        log_progress("dry-run: skipping pending recurrence mutations")
    log_progress("starting workflow discovery")
    retention_days = int(os.environ.get("FLAKY_GOVERNANCE_ARTIFACT_RETENTION_DAYS", "90"))
    runs = discover(client, args, state, now, retention_days)
    summary.scanned_runs = len(runs)
    grouped: Dict[Tuple[str, str, str], List[Occurrence]] = {}
    selector_mode = args.run_id is not None or args.since is not None
    retention_cutoff = now - datetime.timedelta(days=retention_days)
    if not selector_mode:
        state.rejected_artifacts = [record for record in state.rejected_artifacts if utc(record.rejected_at) >= retention_cutoff]
    for workflow_run in runs:
        try:
            branch = branch_for(workflow_run)
            for occurrence in collect_occurrences(client, workflow_run, summary, state, now, not selector_mode):
                grouped.setdefault((branch, occurrence.classname, occurrence.testcase_name), []).append(occurrence)
        except TransientBranchError as error:
            log_progress("skipping run: %s" % error)
        except ArtifactError:
            summary.rejected_artifacts += 1
    log_progress("classification found %d occurrence(s) across %d test identity group(s)" % (sum(len(items) for items in grouped.values()), len(grouped)))
    summary.detected_occurrences = sum(len(items) for items in grouped.values())
    log_progress("loading GitHub issue candidates")
    issue_list = lookup_issue_candidates(client, grouped)
    log_progress("loaded %d GitHub issue candidate(s)" % len(issue_list))
    plans = [plan_issue(branch, occurrences, issue_list) for (branch, _, _), occurrences in grouped.items()]
    log_progress("generated %d issue plan(s)" % len(plans))
    for plan in sorted(plans, key=lambda item: item.key):
        if plan.kind == "create":
            summary.creates += 1
            summary.comments += max(0, len(plan.chunks) - 1)
        elif plan.kind == "comment":
            summary.comments += len(plan.chunks)
        elif plan.kind == "reopen":
            summary.reopens += 1
            summary.comments += len(plan.chunks)
        elif plan.kind == "noop":
            summary.noops += 1
        elif plan.kind == "conflict":
            summary.conflicts += 1
            print("conflict for issue key %s" % plan.key, file=sys.stderr)
        if plan.kind not in ("noop", "conflict"):
            issue = "new" if plan.issue_number is None else str(plan.issue_number)
            print("plan %s issue %s with %d comment%s" % (plan.kind, issue, len(plan.chunks), "" if len(plan.chunks) == 1 else "s"), file=sys.stderr)
        if plan.kind == "create":
            print("issue title: %s" % plan.title, file=sys.stderr)
            if args.dry_run:
                print("issue body:\n%s" % plan.chunks[0], file=sys.stderr)
        if args.dry_run or plan.kind in ("noop", "conflict"):
            continue
        if plan.kind == "create":
            created = client.create_issue(plan.title, plan.chunks[0], labels)
            log_progress("created issue https://github.com/%s/issues/%d" % (repository, created.number))
            for chunk in plan.chunks[1:]:
                client.comment(created.number, chunk)
        elif plan.kind == "comment":
            if plan.issue_number is None:
                raise RuntimeError("comment plan is missing issue number")
            for chunk in plan.chunks:
                client.comment(plan.issue_number, chunk)
            log_progress("commented on issue https://github.com/%s/issues/%d" % (repository, plan.issue_number))
        elif plan.kind == "reopen":
            if plan.issue_number is None:
                raise RuntimeError("reopen plan is missing issue number")
            fingerprints: Set[str] = set()
            for chunk in plan.chunks:
                fingerprints.update(occurrence_markers(chunk))
            record = PendingRecurrence(plan.issue_number, plan.previous_closed_at, sorted(fingerprints), list(plan.chunks))
            state.pending_recurrences.append(record)
            save_state(state_path, state)
            client.reopen(plan.issue_number)
            for chunk in plan.chunks:
                client.comment(plan.issue_number, chunk)
            state.pending_recurrences.remove(record)
            save_state(state_path, state)
            log_progress("reopened issue https://github.com/%s/issues/%d" % (repository, plan.issue_number))
    if not args.dry_run and args.run_id is None and args.since is None and not summary.conflicts:
        state.successful_scan_started_at = iso(now)
        save_state(state_path, state)
    if summary.rejected_artifacts:
        summary.exit_code = 2
    summary.state_path = None
    log_progress("pipeline completed in %.1fs" % (time.monotonic() - started))
    return summary


def main() -> int:
    summary = new_summary()
    try:
        summary = run(sys.argv[1:], summary)
    except Exception as error:  # noqa: BROAD_EXCEPT_OK
        log_progress("pipeline failed: %s: %s" % (type(error).__name__, error))
        summary.state_path = None
        summary.exit_code = 1
    print(json.dumps(summary.as_json(), sort_keys=True, separators=(",", ":")))
    return summary.exit_code


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