app.py

import os
import secrets
import sqlite3
from calendar import monthrange
from contextlib import closing
from datetime import date, datetime, timedelta
from urllib.parse import urlparse
from zoneinfo import ZoneInfo, ZoneInfoNotFoundError

from flask import Flask, g, has_request_context, redirect, render_template, request, session, url_for
from werkzeug.security import check_password_hash, generate_password_hash

VALID_TABS = {"register", "recurring"}
VALID_TX_TYPES = {"debit", "credit"}
VALID_FREQUENCIES = {"weekly", "biweekly", "monthly", "quarterly", "yearly"}
DATE_FORMAT = "%Y-%m-%d"
AUTH_PASSWORD_KEY = "auth_password_hash"
RECURRING_PROJECTION_MONTHS = 2


def resolve_secret_key() -> str:
    env_key = os.environ.get("CHECKBOOK_SECRET_KEY")
    if env_key:
        return env_key

    secret_file = os.environ.get("CHECKBOOK_SECRET_FILE", "/data/.checkbook_secret_key")
    secret_dir = os.path.dirname(secret_file)

    try:
        if secret_dir:
            os.makedirs(secret_dir, exist_ok=True)

        if os.path.exists(secret_file):
            with open(secret_file, "r", encoding="utf-8") as file_obj:
                existing = file_obj.read().strip()
            if existing:
                return existing

        generated = secrets.token_hex(32)
        with open(secret_file, "w", encoding="utf-8") as file_obj:
            file_obj.write(generated)
        return generated
    except OSError:
        return "change-me-checkbook-secret-key"


app = Flask(__name__)
app.config["DATABASE"] = os.environ.get("CHECKBOOK_DB", "/data/checkbook.db")
app.config["SECRET_KEY"] = resolve_secret_key()
app.config["SESSION_COOKIE_HTTPONLY"] = True
app.config["SESSION_COOKIE_SAMESITE"] = "Lax"
app.config["PERMANENT_SESSION_LIFETIME"] = timedelta(days=14)
app.config["SESSION_REFRESH_EACH_REQUEST"] = False


def get_db() -> sqlite3.Connection:
    if "db" not in g:
        db_path = app.config["DATABASE"]
        db_dir = os.path.dirname(db_path)
        if db_dir:
            os.makedirs(db_dir, exist_ok=True)

        g.db = sqlite3.connect(db_path)
        g.db.row_factory = sqlite3.Row

    return g.db


def app_now() -> datetime:
    tz_name = ""
    if has_request_context():
        tz_name = request.cookies.get("checkbook_tz", "").strip()
    if not tz_name:
        tz_name = os.environ.get("CHECKBOOK_TIMEZONE", "").strip()
    if tz_name:
        try:
            return datetime.now(ZoneInfo(tz_name))
        except ZoneInfoNotFoundError:
            pass
    return datetime.now().astimezone()


def today_iso() -> str:
    return app_now().strftime(DATE_FORMAT)


def parse_date(value: str) -> str | None:
    cleaned = value.strip()
    if not cleaned:
        return None

    try:
        datetime.strptime(cleaned, DATE_FORMAT)
    except ValueError:
        return None

    return cleaned


def parse_optional_int(value: str | None) -> int | None:
    if not value:
        return None

    try:
        return int(value)
    except ValueError:
        return None


def get_setting(key: str) -> str | None:
    row = get_db().execute("SELECT value FROM app_settings WHERE key = ?", (key,)).fetchone()
    if row is None:
        return None
    return row["value"]


def set_setting(key: str, value: str) -> None:
    get_db().execute(
        """
        INSERT INTO app_settings (key, value)
        VALUES (?, ?)
        ON CONFLICT(key) DO UPDATE SET value = excluded.value
        """,
        (key, value),
    )
    get_db().commit()


def has_password() -> bool:
    return bool(get_setting(AUTH_PASSWORD_KEY))


def is_safe_path(path_value: str) -> bool:
    if not path_value:
        return False
    parsed = urlparse(path_value)
    return parsed.scheme == "" and parsed.netloc == "" and path_value.startswith("/")


def add_months(base_date: date, months: int) -> date:
    month_idx = base_date.month - 1 + months
    year = base_date.year + month_idx // 12
    month = month_idx % 12 + 1
    day = min(base_date.day, monthrange(year, month)[1])
    return date(year, month, day)


