# /// script
# dependencies = [
#     "aiohttp",
# ]
# ///
"""
Script to visualize cross-version test results for MLflow autologging and models.

This script fetches scheduled workflow run results from GitHub Actions and generates
a markdown table showing the test status for different package versions across
different dates.

Usage:
    uv run dev/xtest_viz.py                      # Fetch last 14 days from mlflow/dev
    uv run dev/xtest_viz.py --days 30            # Fetch last 30 days
    uv run dev/xtest_viz.py --repo mlflow/mlflow  # Use different repo

Example output (truncated for brevity):
    | Name                                   | 2024-01-15 | 2024-01-14 | 2024-01-13 |
    |----------------------------------------|------------|------------|------------|
    | test1 (sklearn, 1.3.1, autologging...) | [✅](link) | [✅](link) | [❌](link) |
    | test1 (pytorch, 2.1.0, models...)      | [✅](link) | [⚠️](link) | [✅](link) |
    | test2 (xgboost, 2.0.0, autologging...) | [❌](link) | [✅](link) | —          |

Where:
    ✅ = success
    ❌ = failure
    ⚠️ = cancelled
    ❓ = unknown status
    — = no data
"""

import argparse
import asyncio
import json
import os
import re
import sys
from dataclasses import asdict, dataclass
from datetime import datetime, timedelta
from pathlib import Path
from typing import Any, cast

import aiohttp


@dataclass
class JobResult:
    name: str
    conclusion: str
    date: str
    started_at: str
    completed_at: str
    failed_step: str | None
    html_url: str
    logs_url: str


