#!/usr/bin/env python3
import argparse
import os
import re
import subprocess
from datetime import date, datetime, timedelta
from pathlib import Path


H1_RE = re.compile(r"^# (?!#)")
DATE_H1_RE = re.compile(r"^# (\d{4}-\d{2}-\d{2})\s*$")
CYCLE_RE = re.compile(r"^# Cycle \d+ \((\d{2}/\d{2}/\d{4}) - (\d{2}/\d{2}/\d{4})\)\s*$")
CHECK_RE = re.compile(r"^- \[( |x|X)\] (.+)$")
METRIC_RE = re.compile(r"^-?\s*([^:]+):\s*(.*)$")


# Edit this markdown to change the weekly review template appended to each file.
REVIEW_TEMPLATE = """## Accomplishments

## Disapointments 

## Highs

## Lows

## What I learned

"""


def parse_iso(s: str) -> date:
    return datetime.strptime(s, "%Y-%m-%d").date()


def parse_us(s: str) -> date:
    return datetime.strptime(s, "%m/%d/%Y").date()


def fmt_us(d: date) -> str:
    return d.strftime("%m/%d/%Y")


def fmt_file(d: date) -> str:
    return d.strftime("%m-%d-%Y")


def last_completed_sun_sat(today: date) -> tuple[date, date]:
    days_since_saturday = (today.weekday() + 2) % 7
    end = today - timedelta(days=days_since_saturday)
    start = end - timedelta(days=6)
    return start, end


def dates_between(start: date, end: date) -> list[date]:
    days = []
    current = start
    while current <= end:
        days.append(current)
        current += timedelta(days=1)
    return days


def section_lines(lines: list[str], header_index: int) -> list[str]:
    end = len(lines)
    for i in range(header_index + 1, len(lines)):
        if H1_RE.match(lines[i]):
            end = i
            break
    return lines[header_index + 1 : end]


def parse_habits(path: Path, start: date, end: date):
    entries = {d: {"habits": {}, "metrics": {}} for d in dates_between(start, end)}
    habit_names = []
    metric_names = []

    if not path.exists():
        return entries, habit_names, metric_names

    lines = path.read_text().splitlines()
    for i, line in enumerate(lines):
        m = DATE_H1_RE.match(line)
        if not m:
            continue

        day = parse_iso(m.group(1))
        if day < start or day > end:
            continue

        current_section = None
        for entry_line in section_lines(lines, i):
            stripped = entry_line.strip()
            if stripped == "## Habits":
                current_section = "habits"
                continue
            if stripped == "## Metrics":
                current_section = "metrics"
                continue
            if stripped.startswith("## "):
                current_section = None
                continue

            if current_section == "habits":
                hm = CHECK_RE.match(stripped)
                if hm:
                    name = hm.group(2).strip()
                    entries[day]["habits"][name] = hm.group(1).lower() == "x"
                    if name not in habit_names:
                        habit_names.append(name)
            elif current_section == "metrics":
                mm = METRIC_RE.match(stripped)
                if mm:
                    name = mm.group(1).strip()
                    value = mm.group(2).strip()
                    entries[day]["metrics"][name] = value
                    if name not in metric_names:
                        metric_names.append(name)

    return entries, habit_names, metric_names


def extract_done_tasks(path: Path, start: date, end: date) -> list[str]:
    if not path.exists():
        return []

    lines = path.read_text().splitlines()
    for i, line in enumerate(lines):
        m = CYCLE_RE.match(line)
        if not m:
            continue
        if parse_us(m.group(1)) == start and parse_us(m.group(2)) == end:
            body = section_lines(lines, i)
            while body and not body[0].strip():
                body.pop(0)
            while body and not body[-1].strip():
                body.pop()
            if body and body[0].strip() == "## Done":
                body = body[1:]
                while body and not body[0].strip():
                    body.pop(0)
            return body

    return []