def increment_due_date(base_date: date, frequency: str) -> date:
    if frequency == "weekly":
        return base_date + timedelta(days=7)
    if frequency == "biweekly":
        return base_date + timedelta(days=14)
    if frequency == "monthly":
        return add_months(base_date, 1)
    if frequency == "quarterly":
        return add_months(base_date, 3)
    if frequency == "yearly":
        return add_months(base_date, 12)
    return base_date + timedelta(days=30)


def next_due_date_after_today(current_due_date: str, frequency: str) -> str:
    try:
        next_due = datetime.strptime(current_due_date, DATE_FORMAT).date()
    except ValueError:
        next_due = app_now().date()

    today = app_now().date()
    next_due = increment_due_date(next_due, frequency)
    while next_due <= today:
        next_due = increment_due_date(next_due, frequency)
    return next_due.strftime(DATE_FORMAT)


def project_recurring_entries(
    recurring_rows: list[sqlite3.Row], months_ahead: int = RECURRING_PROJECTION_MONTHS
) -> tuple[list[dict], str, str, float]:
    window_start = app_now().date()
    window_end = add_months(window_start, months_ahead)
    projected_entries: list[dict] = []

    for row in recurring_rows:
        try:
            due_date = datetime.strptime(row["next_due_date"], DATE_FORMAT).date()
        except ValueError:
            due_date = window_start

        while due_date < window_start:
            due_date = increment_due_date(due_date, row["frequency"])

        while due_date <= window_end:
            signed_amount = row["amount"] if row["tx_type"] == "credit" else -row["amount"]
            projected_entries.append(
                {
                    "id": row["id"],
                    "due_date": due_date.strftime(DATE_FORMAT),
                    "description": row["description"],
                    "category": row["category"],
                    "frequency": row["frequency"],
                    "tx_type": row["tx_type"],
                    "amount": row["amount"],
                    "signed_amount": signed_amount,
                }
            )
            due_date = increment_due_date(due_date, row["frequency"])

    projected_entries.sort(key=lambda entry: (entry["due_date"], entry["id"]))
    net_total = sum(entry["signed_amount"] for entry in projected_entries)

    return (
        projected_entries,
        window_start.strftime(DATE_FORMAT),
        window_end.strftime(DATE_FORMAT),
        net_total,
    )


def init_db() -> None:
    db = get_db()
    with closing(db.cursor()) as cursor:
        cursor.execute(
            """
            CREATE TABLE IF NOT EXISTS app_settings (
                key TEXT PRIMARY KEY,
                value TEXT NOT NULL
            )
            """
        )

        cursor.execute(
            """
            CREATE TABLE IF NOT EXISTS transactions (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                tx_date TEXT NOT NULL,
                description TEXT NOT NULL,
                category TEXT NOT NULL,
                amount REAL NOT NULL,
                tx_type TEXT NOT NULL CHECK(tx_type IN ('debit', 'credit')),
                created_at TEXT NOT NULL
            )
            """
        )

        cursor.execute(
            """
            CREATE TABLE IF NOT EXISTS recurring_transactions (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                next_due_date TEXT NOT NULL,
                description TEXT NOT NULL,
                category TEXT NOT NULL,
                amount REAL NOT NULL,
                tx_type TEXT NOT NULL CHECK(tx_type IN ('debit', 'credit')),
                frequency TEXT NOT NULL,
                is_completed INTEGER NOT NULL DEFAULT 0,
                created_at TEXT NOT NULL
            )
            """
        )

        cursor.execute("PRAGMA table_info(recurring_transactions)")
        columns = {row["name"] for row in cursor.fetchall()}
        if "next_due_date" not in columns:
            cursor.execute("ALTER TABLE recurring_transactions ADD COLUMN next_due_date TEXT")

        cursor.execute(
            """
            UPDATE recurring_transactions
            SET next_due_date = ?
            WHERE next_due_date IS NULL OR TRIM(next_due_date) = ''
            """,
            (today_iso(),),
        )

        db.commit()