class XTestViz:
    def __init__(self, github_token: str | None = None, repo: str = "mlflow/dev"):
        self.github_token = github_token or os.environ.get("GH_TOKEN")
        self.repo = repo
        self.per_page = 30
        self.headers: dict[str, str] = {}
        if self.github_token:
            self.headers["Authorization"] = f"token {self.github_token}"
            self.headers["Accept"] = "application/vnd.github.v3+json"

    def status_to_emoji(self, status: str) -> str | None:
        """Convert job status to emoji representation.

        Returns None for skipped status to indicate it should be filtered out.
        """
        match status:
            case "success":
                return "✅"
            case "failure":
                return "❌"
            case "cancelled":
                return "⚠️"
            case "skipped":
                return None
            case _:
                return "❓"

    def parse_job_name(self, job_name: str) -> str:
        """Extract string inside parentheses from job name.

        Examples:
        - "test1 (sklearn / autologging / 1.3.1)" -> "sklearn / autologging / 1.3.1"
        - "test2 (pytorch / models / 2.1.0)" -> "pytorch / models / 2.1.0"

        Returns:
            str: Content inside parentheses, or original name if no parentheses found
        """
        # Pattern to match: anything (content)
        pattern = r"\(([^)]+)\)"
        if match := re.search(pattern, job_name.strip()):
            return match.group(1).strip()

        return job_name

    async def _make_request(
        self,
        session: aiohttp.ClientSession,
        url: str,
        params: dict[str, str] | None = None,
    ) -> dict[str, Any]:
        """Make an async HTTP GET request and return JSON response."""
        async with session.get(url, headers=self.headers, params=params) as response:
            response.raise_for_status()
            return cast(dict[str, Any], await response.json())

    async def get_workflow_runs(
        self, session: aiohttp.ClientSession, days_back: int = 30
    ) -> list[dict[str, Any]]:
        """Fetch cross-version test workflow runs from the last N days."""
        since_date = (datetime.now() - timedelta(days=days_back)).isoformat()

        print(f"Fetching scheduled workflow runs from last {days_back} days...", file=sys.stderr)

        all_runs: list[dict[str, Any]] = []
        page = 1

        while True:
            params = {
                "per_page": str(self.per_page),
                "page": str(page),
                "created": f">={since_date}",
                "status": "completed",
                "event": "schedule",
            }
            url = f"https://api.github.com/repos/{self.repo}/actions/workflows/cross-version-tests.yml/runs"

            data = await self._make_request(session, url, params=params)
            runs = data.get("workflow_runs", [])

            if not runs:
                break

            all_runs.extend(runs)

            print(f"  Fetched page {page} ({len(runs)} runs)", file=sys.stderr)

            if len(runs) < self.per_page:
                break

            page += 1

        print(f"Found {len(all_runs)} scheduled workflow runs total", file=sys.stderr)

        return all_runs

    async def get_workflow_jobs(
        self, session: aiohttp.ClientSession, run_id: int
    ) -> list[dict[str, Any]]:
        """Get jobs for a specific workflow run."""
        all_jobs: list[dict[str, Any]] = []
        page = 1

        while True:
            params = {"per_page": str(self.per_page), "page": str(page)}
            url = f"https://api.github.com/repos/{self.repo}/actions/runs/{run_id}/jobs"

            data = await self._make_request(session, url, params=params)
            jobs = data.get("jobs", [])

            if not jobs:
                break

            all_jobs.extend(jobs)

            if len(jobs) < self.per_page:
                break

            page += 1

        return all_jobs

    async def _fetch_run_jobs(
        self, session: aiohttp.ClientSession, run: dict[str, Any]
    ) -> list[JobResult]:
        """Fetch jobs for a single workflow run."""
        run_id = run["id"]
        run_date = datetime.fromisoformat(run["created_at"].replace("Z", "+00:00")).strftime(
            "%m/%d"
        )

        jobs = await self.get_workflow_jobs(session, run_id)
        data_rows = []

        for job in jobs:
            conclusion = job["conclusion"]
            if self.status_to_emoji(conclusion) is None:  # Skip this job
                continue

            failed_step = next(
                (s["name"] for s in job.get("steps", []) if s.get("conclusion") == "failure"),
                None,
            )
            data_rows.append(
                JobResult(
                    name=self.parse_job_name(job["name"]),
                    conclusion=conclusion,
                    date=run_date,
                    started_at=job["started_at"],
                    completed_at=job["completed_at"],
                    failed_step=failed_step,
                    html_url=job["html_url"],
                    logs_url=f"{job['url']}/logs",
                )
            )

        return data_rows

    async def fetch_all_jobs(self, days_back: int = 30) -> list[JobResult]:
        """Fetch all jobs from workflow runs in the specified time period."""
        async with aiohttp.ClientSession() as session:
            workflow_runs = await self.get_workflow_runs(session, days_back)

            if not workflow_runs:
                return []

            print(
                f"Fetching jobs for {len(workflow_runs)} workflow runs concurrently...",
                file=sys.stderr,
            )

            tasks = [self._fetch_run_jobs(session, run) for run in workflow_runs]

            results = await asyncio.gather(*tasks, return_exceptions=True)
            data_rows: list[JobResult] = []

            for i, result in enumerate(results, 1):
                if isinstance(result, BaseException):
                    print(f"  Error fetching jobs for run {i}: {result}", file=sys.stderr)
                else:
                    data_rows.extend(result)
                    print(
                        f"  Completed {i}/{len(workflow_runs)} ({len(result)} jobs)",
                        file=sys.stderr,
                    )

            return data_rows

    def _pivot_job_results(
        self, data_rows: list[JobResult]
    ) -> tuple[dict[str, dict[str, str]], list[str], list[str]]:
        """Pivot job results data into a format suitable for table rendering.

        Args:
            data_rows: List of job results to pivot

        Returns:
            Tuple of (pivot_data, sorted_dates, sorted_names) where:
            - pivot_data: Dictionary mapping name -> date -> status
            - sorted_dates: List of dates sorted in reverse chronological order
            - sorted_names: List of test names sorted alphabetically
        """
        pivot_data: dict[str, dict[str, str]] = {}
        all_dates: set[str] = set()

        for row in data_rows:
            if row.name not in pivot_data:
                pivot_data[row.name] = {}
            # Use first occurrence for each name-date combination
            if row.date not in pivot_data[row.name]:
                emoji = self.status_to_emoji(row.conclusion)
                pivot_data[row.name][row.date] = f"[{emoji}]({row.html_url})"
            all_dates.add(row.date)

        # Sort dates in reverse order (newest first)
        sorted_dates = sorted(all_dates, reverse=True)

        # Sort names alphabetically
        sorted_names = sorted(pivot_data.keys())

        return pivot_data, sorted_dates, sorted_names

    def _build_markdown_table(
        self,
        pivot_data: dict[str, dict[str, str]],
        sorted_dates: list[str],
        sorted_names: list[str],
    ) -> str:
        """Build a markdown table from pivoted data.

        Args:
            pivot_data: Dictionary mapping name -> date -> status
            sorted_dates: List of dates (columns) in desired order
            sorted_names: List of test names (rows) in desired order

        Returns:
            Markdown-formatted table as a string
        """
        headers = ["Name"] + sorted_dates

        # Calculate column widths
        col_widths = [len(h) for h in headers]
        for name in sorted_names:
            col_widths[0] = max(col_widths[0], len(name))
            for i, date in enumerate(sorted_dates, 1):
                value = pivot_data[name].get(date, "—")
                col_widths[i] = max(col_widths[i], len(value))

        # Build table rows
        lines = []

        # Header row
        header_row = "| " + " | ".join(h.ljust(col_widths[i]) for i, h in enumerate(headers)) + " |"
        lines.append(header_row)

        # Separator row
        separator = "| " + " | ".join("-" * w for w in col_widths) + " |"
        lines.append(separator)

        # Data rows
        for name in sorted_names:
            row_values = [name.ljust(col_widths[0])]
            for i, date in enumerate(sorted_dates, 1):
                value = pivot_data[name].get(date, "—")
                row_values.append(value.ljust(col_widths[i]))
            lines.append("| " + " | ".join(row_values) + " |")

        return "\n".join(lines)

    def filter_latest_not_success(self, data_rows: list[JobResult]) -> list[JobResult]:
        """Keep only jobs whose latest run for each name was not a success."""
        latest_per_name: dict[str, JobResult] = {}
        for r in data_rows:
            cur = latest_per_name.get(r.name)
            if cur is None or r.started_at > cur.started_at:
                latest_per_name[r.name] = r
        keep = {name for name, r in latest_per_name.items() if r.conclusion != "success"}
        return [r for r in data_rows if r.name in keep]

    def render_results_table(self, data_rows: list[JobResult]) -> str:
        """Render job data as a markdown table."""
        if not data_rows:
            return "No test jobs found."

        pivot_data, sorted_dates, sorted_names = self._pivot_job_results(data_rows)
        return self._build_markdown_table(pivot_data, sorted_dates, sorted_names)

    def render_json(self, data_rows: list[JobResult]) -> str:
        return json.dumps([asdict(r) for r in data_rows], indent=2)


