app.py

import os
import re
import time
import logging
import threading
from flask import Flask, render_template, request, jsonify, redirect, url_for
from flask_sqlalchemy import SQLAlchemy
from werkzeug.utils import secure_filename
from sqlalchemy import create_engine, text, inspect
from huggingface_hub import HfApi, hf_hub_download

# Set up logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

app = Flask(__name__)

# --- App Settings ---
UPLOAD_FOLDER = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'uploads')
app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER
app.config['MAX_CONTENT_LENGTH'] = 32 * 1024 * 1024  # 32MB upload limit
os.makedirs(UPLOAD_FOLDER, exist_ok=True)

# --- Hugging Face Dataset Persistence Config ---
HF_TOKEN = os.environ.get("HF_TOKEN")
DATASET_REPO_ID = os.environ.get("DATASET_REPO_ID")
if DATASET_REPO_ID:
    DATASET_REPO_ID = DATASET_REPO_ID.replace("https://huggingface.co/datasets/", "")
    DATASET_REPO_ID = DATASET_REPO_ID.replace("https://huggingface.co/", "")
    DATASET_REPO_ID = DATASET_REPO_ID.strip("/")
DB_FILENAME = "app_records.db"
DB_FILE_PATH = os.path.join(os.path.dirname(os.path.abspath(__file__)), DB_FILENAME)

# 1. Download database from Hugging Face Dataset on startup
if HF_TOKEN and DATASET_REPO_ID:
    logger.info("Attempting to download database from Hugging Face Dataset repo...")
    try:
        hf_hub_download(
            repo_id=DATASET_REPO_ID,
            filename=DB_FILENAME,
            repo_type="dataset",
            token=HF_TOKEN,
            local_dir=os.path.dirname(os.path.abspath(__file__))
        )
        logger.info("Database downloaded successfully.")
    except Exception as e:
        logger.warning(f"Could not download database (might be first run or empty dataset): {e}")
else:
    logger.warning("HF_TOKEN or DATASET_REPO_ID missing. Starting with a fresh database locally.")

# 2. Configure Flask-SQLAlchemy to use the SQLite file
app.config['SQLALCHEMY_DATABASE_URI'] = f"sqlite:///{DB_FILE_PATH}"
app.config['SQLALCHEMY_TRACK_MODIFICATIONS'] = False

db = SQLAlchemy(app)

# --- Database Schema Models ---
class CaseRecord(db.Model):
    __tablename__ = 'case_records'
    id = db.Column(db.Integer, primary_key=True)
    case_name = db.Column(db.String(100), nullable=False)
    description = db.Column(db.Text, nullable=True)
    created_at = db.Column(db.Float, default=time.time)
    
    # Cascade delete files when a case is deleted
    files = db.relationship('EvidenceFile', backref='case', cascade='all, delete-orphan', lazy=True)

    def __init__(self, case_name: str, description: str = None):
        self.case_name = case_name
        self.description = description
        self.created_at = time.time()

class EvidenceFile(db.Model):
    __tablename__ = 'evidence_files'
    id = db.Column(db.Integer, primary_key=True)
    case_id = db.Column(db.Integer, db.ForeignKey('case_records.id', ondelete='CASCADE'), nullable=True)
    filename = db.Column(db.String(255), nullable=False)
    size_bytes = db.Column(db.Integer, nullable=False)
    created_at = db.Column(db.Float, default=time.time)

    def __init__(self, filename: str, size_bytes: int, case_id: int = None):
        self.filename = filename
        self.size_bytes = size_bytes
        self.case_id = case_id
        self.created_at = time.time()

# 3. Synchronize database back to Hugging Face Dataset in the background
def sync_db_to_hf():
    if HF_TOKEN and DATASET_REPO_ID:
        try:
            logger.info("Syncing database to Hugging Face Dataset...")
            api = HfApi()
            api.upload_file(
                path_or_fileobj=DB_FILE_PATH,
                path_in_repo=DB_FILENAME,
                repo_id=DATASET_REPO_ID,
                repo_type="dataset",
                token=HF_TOKEN
            )
            logger.info("Database synced to Hugging Face successfully.")
        except Exception as e:
            logger.error(f"Error syncing database to Hugging Face Dataset: {e}")

# Initialize database schema tables inside application context
try:
    with app.app_context():
        inspector = inspect(db.engine)
        if not inspector.has_table("case_records") or not inspector.has_table("evidence_files"):
            db.create_all()
            logger.info("Database schemas initialized successfully.")
        else:
            logger.info("Database schemas already exist. Skipping creation.")
except Exception as e:
    logger.warning(f"Database table verification/creation skipped: {e}")