@app.before_request
def before_request() -> None:
    init_db()
    endpoint = request.endpoint
    if endpoint is None or endpoint == "static":
        return

    password_is_set = has_password()

    if not password_is_set:
        if endpoint != "setup":
            return redirect(url_for("setup"))
        return

    if endpoint == "setup":
        return redirect(url_for("index") if session.get("authenticated") else url_for("login"))

    if endpoint == "login":
        if session.get("authenticated"):
            return redirect(url_for("index"))
        return

    if not session.get("authenticated"):
        next_path = request.full_path.rstrip("?")
        if not next_path:
            next_path = url_for("index")
        return redirect(url_for("login", next=next_path))


@app.teardown_appcontext
def close_db(_error: BaseException | None) -> None:
    db = g.pop("db", None)
    if db is not None:
        db.close()


@app.route("/setup", methods=["GET", "POST"])
def setup():
    if has_password():
        return redirect(url_for("index") if session.get("authenticated") else url_for("login"))

    error = None
    if request.method == "POST":
        password = request.form.get("password", "")
        confirm_password = request.form.get("confirm_password", "")

        if len(password) < 8:
            error = "Password must be at least 8 characters."
        elif password != confirm_password:
            error = "Passwords do not match."
        else:
            set_setting(AUTH_PASSWORD_KEY, generate_password_hash(password))
            session.permanent = True
            session["authenticated"] = True
            return redirect(url_for("index"))

    return render_template("setup.html", error=error)


@app.route("/login", methods=["GET", "POST"])
def login():
    password_hash = get_setting(AUTH_PASSWORD_KEY)
    if not password_hash:
        return redirect(url_for("setup"))

    if session.get("authenticated"):
        return redirect(url_for("index"))

    next_path = request.args.get("next", "")
    if request.method == "POST":
        next_path = request.form.get("next_path", next_path)

    if not is_safe_path(next_path):
        next_path = url_for("index")

    error = None
    if request.method == "POST":
        password = request.form.get("password", "")
        if check_password_hash(password_hash, password):
            session.permanent = True
            session["authenticated"] = True
            return redirect(next_path)
        error = "Invalid password."

    return render_template("login.html", error=error, next_path=next_path)


@app.post("/logout")
def logout():
    session.clear()
    return redirect(url_for("login"))


@app.get("/")
def index():
    active_tab = request.args.get("tab", "register").strip().lower()
    if active_tab not in VALID_TABS:
        active_tab = "register"

    db = get_db()

    rows = db.execute(
        """
        SELECT id, tx_date, description, category, amount, tx_type
        FROM transactions
        ORDER BY tx_date DESC, id DESC
        """
    ).fetchall()

    running_balance = 0.0
    entries = []
    for row in rows:
        signed_amount = row["amount"] if row["tx_type"] == "credit" else -row["amount"]
        running_balance += signed_amount
        entries.append(
            {
                "id": row["id"],
                "tx_date": row["tx_date"],
                "description": row["description"],
                "category": row["category"],
                "tx_type": row["tx_type"],
                "amount": row["amount"],
                "signed_amount": signed_amount,
            }
        )

    edit_transaction = None
    edit_tx_id = parse_optional_int(request.args.get("edit_tx_id"))
    if edit_tx_id is not None:
        edit_row = db.execute(
            """
            SELECT id, tx_date, description, category, amount, tx_type
            FROM transactions
            WHERE id = ?
            """,
            (edit_tx_id,),
        ).fetchone()
        if edit_row is not None:
            edit_transaction = {
                "id": edit_row["id"],
                "tx_date": edit_row["tx_date"],
                "description": edit_row["description"],
                "category": edit_row["category"],
                "amount": edit_row["amount"],
                "tx_type": edit_row["tx_type"],
            }

    recurring_rows = db.execute(
        """
        SELECT id, next_due_date, description, category, amount, tx_type, frequency
        FROM recurring_transactions
        ORDER BY next_due_date ASC, id ASC
        """
    ).fetchall()
    recurring_entries = []
    for row in recurring_rows:
        recurring_entries.append(
            {
                "id": row["id"],
                "next_due_date": row["next_due_date"],
                "description": row["description"],
                "category": row["category"],
                "amount": row["amount"],
                "tx_type": row["tx_type"],
                "frequency": row["frequency"],
            }
        )
    recurring_projection_entries, _, _, _ = project_recurring_entries(recurring_rows)
    recurring_projection_by_id: dict[int, list[dict]] = {}
    for projection in recurring_projection_entries:
        recurring_projection_by_id.setdefault(projection["id"], []).append(projection)
    recurring_timeline_entries = []
    for recurring in recurring_entries:
        recurring_timeline_entries.append(
            {
                "id": recurring["id"],
                "due_date": recurring["next_due_date"],
                "description": recurring["description"],
                "frequency": recurring["frequency"],
                "tx_type": recurring["tx_type"],
                "amount": recurring["amount"],
                "is_projected": False,
            }
        )
        for projection in recurring_projection_by_id.get(recurring["id"], []):
            if projection["due_date"] == recurring["next_due_date"]:
                continue
            recurring_timeline_entries.append(
                {
                    "id": projection["id"],
                    "due_date": projection["due_date"],
                    "description": projection["description"],
                    "frequency": projection["frequency"],
                    "tx_type": projection["tx_type"],
                    "amount": projection["amount"],
                    "is_projected": True,
                }
            )
    recurring_timeline_entries.sort(key=lambda entry: (entry["due_date"], entry["is_projected"], entry["id"]))

    edit_recurring = None
    edit_recurring_id = parse_optional_int(request.args.get("edit_recurring_id"))
    if edit_recurring_id is not None:
        edit_row = db.execute(
            """
            SELECT id, next_due_date, description, category, amount, tx_type, frequency
            FROM recurring_transactions
            WHERE id = ?
            """,
            (edit_recurring_id,),
        ).fetchone()
        if edit_row is not None:
            edit_recurring = {
                "id": edit_row["id"],
                "next_due_date": edit_row["next_due_date"],
                "description": edit_row["description"],
                "category": edit_row["category"],
                "amount": edit_row["amount"],
                "tx_type": edit_row["tx_type"],
                "frequency": edit_row["frequency"],
            }

    return render_template(
        "index.html",
        entries=entries,
        recurring_entries=recurring_entries,
        balance=running_balance,
        active_tab=active_tab,
        edit_transaction=edit_transaction,
        edit_recurring=edit_recurring,
        today=today_iso(),
        recurring_timeline_entries=recurring_timeline_entries,
        recurring_projection_months=RECURRING_PROJECTION_MONTHS,
    )


