DynaNote/dynanote_tui.py

#!/usr/bin/env python3
"""
DynaNote TUI - 使用 SM-2 算法的间隔重复闪卡系统 (Textual 版本)
"""

import os
import sys
import toml
import tempfile
import subprocess
from typing import Dict, List, Optional, Tuple
from datetime import datetime

# 导入现有模块
import hasher
import timer
import sm2
import pathlib

from textual import on
from textual.app import App, ComposeResult
from textual.containers import Container, Horizontal, Vertical
from textual.widgets import (
    Header, Footer, Button, Static, Input,
    DataTable, Label, Select, TextArea,
    TabbedContent, TabPane, ContentSwitcher
)
from textual.screen import Screen, ModalScreen
from textual.reactive import reactive

os.chdir(pathlib.Path(__file__).resolve().parent)

class MemoryUnit:
    """表示单个记忆单元（闪卡）"""

    def __init__(self, name: str, unit_id: Optional[str] = None):
        self.name = name
        self.unit_id = unit_id or hasher.get_md5(name)
        self.attachments = set()  # 元组集合 (字符串, 创建时间戳)
        self.created_time = timer.get_daystamp()
        self.algodata = {
            sm2.SM2Algorithm.algo_name: {
                'efactor': 2.5,
                'real_rept': 0,
                'rept': 0,
                'interval': 0,
                'last_date': 0,
                'next_date': 0,
                'is_activated': 0,
                'last_modify': timer.get_timestamp()
            }
        }

    def to_dict(self) -> Dict:
        """将记忆单元转换为字典以便 TOML 序列化"""
        return {
            'name': self.name,
            'attachments': list(self.attachments),
            'created_time': self.created_time,
            'algodata': self.algodata
        }

    @classmethod
    def from_dict(cls, unit_id: str, data: Dict) -> 'MemoryUnit':
        """从字典创建记忆单元"""
        unit = cls(data['name'], unit_id)
        unit.attachments = set(tuple(att) for att in data['attachments'])
        unit.created_time = data['created_time']
        unit.algodata = data['algodata']
        return unit


class NewUnitScreen(ModalScreen):
    """创建新记忆单元的模态屏幕"""

    def compose(self) -> ComposeResult:
        with Container():
            yield Label("创建新记忆单元", classes="modal-title")
            yield Input(placeholder="输入单元名称...", id="unit-name")
            with Horizontal():
                yield Button("创建", variant="primary", id="create-btn")
                yield Button("取消", variant="default", id="cancel-btn")

    @on(Button.Pressed, "#create-btn")
    def create_unit(self) -> None:
        name_input = self.query_one("#unit-name", Input)
        name = name_input.value.strip()
        if name:
            self.dismiss(name)

    @on(Button.Pressed, "#cancel-btn")
    def cancel(self) -> None:
        self.dismiss(None)


class MarkRatingScreen(ModalScreen):
    """评分记忆单元的模态屏幕"""

    def __init__(self, unit_name: str):
        super().__init__()
        self.unit_name = unit_name

    def compose(self) -> ComposeResult:
        with Container():
            yield Label(f"为 '{self.unit_name}' 评分", classes="modal-title")
            yield Label("\nSM-2 评分标准:", classes="rating-title")
            yield Label("5 - 完美回答", classes="rating-item")
            yield Label("4 - 犹豫后正确回答", classes="rating-item")
            yield Label("3 - 经过严重困难后回忆起正确答案", classes="rating-item")
            yield Label("2 - 错误回答; 但记得正确答案", classes="rating-item")
            yield Label("1 - 错误回答; 但正确答案似乎熟悉", classes="rating-item")
            yield Label("0 - 完全遗忘", classes="rating-item")

            with Horizontal():
                for rating in [5, 4, 3, 2, 1, 0]:
                    yield Button(str(rating), id=f"rating-{rating}")
            yield Button("取消", variant="default", id="cancel-btn")

    @on(Button.Pressed)
    def handle_rating(self, event: Button.Pressed) -> None:
        if event.button.id and event.button.id.startswith("rating-"):
            rating = int(event.button.id.split("-")[1])
            self.dismiss(rating)
        elif event.button.id == "cancel-btn":
            self.dismiss(None)