def markdown_table(headers: list[str], rows: list[list[str]]) -> str:
    table = [headers] + rows
    widths = [max(len(str(row[i])) for row in table) for i in range(len(headers))]

    def fmt_row(row: list[str]) -> str:
        cells = [str(cell).ljust(widths[i]) for i, cell in enumerate(row)]
        return "| " + " | ".join(cells) + " |"

    separator = ["-" * width for width in widths]
    return (
        "\n".join(
            [fmt_row(headers), fmt_row(separator), *[fmt_row(row) for row in rows]]
        )
        + "\n"
    )


def build_habit_chart(entries, habit_names: list[str], start: date, end: date) -> str:
    days = dates_between(start, end)
    headers = ["Habit"] + [d.strftime("%a %m/%d") for d in days] + ["Total"]
    rows = []
    for habit in habit_names:
        total = 0
        row = [habit]
        for day in days:
            value = entries[day]["habits"].get(habit)
            if value is True:
                row.append("✓")
                total += 1
            elif value is False:
                row.append("—")
            else:
                row.append("")
        row.append(f"{total}/7")
        rows.append(row)

    if not rows:
        return "No habit entries found for this week.\n"
    return markdown_table(headers, rows)


def build_metrics_chart(
    entries, metric_names: list[str], start: date, end: date
) -> str:
    days = dates_between(start, end)
    headers = ["Metric"] + [d.strftime("%a %m/%d") for d in days]
    rows = []
    for metric in metric_names:
        rows.append(
            [metric] + [entries[day]["metrics"].get(metric, "") for day in days]
        )

    if not rows:
        return "No metrics found for this week.\n"
    return markdown_table(headers, rows)


def load_review_markdown() -> str:
    return REVIEW_TEMPLATE.rstrip() + "\n"


def main():
    parser = argparse.ArgumentParser(
        description="Create a weekly review markdown file."
    )
    parser.add_argument("--today", help="Override today, format YYYY-MM-DD")
    parser.add_argument("--start", help="Override week start, format YYYY-MM-DD")
    parser.add_argument("--end", help="Override week end, format YYYY-MM-DD")
    parser.add_argument(
        "--force", action="store_true", help="Overwrite an existing review file"
    )
    parser.add_argument(
        "--no-open", action="store_true", help="Do not open the review in $EDITOR"
    )
    parser.add_argument("--output-dir", default="~/.local/share/os/weeks")
    parser.add_argument("--habits-file", default="~/.local/share/os/habits.md")
    parser.add_argument("--tasks-file", default="~/.local/share/os/tasks.md")
    args = parser.parse_args()

    today = parse_iso(args.today) if args.today else date.today()
    if args.start or args.end:
        if not args.start or not args.end:
            raise SystemExit("Use --start and --end together")
        start = parse_iso(args.start)
        end = parse_iso(args.end)
    else:
        start, end = last_completed_sun_sat(today)

    output_dir = Path(args.output_dir).expanduser()
    output_path = output_dir / f"{fmt_file(start)}--{fmt_file(end)}.md"
    if output_path.exists() and not args.force:
        print(f"Review already exists: {output_path}")
        return

    habits_file = Path(args.habits_file).expanduser()
    tasks_file = Path(args.tasks_file).expanduser()

    habit_entries, habit_names, metric_names = parse_habits(habits_file, start, end)
    done_tasks = extract_done_tasks(tasks_file, start, end)
    review_markdown = load_review_markdown()

    content = []
    content.append(f"# Weekly Review ({fmt_us(start)} - {fmt_us(end)})\n\n")
    content.append("## Habits\n\n")
    content.append(build_habit_chart(habit_entries, habit_names, start, end))
    content.append("\n## Metrics\n\n")
    content.append(build_metrics_chart(habit_entries, metric_names, start, end))
    content.append("\n## Done Tasks\n\n")
    if done_tasks:
        content.append("\n".join(done_tasks).rstrip() + "\n")
    else:
        content.append("No done tasks found for this week.\n")
    content.append("\n")
    content.append(review_markdown)

    output_dir.mkdir(parents=True, exist_ok=True)
    output_path.write_text("".join(content))
    print(f"Created {output_path}")

    if not args.no_open:
        subprocess.run([os.environ.get("EDITOR", "vi"), str(output_path)], check=True)


if __name__ == "__main__":
    main()
