#!/usr/bin/env cached-nix-shell
#! nix-shell -i python3
#! nix-shell -p python3 python3Packages.rich python3Packages.typer aoc-cli
# vim: filetype=python

"""
Display Advent of Code puzzle solving times based on web browser history.

Assumptions/Limitations:
- You opened your and submitted your inputs with the qutebrowser web browser
- The last attempt you made at a puzzle was the correct one
- aoc-cli is configured
- There will be a YYYY/MM folder relative to $PWD if you attempted the puzzle
- This script will add a prompt file in above directory (to cache the titles)
"""

import collections
import csv
import datetime
import enum
import os
import pathlib
import re
import sqlite3
import statistics
import subprocess
import typing
import urllib.parse
import zoneinfo

import rich.console
import rich.table
import typer

app = typer.Typer()


def get_qutebrowser_urls() -> (
    typing.Generator[tuple[datetime.datetime, urllib.parse.ParseResult], None, None]
):
    qutebrowser_history_path = os.path.expanduser(
        "~/.local/share/qutebrowser/history.sqlite"
    )
    db = sqlite3.connect(qutebrowser_history_path)

    cur = db.cursor()
    for row in cur.execute(
        "SELECT atime,url FROM History"
        " WHERE url LIKE 'https://adventofcode.com/%/day/%'"
        " ORDER BY atime ASC"
    ):
        tim = datetime.datetime.fromtimestamp(row[0]).astimezone()
        url = urllib.parse.urlparse(row[1])
        yield tim, url


# opened, p1 attempts, p1 last answer, p2 attempts, p2 last answer
PuzzleTimesAttempts = tuple[
    datetime.datetime, int, datetime.datetime | None, int, datetime.datetime | None
]


def get_puzzle_folder(date: datetime.date) -> pathlib.Path | None:
    path = pathlib.Path(f"{date.year}/{date.day}")
    if path.is_dir():
        return path
    else:
        return None


def get_time_attempts(year: int | None) -> dict[datetime.date, PuzzleTimesAttempts]:
    starts: dict[datetime.date, datetime.datetime] = dict()
    part1a: dict[datetime.date, int] = collections.Counter()
    part1t: dict[datetime.date, datetime.datetime] = dict()
    got_part2: set[datetime.date] = set()
    part2a: dict[datetime.date, int] = collections.Counter()
    part2t: dict[datetime.date, datetime.datetime] = dict()

    for tim, url in get_qutebrowser_urls():
        match = re.match(
            r"^/(?P<year>[0-9]{4})/day/(?P<day>[0-9]{1,2})(?P<answer>/answer)?$",
            url.path,
        )
        if not match:
            continue  # Not an interesting URL

        date = datetime.date(int(match["year"]), 12, int(match["day"]))
        if year is not None and date.year != year:
            continue  # Not the wanted year

        answer = bool(match["answer"])
        part2 = url.fragment == "part2"

        if date not in starts:
            starts[date] = tim

        if answer:
            if date in got_part2:
                part2a[date] += 1
                part2t[date] = tim
            else:
                part1a[date] += 1
                part1t[date] = tim
        elif part2:
            got_part2.add(date)

    time_attempts = dict()
    for date in starts:
        if get_puzzle_folder(date) is None:
            continue  # Not attempted

        time_attempts[date] = (
            starts[date],
            part1a.get(date, 0),
            part1t.get(date),
            part2a.get(date, 0),
            part2t.get(date),
        )
    return time_attempts


def get_puzzle_prompt(date: datetime.date) -> str:
    path = get_puzzle_folder(date)
    assert path
    filepath = path / "prompt"
    if not filepath.is_file():
        r = subprocess.run(
            ["aoc", "read", "--year", str(date.year), "--day", str(date.day)],
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
        )
        r.check_returncode()
        with open(filepath, "w") as fd:
            fd.write(r.stdout.decode())
    with open(filepath) as fd:
        return fd.read()


def get_puzzle_title(date: datetime.date) -> str:
    prompt = get_puzzle_prompt(date)
    for line in prompt.splitlines():
        match = re.match(f"^--- Day {date.day}: (.+) ---$", line)
        if match:
            return match[1]
    raise ValueError("Could not find title in prompt")


def bold_time_threshold(times: list[datetime.timedelta]) -> datetime.timedelta:
    times_s = [round(t.total_seconds(), 1) for t in times]
    quantiles = statistics.quantiles(times_s)
    return datetime.timedelta(seconds=quantiles[-1])


def format_puzzle(date: datetime.date) -> str:
    string = str(date)
    if date.weekday() in (5, 6):
        string = f"[bold]{string}"
    return string