class AttachmentScreen(ModalScreen):
    """添加附件的模态屏幕"""

    def compose(self) -> ComposeResult:
        with Container():
            yield Label("添加附件", classes="modal-title")
            yield TextArea(id="attachment-text", language="markdown")
            with Horizontal():
                yield Button("添加", variant="primary", id="add-btn")
                yield Button("取消", variant="default", id="cancel-btn")

    @on(Button.Pressed, "#add-btn")
    def add_attachment(self) -> None:
        text_area = self.query_one("#attachment-text", TextArea)
        text = text_area.text.strip()
        if text:
            self.dismiss(text)

    @on(Button.Pressed, "#cancel-btn")
    def cancel(self) -> None:
        self.dismiss(None)


class MainScreen(Screen):
    """主屏幕"""

    selected_unit = reactive(None)

    def __init__(self):
        super().__init__()
        self.data_file = './data.toml'
        self.memory_units: Dict[str, MemoryUnit] = {}
        self.load_data()

    def load_data(self) -> None:
        """从 data.toml 文件加载记忆单元"""
        if os.path.exists(self.data_file):
            try:
                with open(self.data_file, 'r', encoding='utf-8') as f:
                    data = toml.load(f)
                    # 处理嵌套数据
                    nested_data = dict()
                    for key, value in data.items():
                        if '.' in key:
                            parts = key.split('.')
                            current = nested_data
                            for part in parts[:-1]:
                                if part not in current:
                                    current[part] = {}
                                current = current[part]
                            current[parts[-1]] = value
                        else:
                            nested_data[key] = value
                    data = nested_data

                for unit_id, unit_data in data.items():
                    self.memory_units[unit_id] = MemoryUnit.from_dict(unit_id, unit_data)

                self.notify(f"已加载 {len(self.memory_units)} 个记忆单元", severity="information")
            except Exception as e:
                self.notify(f"加载数据时出错: {e}", severity="error")
        else:
            self.notify("未找到数据文件. 从空数据库开始.", severity="warning")

    def save_data(self) -> None:
        """将记忆单元保存到 data.toml 文件"""
        try:
            data = {}
            for unit_id, unit in self.memory_units.items():
                data[unit_id] = unit.to_dict()

            with open(self.data_file, 'w', encoding='utf-8') as f:
                toml.dump(data, f)

            self.notify("数据保存成功", severity="information")
        except Exception as e:
            self.notify(f"保存数据时出错: {e}", severity="error")

    def compose(self) -> ComposeResult:
        yield Header()
        with TabbedContent():
            with TabPane("记忆单元列表", id="list-tab"):
                with Vertical():
                    yield Label("记忆单元列表", classes="section-title")
                    with Horizontal():
                        yield Button("刷新", id="refresh-btn")
                        yield Button("新建", id="new-btn", variant="primary")
                        yield Button("删除", id="delete-btn", variant="error")
                    yield DataTable(id="units-table")

            with TabPane("单元详情", id="detail-tab"):
                with Vertical():
                    yield Label("记忆单元详情", classes="section-title")
                    yield Static("未选择任何单元", id="unit-info")
                    with Horizontal():
                        yield Button("评分", id="mark-btn", variant="success")
                        yield Button("添加附件", id="attach-btn")
                        yield Button("编辑", id="edit-btn")
                    yield Static("", id="attachments-info")

        yield Footer()

    def on_mount(self) -> None:
        """挂载时初始化"""
        self.setup_table()
        self.update_unit_info()

    def on_resume(self) -> None:
        """从模态屏幕返回时刷新数据"""
        self.refresh_table()

    def setup_table(self) -> None:
        """设置数据表格"""
        table = self.query_one("#units-table", DataTable)

        # 只在表格为空时添加列头，避免重复添加
        if not table.columns:
            table.add_columns(
                "ID", "名称", "复习次数", "EF", "下次复习时间", "附件数"
            )

        # 清除现有行数据
        table.clear()

        today = timer.get_daystamp()
        for unit in self.memory_units.values():
            algodata = unit.algodata[sm2.SM2Algorithm.algo_name]
            table.add_row(
                unit.unit_id[:8] + "...",
                unit.name,
                str(algodata['real_rept']),
                f"{algodata['efactor']:.2f}",
                str(algodata['next_date']),
                str(len(unit.attachments))
            )

    def update_unit_info(self) -> None:
        """更新单元详情显示"""
        info_widget = self.query_one("#unit-info", Static)
        attachments_widget = self.query_one("#attachments-info", Static)

        if self.selected_unit:
            unit = self.selected_unit
            algodata = unit.algodata[sm2.SM2Algorithm.algo_name]

            info_text = f"""
记忆单元: {unit.name}
ID: {unit.unit_id}
创建时间: {unit.created_time}

SM-2 算法数据:
  EFactor: {algodata['efactor']}
  实际复习次数: {algodata['real_rept']}
  当前重复次数: {algodata['rept']}
  间隔: {algodata['interval']} 天
  上次复习: {algodata['last_date']}
  下次复习: {algodata['next_date']}
  是否激活: {algodata['is_activated']}
"""
            info_widget.update(info_text)

            if unit.attachments:
                attachments_text = f"\n附件 ({len(unit.attachments)}):\n"
                sorted_attachments = sorted(unit.attachments, key=lambda x: float(x[1]))
                for i, (text, timestamp) in enumerate(sorted_attachments, 1):
                    attachments_text += f"  {i}. {text[:80]}{'...' if len(text) > 80 else ''} (于 {timestamp})\n"
                attachments_widget.update(attachments_text)
            else:
                attachments_widget.update("\n无附件")
        else:
            info_widget.update("未选择任何单元")
            attachments_widget.update("")

    @on(DataTable.RowSelected, "#units-table")
    def on_row_selected(self, event: DataTable.RowSelected) -> None:
        """处理表格行选择"""
        if event.row_key:
            unit_id_prefix = event.row_key.value
            for unit_id, unit in self.memory_units.items():
                if unit_id.startswith(unit_id_prefix):
                    self.selected_unit = unit
                    self.update_unit_info()
                    self.notify(f"已选择单元: {unit.name}", severity="information")
                    break

    @on(Button.Pressed, "#refresh-btn")
    def refresh_table(self) -> None:
        """刷新表格"""
        self.load_data()  # 重新加载数据
        self.setup_table()
        self.update_unit_info()
        self.notify("表格已刷新", severity="information")

    @on(Button.Pressed, "#new-btn")
    async def new_unit(self) -> None:
        """创建新单元"""
        name = await self.app.push_screen_wait(NewUnitScreen())
        if name:
            unit = MemoryUnit(name)
            self.memory_units[unit.unit_id] = unit
            self.selected_unit = unit
            self.setup_table()
            self.update_unit_info()
            self.save_data()
            self.notify(f"已创建新单元: {unit.name}", severity="success")

    @on(Button.Pressed, "#delete-btn")
    def delete_unit(self) -> None:
        """删除选中的单元"""
        if self.selected_unit:
            unit_name = self.selected_unit.name
            unit_id = self.selected_unit.unit_id
            del self.memory_units[unit_id]
            self.selected_unit = None
            self.setup_table()
            self.update_unit_info()
            self.save_data()
            self.notify(f"已删除单元: {unit_name}", severity="warning")
        else:
            self.notify("请先选择一个单元", severity="error")

    @on(Button.Pressed, "#mark-btn")
    async def mark_unit(self) -> None:
        """为选中的单元评分"""
        if not self.selected_unit:
            self.notify("请先选择一个单元", severity="error")
            return

        rating = await self.app.push_screen_wait(
            MarkRatingScreen(self.selected_unit.name)
        )

        if rating is not None:
            algodata = self.selected_unit.algodata
            if algodata[sm2.SM2Algorithm.algo_name]['is_activated'] == 0:
                sm2.SM2Algorithm.revisor(algodata, rating, is_new_activation=1)
                algodata[sm2.SM2Algorithm.algo_name]['is_activated'] = 1
            else:
                sm2.SM2Algorithm.revisor(algodata, rating)

            self.update_unit_info()
            self.save_data()
            self.notify(f"已使用评分 {rating} 更新 {self.selected_unit.name} 的 SM-2 参数",
                       severity="success")

    @on(Button.Pressed, "#attach-btn")
    async def add_attachment(self) -> None:
        """为选中的单元添加附件"""
        if not self.selected_unit:
            self.notify("请先选择一个单元", severity="error")
            return

        text = await self.app.push_screen_wait(AttachmentScreen())
        if text:
            attachment = (text, str(timer.get_timestamp()))
            self.selected_unit.attachments.add(attachment)
            self.update_unit_info()
            self.save_data()
            self.notify(f"附件已添加到 {self.selected_unit.name}", severity="success")

    @on(Button.Pressed, "#edit-btn")
    def edit_unit(self) -> None:
        """编辑选中的单元"""
        if not self.selected_unit:
            self.notify("请先选择一个单元", severity="error")
            return

        # 创建包含单元数据的临时文件
        with tempfile.NamedTemporaryFile(mode='w', suffix='.toml', delete=False) as f:
            temp_file = f.name
            toml.dump({self.selected_unit.unit_id: self.selected_unit.to_dict()}, f)

        try:
            # 启动 nano 编辑器
            subprocess.run(['nano', temp_file], check=True)

            # 读取编辑后的文件
            with open(temp_file, 'r', encoding='utf-8') as f:
                edited_data = toml.load(f)

            # 更新记忆单元
            if self.selected_unit.unit_id in edited_data:
                updated_data = edited_data[self.selected_unit.unit_id]
                self.memory_units[self.selected_unit.unit_id] = MemoryUnit.from_dict(
                    self.selected_unit.unit_id, updated_data
                )
                self.selected_unit = self.memory_units[self.selected_unit.unit_id]
                self.update_unit_info()
                self.save_data()
                self.notify("单元更新成功", severity="success")
            else:
                self.notify("错误: 在编辑后的文件中未找到单元 ID", severity="error")

        except subprocess.CalledProcessError:
            self.notify("编辑器已关闭但未保存", severity="warning")
        except Exception as e:
            self.notify(f"编辑单元时出错: {e}", severity="error")
        finally:
            # 清理临时文件
            if os.path.exists(temp_file):
                os.unlink(temp_file)