@app.post("/transactions")
def add_transaction():
    raw_tx_date = request.form.get("tx_date", "").strip()
    tx_date = today_iso() if not raw_tx_date else parse_date(raw_tx_date)
    description = request.form.get("description", "").strip()
    category = "General"
    amount_raw = request.form.get("amount", "").strip()
    tx_type = request.form.get("tx_type", "debit").strip().lower()

    if tx_date is None:
        return redirect(url_for("index", tab="register"))

    try:
        amount = abs(float(amount_raw))
    except ValueError:
        return redirect(url_for("index", tab="register"))

    if not description or amount == 0 or tx_type not in VALID_TX_TYPES:
        return redirect(url_for("index", tab="register"))

    get_db().execute(
        """
        INSERT INTO transactions (tx_date, description, category, amount, tx_type, created_at)
        VALUES (?, ?, ?, ?, ?, ?)
        """,
        (
            tx_date,
            description,
            category,
            amount,
            tx_type,
            datetime.utcnow().isoformat(timespec="seconds"),
        ),
    )
    get_db().commit()

    return redirect(url_for("index", tab="register"))


@app.post("/transactions/<int:tx_id>/edit")
def edit_transaction(tx_id: int):
    tx_date = parse_date(request.form.get("tx_date", ""))
    description = request.form.get("description", "").strip()
    amount_raw = request.form.get("amount", "").strip()
    tx_type = request.form.get("tx_type", "debit").strip().lower()

    if tx_date is None:
        return redirect(url_for("index", tab="register", edit_tx_id=tx_id))

    try:
        amount = abs(float(amount_raw))
    except ValueError:
        return redirect(url_for("index", tab="register", edit_tx_id=tx_id))

    if not description or amount == 0 or tx_type not in VALID_TX_TYPES:
        return redirect(url_for("index", tab="register", edit_tx_id=tx_id))

    get_db().execute(
        """
        UPDATE transactions
        SET tx_date = ?, description = ?, amount = ?, tx_type = ?
        WHERE id = ?
        """,
        (tx_date, description, amount, tx_type, tx_id),
    )
    get_db().commit()
    return redirect(url_for("index", tab="register"))