def format_time(
    time: datetime.timedelta | None,
    bold_threshold: datetime.timedelta,
    compare: datetime.timedelta | None = None,
) -> str:
    if time is None:
        return "[bright_black]-"
    string = str(time)
    if time >= bold_threshold:
        string = f"[underline]{string}"
    if compare and time > compare:
        string = f"[bold]{string}"
    return string


def format_attempts(
    attempts: int,
) -> str:
    if attempts == 0:
        return "[bright_black]-"
    string = str(attempts)
    if attempts > 1:
        string = f"[underline]{string}"
    return string


class SortMethod(str, enum.Enum):
    puzzle = "puzzle"
    opened = "opened"
    part1_time = "part1_time"
    part2_time = "part2_time"
    total_time = "total_time"


@app.command("csv")
def write_csv(file: typer.FileTextWrite, year: int | None = None) -> None:
    writer = csv.writer(file)
    writer.writerow(
        (
            "puzzle",
            "opened",
            "part1_attempts",
            "part1_finished",
            "part2_attempts",
            "part2_finished",
        )
    )
    for puzzle, times_attempts in get_time_attempts(year).items():
        all_info = (puzzle,) + times_attempts
        writer.writerow(str(i) for i in all_info)


@app.command("table")
def print_table(
    year: int | None = None,
    assume_early_bird: bool = False,
    show_opened: bool = True,
    sort: SortMethod = SortMethod.opened,
) -> None:

    if assume_early_bird:
        show_opened = False

    puzzles = get_time_attempts(year)
    part1_times = dict()
    part2_times = dict()
    total_times = dict()
    for date, time_attempts in puzzles.items():
        start, _, part1_date, _, part2_date = time_attempts
        if part1_date:
            if assume_early_bird:
                start = datetime.datetime(
                    date.year, date.month, date.day, tzinfo=zoneinfo.ZoneInfo("EST")
                )
            part1_times[date] = part1_date - start
            if part2_date:
                part2_times[date] = part2_date - part1_date
                total_times[date] = part2_date - start

    part1_bold_time_threshold = bold_time_threshold(list(part1_times.values()))
    part2_bold_time_threshold = bold_time_threshold(list(part2_times.values()))
    total_bold_time_threshold = bold_time_threshold(list(total_times.values()))

    part1_total = sum(part1_times.values(), start=datetime.timedelta())
    part1_attempts = sum(puzzle[1] for puzzle in puzzles.values())
    part2_total = sum(part2_times.values(), start=datetime.timedelta())
    part2_attempts = sum(puzzle[3] for puzzle in puzzles.values())
    total_total = sum(total_times.values(), start=datetime.timedelta())

    table = rich.table.Table(
        title="Advent of Code times",
        caption="Bold puzzle date: weekend; "
        "Bold time: worst part; "
        "Underline time: worst 25%; "
        "Underline attempts: >1",
        show_footer=True,
    )
    table.add_column("Puzzle date", "Total", style="blue")
    table.add_column("Puzzle title", str(len(puzzles)), style="blue")
    if show_opened:
        table.add_column("Opened date", style="magenta")
    table.add_column("P1 time", str(part1_total), style="green")
    table.add_column("P1 attempts", str(part1_attempts), style="cyan")
    table.add_column("P2 time", str(part2_total), style="green")
    table.add_column("P2 attempts", str(part2_attempts), style="cyan")
    table.add_column("Total time", str(total_total), style="yellow")

    dates = list(puzzles.keys())
    if sort == SortMethod.puzzle:
        dates.sort()
    elif sort == SortMethod.opened:
        pass
    elif sort == SortMethod.part1_time:
        dates.sort(key=lambda d: part1_times.get(d, datetime.timedelta.max))
    elif sort == SortMethod.part2_time:
        dates.sort(key=lambda d: part2_times.get(d, datetime.timedelta.max))
    elif sort == SortMethod.total_time:
        dates.sort(key=lambda d: total_times.get(d, datetime.timedelta.max))
    else:
        assert False

    for date in dates:
        start, part1_attempts, _, part2_attempts, _ = puzzles[date]
        part1_time = part1_times.get(date)
        part2_time = part2_times.get(date)
        total_time = total_times.get(date)
        row = (
            (
                format_puzzle(date),
                get_puzzle_title(date),
            )
            + ((str(start),) if show_opened else tuple())
            + (
                format_time(part1_time, part1_bold_time_threshold, part2_time),
                format_attempts(part1_attempts),
                format_time(part2_time, part2_bold_time_threshold, part1_time),
                format_attempts(part2_attempts),
                format_time(total_time, total_bold_time_threshold),
            )
        )
        table.add_row(*row)

    console = rich.console.Console()
    console.print(table)


if __name__ == "__main__":
    app()