# --- Routes ---
@app.route('/')
def home():
    try:
        cases = CaseRecord.query.order_by(CaseRecord.created_at.desc()).all()
        files = EvidenceFile.query.order_by(EvidenceFile.created_at.desc()).all()
        db_status = "HF Dataset Connected" if HF_TOKEN and DATASET_REPO_ID else "Local SQLite Mode"
    except Exception as e:
        logger.error(f"Error reading from database: {e}")
        cases, files = [], []
        db_status = "Disconnected (Offline Mode)"

    return render_template('index.html', cases=cases, files=files, db_status=db_status)

@app.route('/submit', methods=['POST'])
def submit_form():
    case_name = request.form.get('case_name', '').strip()
    description = request.form.get('description', '').strip()

    # Validations
    if not case_name or len(case_name) < 3 or len(case_name) > 100:
        return jsonify({"success": False, "error": "Case name must be between 3 and 100 characters."}), 400

    try:
        new_record = CaseRecord(case_name=case_name, description=description)
        db.session.add(new_record)
        db.session.commit()
        
        # Trigger background sync to Hugging Face Dataset
        threading.Thread(target=sync_db_to_hf, daemon=True).start()
        
        return jsonify({"success": True, "message": "Case successfully logged!"})
    except Exception as e:
        logger.error(f"Failed to submit database entry: {e}")
        db.session.rollback()
        return jsonify({"success": False, "error": "Database error occurred."}), 500

@app.route('/delete-case/<int:case_id>', methods=['POST'])
def delete_case(case_id):
    try:
        case = CaseRecord.query.get(case_id)
        if not case:
            return jsonify({"success": False, "error": "Case not found."}), 404
        
        # Delete files associated from disk
        for file_record in case.files:
            file_path = os.path.join(app.config['UPLOAD_FOLDER'], file_record.filename)
            if os.path.exists(file_path):
                os.remove(file_path)

        db.session.delete(case)
        db.session.commit()

        # Trigger background sync
        threading.Thread(target=sync_db_to_hf, daemon=True).start()
        return jsonify({"success": True, "message": "Case and associated files deleted."})
    except Exception as e:
        logger.error(f"Error deleting case: {e}")
        db.session.rollback()
        return jsonify({"success": False, "error": "Failed to delete case."}), 500

@app.route('/upload', methods=['POST'])
def upload_file():
    if 'file' not in request.files:
        return jsonify({"success": False, "error": "No file part in the request."}), 400
    
    file = request.files['file']
    case_id_val = request.form.get('case_id', '').strip()
    
    if file.filename == '':
        return jsonify({"success": False, "error": "No file selected."}), 400

    case_id = int(case_id_val) if case_id_val and case_id_val.isdigit() else None

    if file:
        filename = secure_filename(file.filename)
        save_path = os.path.join(app.config['UPLOAD_FOLDER'], filename)
        try:
            file.save(save_path)
            size_bytes = os.path.getsize(save_path)

            # Create Database Record
            new_file = EvidenceFile(filename=filename, size_bytes=size_bytes, case_id=case_id)
            db.session.add(new_file)
            db.session.commit()

            # Trigger background sync
            threading.Thread(target=sync_db_to_hf, daemon=True).start()

            return jsonify({
                "success": True, 
                "filename": filename, 
                "size_bytes": size_bytes,
                "message": "File uploaded successfully."
            })
        except Exception as e:
            logger.error(f"File upload error: {e}")
            db.session.rollback()
            return jsonify({"success": False, "error": "Failed to save file on server."}), 500

@app.route('/delete-file/<int:file_id>', methods=['POST'])
def delete_file(file_id):
    try:
        file_record = EvidenceFile.query.get(file_id)
        if not file_record:
            return jsonify({"success": False, "error": "File record not found."}), 404

        file_path = os.path.join(app.config['UPLOAD_FOLDER'], file_record.filename)
        if os.path.exists(file_path):
            os.remove(file_path)

        db.session.delete(file_record)
        db.session.commit()

        # Trigger background sync
        threading.Thread(target=sync_db_to_hf, daemon=True).start()
        return jsonify({"success": True, "message": "Evidence file deleted."})
    except Exception as e:
        logger.error(f"Error deleting file: {e}")
        db.session.rollback()
        return jsonify({"success": False, "error": "Failed to delete file."}), 500

@app.route('/health')
def health():
    try:
        db.session.execute(text('SELECT 1'))
        db_healthy = True
    except Exception as e:
        logger.error(f"Health check failed for database: {e}")
        db_healthy = False

    return jsonify({
        "status": "healthy",
        "database": "connected" if db_healthy else "disconnected",
        "storage": "writable" if os.access(UPLOAD_FOLDER, os.W_OK) else "readonly",
        "sync_enabled": bool(HF_TOKEN and DATASET_REPO_ID)
    })

if __name__ == '__main__':
    # Run local server
    app.run(host='0.0.0.0', port=int(os.environ.get('PORT', 7860)), debug=True)