#!/usr/bin/env -S uv run --script
from datetime import date, datetime, timedelta
from peewee import (
    Model,
    SqliteDatabase,
    TextField,
    DateTimeField,
    FloatField,
    IntegerField,
    DateField,
    fn,
)
from rich.console import Console
from rich.markdown import Markdown
from typing import Optional
from xdg_base_dirs import xdg_config_home, xdg_data_home

import argparse
import subprocess

console = Console(style="#f8f8f8")

db_dir = xdg_data_home() / "memories"
db_dir.mkdir(parents=True, exist_ok=True)
db_file = db_dir / "memories.db"
db = SqliteDatabase(db_file)


class BaseModel(Model):
    class Meta:
        database = db


class Memory(BaseModel):
    text = TextField(index=True, unique=True)
    created = DateTimeField(default=datetime.now)
    last_reviewed = DateTimeField(default=datetime.now)
    ease = FloatField(default=2.5)
    interval = IntegerField(default=1)
    repetitions = IntegerField(default=0)
    due = DateField(default=date.today, index=True)

    def reschedule(self, score: int):
        if score < 3:
            self.repetitions = 0
            self.interval = 1
        else:
            self.repetitions += 1
            if self.repetitions == 1:
                self.interval = 1
            elif self.repetitions == 2:
                self.interval = 6
            else:
                self.interval = round(self.interval * self.ease)
            self.ease += 0.1 - (5 - score) * (0.08 + (5 - score) * 0.02)
            self.ease = max(1.3, self.ease)
        self.due = date.today() + timedelta(days=self.interval)
        self.save()

    def get_score(self):
        while True:
            try:
                response = input("\033[1mRate 0-5 (5=perfect, 4=hesitation, 3=slow, 2=wrongish, 1=wrong, 0=forgot):\033[m ")
            except (Exception, KeyboardInterrupt):
                print()
                return

            if response == "" or response == "q":
                exit(0)
            elif response == "d":
                self.delete_instance()
                return

            try:
                if response == "":
                    score = 5
                else:
                    score = int(response)
            except ValueError:
                print("Not a number! Try again.")
                continue
            if score < 1 or score > 5:
                print("Not in the range 1-5! Try again.")
                continue

            self.reschedule(score)
            return

    def print(self):
        renderable_markup = Markdown(self.text)
        self.last_reviewed = datetime.now()
        self.save()
        console.print(renderable_markup)


# Setup
db.connect()
db.create_tables([Memory])

def main():
    import sys
    args = sys.argv[1:]

    memories_dir = xdg_config_home() / "memories"
    memories_dir.mkdir(parents=True, exist_ok=True)
    memories_file = memories_dir / "memories.md"

    parser = argparse.ArgumentParser(prog="memorize")
    parser.add_argument("--ask-score", type=bool, default=False, help="Ask user for score")
    sub = parser.add_subparsers(dest="command")

    sub.add_parser("review")
    sub.add_parser("edit")

    args = parser.parse_args()
    if args.command == "edit":
        edit(memories_file)
    else:
        review(memories_file, ask_score=args.ask_score)

def edit(memories_file:"Path"):
    subprocess.call(["vim", memories_file])
    update_database(memories_file)

def update_database(memories_file:"Path"):
    if memories_file.exists():
        new_memories = 0
        with memories_file.open() as memories:
            valid_texts = {text.strip() for text in memories.read().split("\n---\n")}
            for text in valid_texts:
                memory, created = Memory.get_or_create(text=text)
                if created:
                    new_memories += 1

            deleted = 0
            for memory in Memory.select():
                if memory.text not in valid_texts:
                    memory.delete_instance()
                    deleted += 1

            if new_memories > 0 or deleted > 0:
                print(f"\033[2;3m{new_memories} new memories, {deleted} deleted\033[m\n")

def review(memories_file:"Path", ask_score:bool = False):
    update_database(memories_file)

    memory = (
        Memory.select()
        .where(Memory.due <= date.today())
        .order_by(fn.Random())
        .get_or_none()
    )
    if memory is None:
        memory = Memory.select().order_by(Memory.last_reviewed).get_or_none()

    if memory is None:
        print("No memories!")
        exit(1)

    memory.print()
    if ask_score:
        print()
        memory.get_score()
    else:
        memory.reschedule(4)

if __name__ == "__main__":
    main()