class DynaNoteTUI(App):
    """DynaNote TUI 应用"""

    CSS = """
    .modal-title {
        text-align: center;
        text-style: bold;
        margin: 1;
    }

    .section-title {
        text-align: center;
        text-style: bold;
        margin: 1;
    }

    .rating-title {
        text-style: bold;
        margin: 1 0;
    }

    .rating-item {
        margin: 0 0 0 2;
    }

    #unit-info {
        margin: 1;
        border: solid $primary;
        padding: 1;
    }

    #attachments-info {
        margin: 1;
        border: solid $accent;
        padding: 1;
    }

    DataTable {
        height: 1fr;
    }
    """

    BINDINGS = [
        ("q", "quit", "退出"),
        ("s", "save", "保存数据"),
        ("r", "refresh", "刷新"),
        ("n", "new", "新建单元"),
    ]

    def on_mount(self) -> None:
        """挂载时切换到主屏幕"""
        self.push_screen(MainScreen())

    def action_quit(self) -> None:
        """退出应用"""
        self.exit()

    def action_save(self) -> None:
        """保存数据"""
        if hasattr(self.screen, 'save_data'):
            self.screen.save_data()

    def action_refresh(self) -> None:
        """刷新"""
        if hasattr(self.screen, 'refresh_table'):
            self.screen.refresh_table()

    def action_new(self) -> None:
        """新建单元"""
        if hasattr(self.screen, 'new_unit'):
            self.screen.new_unit()


def main():
    """主入口点"""
    app = DynaNoteTUI()
    app.run()


if __name__ == '__main__':
    main()