Project_Journal-Csharp_backend-Python-Frontend/journal/core/storage.py

import sys
import hashlib
import threading
import time
from cryptography.exceptions import InvalidTag
import shutil
import zipfile
from datetime import datetime
from pathlib import Path

sys.path.append(str(Path(__file__).resolve().parent.parent.parent))
from .parser import parse_journal_content, parse_journal_file
from .database import get_db_connection, hydrate_database
from .encryption import encrypt_data, decrypt_data

from .config import (
    DATA_DIR,
    VAULT_DIR,
    MONTHLY_VAULT_FORMAT,
)

_month_fingerprint_cache: dict[str, str] = {}
_vault_io_lock = threading.RLock()


# --- Monthly Vault Management ---


def _get_monthly_vault_path(date: datetime) -> Path:
    """Returns the path for the monthly vault file."""
    return VAULT_DIR / date.strftime(MONTHLY_VAULT_FORMAT)


def _create_monthly_archive(month_path: Path, archive_path: Path):
    """Creates a zip archive of a temporary monthly directory."""
    with zipfile.ZipFile(archive_path, "w", zipfile.ZIP_DEFLATED) as zipf:
        for file_path in month_path.iterdir():
            zipf.write(file_path, arcname=file_path.name)  # Store only filename in zip


def _extract_monthly_archive(archive_path: Path, extract_to_path: Path):
    """Extracts a zip archive to a specified directory."""
    with zipfile.ZipFile(archive_path, "r") as zipf:
        zipf.extractall(extract_to_path)


# --- Public API for Journal Storage ---


def _save_month(password: str, month_key: str, files_in_month: list[Path]):
    """Helper function to save a single month's vault."""
    # We need a datetime object to generate the vault path, strptime is perfect.
    month_as_date = datetime.strptime(month_key, "%Y-%m")
    monthly_vault_path = _get_monthly_vault_path(month_as_date)

    # Create a temporary directory to stage files for zipping
    temp_month_dir = VAULT_DIR / f"temp_{month_key}"
    temp_month_dir.mkdir(exist_ok=True)
    temp_zip_path: Path | None = None

    try:
        for file_path in files_in_month:
            _ = shutil.copy(file_path, temp_month_dir)

        # Create a temporary zip archive
        temp_zip_path = VAULT_DIR / f"temp_{month_key}.zip"
        _create_monthly_archive(temp_month_dir, temp_zip_path)

        with open(temp_zip_path, "rb") as f_in:
            zip_content = f_in.read()

        encrypted_vault_content = encrypt_data(zip_content, password)
        with open(monthly_vault_path, "wb") as f_out:
            _ = f_out.write(encrypted_vault_content)
        _month_fingerprint_cache[month_key] = _compute_month_fingerprint(files_in_month)
        print(f"Successfully saved {monthly_vault_path.name}")
    except Exception as e:
        print(f"Error saving month {month_key}: {e}")
    finally:
        shutil.rmtree(temp_month_dir, ignore_errors=True)
        if temp_zip_path and temp_zip_path.exists():
            temp_zip_path.unlink()


def _compute_month_fingerprint(files: list[Path]) -> str:
    fingerprint = hashlib.sha256()
    for file_path in sorted(files, key=lambda p: p.name):
        try:
            stat = file_path.stat()
        except OSError:
            continue
        fingerprint.update(file_path.name.encode("utf-8"))
        fingerprint.update(str(stat.st_mtime_ns).encode("ascii"))
        fingerprint.update(str(stat.st_size).encode("ascii"))
    return fingerprint.hexdigest()


def get_today_filename() -> Path:
    """Returns the path for today's journal entry in the active DATA_DIR."""
    return DATA_DIR / f"{datetime.now().strftime('%Y-%m-%d')}.md"


def save_entry_content(
    content: str, file_path: Path | None = None, mode: str = "Daily"
):
    target_file = file_path or get_today_filename()
    target_file.parent.mkdir(parents=True, exist_ok=True)

    if mode == "Fragment":
        print(f"Appending fragment to {target_file.name}...")
        with open(target_file, "a", encoding="utf-8") as f:
            # Ensure there's a newline before the new content
            _ = f.write("\n\n" + content.strip())
        return

    # For Daily, Deep, etc., perform a merge
    if target_file.exists():
        print(f"Merging content into existing file: {target_file.name}")
        existing_entry = parse_journal_file(str(target_file))
        new_entry_data = parse_journal_content(content, target_file.stem)
        existing_entry.merge_with(new_entry_data)
        final_content = existing_entry.to_markdown()
    else:
        print(f"Creating new entry: {target_file.name}")
        final_content = content

    _ = target_file.write_text(final_content, encoding="utf-8")