@app.route("/transactions/<int:tx_id>/delete", methods=["GET", "POST"])
def delete_transaction(tx_id: int):
    get_db().execute("DELETE FROM transactions WHERE id = ?", (tx_id,))
    get_db().commit()
    return redirect(url_for("index", tab="register"))


@app.post("/recurring")
def add_recurring_transaction():
    next_due_date = parse_date(request.form.get("next_due_date", ""))
    description = request.form.get("description", "").strip()
    category = "General"
    amount_raw = request.form.get("amount", "").strip()
    tx_type = request.form.get("tx_type", "debit").strip().lower()
    frequency = request.form.get("frequency", "monthly").strip().lower()

    if next_due_date is None:
        return redirect(url_for("index", tab="recurring"))

    try:
        amount = abs(float(amount_raw))
    except ValueError:
        return redirect(url_for("index", tab="recurring"))

    if (
        not description
        or amount == 0
        or tx_type not in VALID_TX_TYPES
        or frequency not in VALID_FREQUENCIES
    ):
        return redirect(url_for("index", tab="recurring"))

    get_db().execute(
        """
        INSERT INTO recurring_transactions
        (next_due_date, description, category, amount, tx_type, frequency, is_completed, created_at)
        VALUES (?, ?, ?, ?, ?, ?, 0, ?)
        """,
        (
            next_due_date,
            description,
            category,
            amount,
            tx_type,
            frequency,
            datetime.utcnow().isoformat(timespec="seconds"),
        ),
    )
    get_db().commit()
    return redirect(url_for("index", tab="recurring"))


@app.post("/recurring/<int:recurring_id>/edit")
def edit_recurring_transaction(recurring_id: int):
    next_due_date = parse_date(request.form.get("next_due_date", ""))
    description = request.form.get("description", "").strip()
    amount_raw = request.form.get("amount", "").strip()
    tx_type = request.form.get("tx_type", "debit").strip().lower()
    frequency = request.form.get("frequency", "monthly").strip().lower()

    if next_due_date is None:
        return redirect(url_for("index", tab="recurring", edit_recurring_id=recurring_id))

    try:
        amount = abs(float(amount_raw))
    except ValueError:
        return redirect(url_for("index", tab="recurring", edit_recurring_id=recurring_id))

    if (
        not description
        or amount == 0
        or tx_type not in VALID_TX_TYPES
        or frequency not in VALID_FREQUENCIES
    ):
        return redirect(url_for("index", tab="recurring", edit_recurring_id=recurring_id))

    get_db().execute(
        """
        UPDATE recurring_transactions
        SET next_due_date = ?, description = ?, amount = ?, tx_type = ?, frequency = ?
        WHERE id = ?
        """,
        (next_due_date, description, amount, tx_type, frequency, recurring_id),
    )
    get_db().commit()
    return redirect(url_for("index", tab="recurring"))


@app.route("/recurring/<int:recurring_id>/delete", methods=["GET", "POST"])
def delete_recurring_transaction(recurring_id: int):
    get_db().execute("DELETE FROM recurring_transactions WHERE id = ?", (recurring_id,))
    get_db().commit()
    return redirect(url_for("index", tab="recurring"))


@app.route("/recurring/<int:recurring_id>/add", methods=["GET", "POST"])
def add_recurring_to_register(recurring_id: int):
    recurring = (
        get_db()
        .execute(
            """
            SELECT next_due_date, description, category, amount, tx_type, frequency
            FROM recurring_transactions
            WHERE id = ?
            """,
            (recurring_id,),
        )
        .fetchone()
    )
    if recurring is None:
        return redirect(url_for("index", tab="recurring"))

    get_db().execute(
        """
        INSERT INTO transactions (tx_date, description, category, amount, tx_type, created_at)
        VALUES (?, ?, ?, ?, ?, ?)
        """,
        (
            recurring["next_due_date"],
            recurring["description"],
            recurring["category"],
            recurring["amount"],
            recurring["tx_type"],
            datetime.utcnow().isoformat(timespec="seconds"),
        ),
    )
    get_db().execute(
        """
        UPDATE recurring_transactions
        SET next_due_date = ?
        WHERE id = ?
        """,
        (
            next_due_date_after_today(recurring["next_due_date"], recurring["frequency"]),
            recurring_id,
        ),
    )
    get_db().commit()
    return redirect(url_for("index", tab="recurring"))


if __name__ == "__main__":
    app.run(host="0.0.0.0", port=7070)