async def main() -> None:
    parser = argparse.ArgumentParser(description="Visualize MLflow cross-version test results")
    parser.add_argument(
        "--days", type=int, default=14, help="Number of days back to fetch results (default: 14)"
    )
    parser.add_argument(
        "--repo",
        default="mlflow/dev",
        help="GitHub repository in owner/repo format (default: mlflow/dev)",
    )
    parser.add_argument("--token", help="GitHub token (default: use GH_TOKEN env var)")
    parser.add_argument(
        "--json-output",
        help="If set, also write JSON results to this path.",
    )

    args = parser.parse_args()

    token = args.token or os.environ.get("GH_TOKEN")
    if not token:
        print(
            "Warning: No GitHub token provided. API requests may be rate-limited.", file=sys.stderr
        )
        print("Set GH_TOKEN environment variable or use --token option.", file=sys.stderr)

    visualizer = XTestViz(github_token=token, repo=args.repo)
    raw_rows = await visualizer.fetch_all_jobs(args.days)
    data_rows = visualizer.filter_latest_not_success(raw_rows)
    if args.json_output:
        Path(args.json_output).write_text(visualizer.render_json(data_rows))
    if not data_rows:
        if raw_rows:
            print("All latest cross-version tests succeeded.")
        else:
            print("No workflow runs found in the specified time period.")
    else:
        print(visualizer.render_results_table(data_rows))


if __name__ == "__main__":
    asyncio.run(main())