def load_all_vaults(password: str) -> bool:
    """
    Decrypts and extracts all monthly vaults into the DATA_DIR.
    Cleans DATA_DIR before extraction.
    Returns True on success, False if password is incorrect for existing vaults.
    """
    if not password:
        raise ValueError("Password cannot be empty.")

    with _vault_io_lock:
        _month_fingerprint_cache.clear()

        # Clear DATA_DIR first
        _clear_data_dir_with_retries()
        DATA_DIR.mkdir(parents=True, exist_ok=True)

        if not VAULT_DIR.exists() or not any(VAULT_DIR.iterdir()):
            print("Vault directory is empty or does not exist. Assuming new vault.")
            return True  # No vaults to load, so it's a success (new vault)

        decryption_successful = False
        for vault_file in VAULT_DIR.glob("*.vault"):
            if vault_file.name == "_init_vault.vault":
                print(f"Deleting old dummy vault file: {vault_file.name}")
                vault_file.unlink()
                continue
            try:
                with open(vault_file, "rb") as f_in:
                    encrypted_data = f_in.read()

                decrypted_zip_content = decrypt_data(encrypted_data, password)

                # Write decrypted content to a temporary zip file
                temp_zip_path = VAULT_DIR / f"temp_{vault_file.name}.zip"
                with open(temp_zip_path, "wb") as f_out:
                    _ = f_out.write(decrypted_zip_content)

                _extract_monthly_archive(temp_zip_path, DATA_DIR)
                temp_zip_path.unlink()  # Clean up temp zip
                decryption_successful = True
                print(f"Successfully loaded {vault_file.name}")
                print(
                    f"Contents of DATA_DIR after loading {vault_file.name}: {list(DATA_DIR.iterdir())}"
                )
            except InvalidTag:
                print(
                    f"Warning: Could not decrypt '{vault_file.name}'. Invalid password for this file."
                )
                # Do not set decryption_successful to True if only some files fail
            except Exception as e:
                print(f"Error loading vault '{vault_file.name}': {e}")
                # If any other error occurs, it's not necessarily a password issue

        if not decryption_successful and any(VAULT_DIR.iterdir()):
            # If there are vault files, but none could be decrypted, password is wrong
            print("Error: No vault files could be decrypted with the provided password.")
            return False

        # --- Database Hydration ---
        # After successfully decrypting files, hydrate the live, encrypted database.
        conn = None
        try:
            all_entries = [parse_journal_file(str(f)) for f in DATA_DIR.glob("*.md")]
            if all_entries:
                conn = get_db_connection(password)
                hydrate_database(conn, all_entries)
        except Exception as e:
            print(f"Fatal error during database hydration: {e}")
            return False  # Treat DB hydration failure as a critical error
        finally:
            if conn is not None:
                conn.close()

        return True


def rebuild_all_vaults(password: str):
    """
    Rebuilds all monthly vaults from the files in the DATA_DIR.

    This is a comprehensive but slower operation, intended for use on shutdown
    or via the CLI to ensure all changes, including to older entries, are
    persisted. It iterates through all decrypted files and saves them to their
    respective monthly vaults.
    """
    print("rebuild_all_vaults called.")
    if not password:
        raise ValueError("Password cannot be empty.")

    with _vault_io_lock:
        # Group files by month
        monthly_files: dict[str, list[Path]] = {}
        for file_path in DATA_DIR.glob("*.md"):
            try:
                file_date = datetime.strptime(file_path.stem, "%Y-%m-%d")
                month_key = file_date.strftime("%Y-%m")
                if month_key not in monthly_files:
                    monthly_files[month_key] = []
                monthly_files[month_key].append(file_path)
            except ValueError:  # Skip files that don't match YYYY-MM-DD format
                print(f"Skipping non-journal file in DATA_DIR: {file_path.name}")
                continue

        # Ensure VAULT_DIR exists
        VAULT_DIR.mkdir(parents=True, exist_ok=True)

        for month_key, files_in_month in monthly_files.items():
            _save_month(password, month_key, files_in_month)


def save_current_month_vault(password: str):
    """
    Optimized save function that only rebuilds the current month's vault.

    This is used for frequent, in-session saves from the UI to provide better
    performance, as it only operates on the files for the current month.
    """
    print("save_current_month_vault called.")
    if not password:
        raise ValueError("Password cannot be empty.")

    with _vault_io_lock:
        # Determine current month
        now = datetime.now()
        month_key = now.strftime("%Y-%m")

        # Collect files for the current month
        files_in_month: list[Path] = []
        for file_path in DATA_DIR.glob("*.md"):
            if file_path.stem.startswith(month_key):
                files_in_month.append(file_path)

        if not files_in_month:
            print(f"No files found for the current month ({month_key}) to save.")
            return

        current_fingerprint = _compute_month_fingerprint(files_in_month)
        cached_fingerprint = _month_fingerprint_cache.get(month_key)
        if cached_fingerprint == current_fingerprint:
            print(f"Skipping vault save for {month_key}; no file changes detected.")
            return

        _save_month(password, month_key, files_in_month)


def initialize_vault(password: str):
    """
    Ensures the VAULT_DIR exists. The first save operation will create the initial vault files.
    """
    if not password:
        raise ValueError("Password cannot be empty.")

    VAULT_DIR.mkdir(parents=True, exist_ok=True)
    print("Vault directory ensured to exist.")


def clear_data_directory():
    """
    Clears the DATA_DIR. This should only be called on application shutdown.
    """
    print("Clearing DATA_DIR...")
    with _vault_io_lock:
        # The encrypted database file lives in DATA_DIR, so this function
        # will securely delete it along with all the decrypted .md files.
        _clear_data_dir_with_retries()
        DATA_DIR.mkdir(parents=True, exist_ok=True)
        _month_fingerprint_cache.clear()
        print("DATA_DIR cleared.")


def _clear_data_dir_with_retries(retries: int = 5, delay_seconds: float = 0.2) -> None:
    if not DATA_DIR.exists():
        return
    for attempt in range(retries):
        try:
            shutil.rmtree(DATA_DIR)
            return
        except PermissionError:
            if attempt == retries - 1:
                raise
            time.sleep(delay_seconds)