from flask import Flask, render_template, request, redirect, url_for, session, flash, jsonify
from flask_bcrypt import Bcrypt
from pymongo import MongoClient
from bson.objectid import ObjectId
from datetime import datetime
import os
from dotenv import load_dotenv

load_dotenv()

app = Flask(__name__)
app.config['SECRET_KEY'] = os.getenv('SECRET_KEY')

client = MongoClient(os.getenv('MONGO_URI'))
db = client.get_default_database()
bcrypt = Bcrypt(app)

# --- Helpers ---
def get_user_id():
    return session.get('user_id')

def is_logged_in():
    return 'user_id' in session

# --- Routes ---

@app.route('/')
def index():
    if not is_logged_in():
        return redirect(url_for('login'))
    
    user_id = get_user_id()
    
    # Get all activities
    activities = list(db.activities.find({'user_id': user_id}))
    
    # Check for active time entry
    current_entry = db.time_entries.find_one({
        'user_id': user_id,
        'end_time': None
    })
    
    active_tasks = []
    if current_entry:
        # Get tasks associated with the active activity type
        active_activity = db.activities.find_one({'_id': current_entry['activity_id']})
        if active_activity:
            current_entry['activity_name'] = active_activity['name']
            current_entry['activity_color'] = active_activity.get('color', '#3498db')
        
        # 1. Tasks generated specifically for this session (time_entry_id linked)
        # 2. Open tasks linked to this activity type (contextual todos)
        active_tasks = list(db.tasks.find({
            'user_id': user_id,
            'status': 'open',
            '$or': [
                {'time_entry_id': current_entry['_id']},
                {'activity_id': current_entry['activity_id'], 'is_template': False, 'time_entry_id': None}
            ]
        }))

    return render_template('dashboard.html', 
                           activities=activities, 
                           current_entry=current_entry,
                           tasks=active_tasks)

@app.route('/register', methods=['GET', 'POST'])
def register():
    if request.method == 'POST':
        username = request.form['username']
        password = request.form['password']
        
        if db.users.find_one({'username': username}):
            flash('Username already exists')
            return redirect(url_for('register'))
            
        hashed_password = bcrypt.generate_password_hash(password).decode('utf-8')
        user_id = db.users.insert_one({'username': username, 'password': hashed_password}).inserted_id
        
        # Auto login
        session['user_id'] = str(user_id)
        flash('Account created! You are now logged in.')
        return redirect(url_for('index'))
        
    return render_template('register.html')

@app.route('/login', methods=['GET', 'POST'])
def login():
    if request.method == 'POST':
        user = db.users.find_one({'username': request.form['username']})
        if user and bcrypt.check_password_hash(user['password'], request.form['password']):
            session['user_id'] = str(user['_id'])
            return redirect(url_for('index'))
        flash('Login Unsuccessful. Please check username and password')
    return render_template('login.html')

@app.route('/logout')
def logout():
    session.clear()
    return redirect(url_for('login'))

@app.route('/add_activity', methods=['POST'])
def add_activity():
    if not is_logged_in(): return redirect(url_for('login'))
    name = request.form['name']
    color = request.form.get('color', '#3498db')
    
    activity_id = db.activities.insert_one({
        'user_id': get_user_id(),
        'name': name,
        'color': color
    }).inserted_id
    
    # Add optional tasks
    tasks_text = request.form.get('tasks', '')
    if tasks_text:
        tasks = [t.strip() for t in tasks_text.split(',') if t.strip()]
        for t in tasks:
            db.task_templates.insert_one({
                'activity_id': activity_id,
                'user_id': get_user_id(),
                'name': t
            })
            
    return redirect(url_for('index'))

@app.route('/toggle_timer/<activity_id>', methods=['POST'])
def toggle_timer(activity_id):
    if not is_logged_in(): return redirect(url_for('login'))
    user_id = get_user_id()
    note = request.form.get('note', '')

    # Check if this activity is currently running
    current = db.time_entries.find_one({
        'user_id': user_id, 
        'activity_id': ObjectId(activity_id),
        'end_time': None
    })

    if current:
        # Stop it
        db.time_entries.update_one(
            {'_id': current['_id']},
            {'$set': {'end_time': datetime.now()}}
        )
    else:
        # Stop ANY other running timer first
        db.time_entries.update_many(
            {'user_id': user_id, 'end_time': None},
            {'$set': {'end_time': datetime.now()}}
        )
        # Start new
        new_entry_id = db.time_entries.insert_one({
            'user_id': user_id,
            'activity_id': ObjectId(activity_id),
            'start_time': datetime.now(),
            'end_time': None,
            'note': note
        }).inserted_id

        # Check for Task Templates and auto-add them
        templates = db.tasks.find({
            'user_id': user_id,
            'activity_id': ObjectId(activity_id),
            'is_template': True
        })
        
        for t in templates:
            db.tasks.insert_one({
                'user_id': user_id,
                'name': t['name'],
                'activity_id': ObjectId(activity_id),
                'time_entry_id': new_entry_id, # Link to this specific session
                'status': 'open',
                'is_template': False,
                'created_at': datetime.now(),
                'comments': []
            })
        
    return redirect(url_for('index'))

@app.route('/stop_timer', methods=['POST'])
def stop_timer():
    if not is_logged_in(): return redirect(url_for('login'))
    db.time_entries.update_many(
        {'user_id': get_user_id(), 'end_time': None},
        {'$set': {'end_time': datetime.now()}}
    )
    return redirect(url_for('index'))

@app.route('/complete_task', methods=['POST'])
def complete_task():
    if not is_logged_in(): return jsonify({'error': 'auth'}), 401
    
    task_id = request.form['task_id']
    is_checked = request.form['is_checked'] == 'true'
    
    status = 'completed' if is_checked else 'open'
    completed_at = datetime.now() if is_checked else None

    db.tasks.update_one(
        {'_id': ObjectId(task_id), 'user_id': get_user_id()},
        {'$set': {'status': status, 'completed_at': completed_at}}
    )

    # If it was a generic activity task, bind it to current entry if one acts so it shows in log
    # (Optional logic, skipping for simplicity)

    return jsonify({'status': 'success'})

# --- New Task Management Routes ---

@app.route('/tasks')
def tasks():
    if not is_logged_in(): return redirect(url_for('login'))
    user_id = get_user_id()
    
    # Fetch all activities for the dropdown
    activities = list(db.activities.find({'user_id': user_id}))
    
    # Categorize tasks
    # 1. Standalone / Manual Tasks
    tasks_list = list(db.tasks.find({
        'user_id': user_id, 
        'is_template': False, 
        'time_entry_id': None, # Don't show session-specific generated tasks here to avoid clutter? or show all? 
        # Let's show manual tasks only. Session tasks are transient.
        'status': 'open'
    }).sort('due_date', 1))

    # 2. Templates
    templates = list(db.tasks.find({'user_id': user_id, 'is_template': True}))
    
    # Enhance tasks with activity names
    act_map = {str(a['_id']): a for a in activities}
    for t in tasks_list + templates:
        if t.get('activity_id'):
            aid = str(t['activity_id'])
            if aid in act_map:
                t['activity_name'] = act_map[aid]['name']
                t['activity_color'] = act_map[aid].get('color', '#ccc')

    return render_template('tasks.html', tasks=tasks_list, templates=templates, activities=activities)

@app.route('/create_task', methods=['POST'])
def create_task():
    if not is_logged_in(): return redirect(url_for('login'))
    
    name = request.form['name']
    activity_id = request.form.get('activity_id')
    due_date_str = request.form.get('due_date')
    is_template = 'is_template' in request.form
    
    task_doc = {
        'user_id': get_user_id(),
        'name': name,
        'status': 'open',
        'is_template': is_template,
        'created_at': datetime.now(),
        'comments': []
    }
    
    if activity_id:
        task_doc['activity_id'] = ObjectId(activity_id)
        
    if due_date_str:
        task_doc['due_date'] = datetime.strptime(due_date_str, '%Y-%m-%dT%H:%M')
        
    # Standard task or template doesn't belong to a specific time entry yet
    task_doc['time_entry_id'] = None
    
    db.tasks.insert_one(task_doc)
    return redirect(url_for('tasks'))

@app.route('/task/<task_id>', methods=['GET', 'POST'])
def task_detail(task_id):
    if not is_logged_in(): return redirect(url_for('login'))
    
    task = db.tasks.find_one({'_id': ObjectId(task_id), 'user_id': get_user_id()})
    if not task: return "Task not found", 404
    
    if request.method == 'POST':
        comment = request.form['comment']
        if comment:
            db.tasks.update_one(
                {'_id': ObjectId(task_id)},
                {'$push': {'comments': {
                    'text': comment, 
                    'created_at': datetime.now(),
                    'user_id': get_user_id() # In case we add teams later
                }}}
            )
            return redirect(url_for('task_detail', task_id=task_id))

    return render_template('task_detail.html', task=task)

@app.route('/logbook')
def logbook():
    if not is_logged_in(): return redirect(url_for('login'))
    
    # Agrregation to join activities and tasks
    pipeline = [
        {'$match': {'user_id': get_user_id(), 'end_time': {'$ne': None}}},
        {'$sort': {'start_time': -1}},
        {'$lookup': {
            'from': 'activities',
            'localField': 'activity_id',
            'foreignField': '_id',
            'as': 'activity'
        }},
        {'$unwind': '$activity'},
        # Join with tasks collection instead of completed_tasks
        {'$lookup': {
            'from': 'tasks', 
            'localField': '_id',
            'foreignField': 'time_entry_id',
            'as': 'tasks'
        }},
    ]
    
    log = list(db.time_entries.aggregate(pipeline))
    
    # Basic formatting
    for entry in log:
        # Filter tasks to only completed ones for display cleanly in list view
        entry['tasks'] = [t for t in entry['tasks'] if t['status'] == 'completed']
        duration = entry['end_time'] - entry['start_time']
        entry['duration_str'] = str(duration).split('.')[0] # HH:MM:SS
        
    return render_template('logbook.html', log=log)

@app.route('/logbook/<entry_id>', methods=['GET', 'POST'])
def log_entry_detail(entry_id):
    if not is_logged_in(): return redirect(url_for('login'))
    
    # Handle Note Update
    if request.method == 'POST':
        new_note = request.form.get('note')
        db.time_entries.update_one(
            {'_id': ObjectId(entry_id), 'user_id': get_user_id()},
            {'$set': {'note': new_note}}
        )
        flash('Session note updated')
        return redirect(url_for('log_entry_detail', entry_id=entry_id))

    # Fetch Entry with Activity info
    pipeline = [
        {'$match': {'_id': ObjectId(entry_id), 'user_id': get_user_id()}},
        {'$lookup': {
            'from': 'activities',
            'localField': 'activity_id',
            'foreignField': '_id',
            'as': 'activity'
        }},
        {'$unwind': '$activity'}
    ]
    
    results = list(db.time_entries.aggregate(pipeline))
    if not results:
        return "Entry not found", 404
    
    entry = results[0]
    
    # Calculate duration
    if entry.get('end_time') and entry.get('start_time'):
        duration = entry['end_time'] - entry['start_time']
        entry['duration_str'] = str(duration).split('.')[0]
    else:
        entry['duration_str'] = "Running..."

    # Fetch ALL Tasks linked to this entry (completed or not)
    tasks = list(db.tasks.find({'time_entry_id': ObjectId(entry_id)}))
    
    return render_template('logbook_detail.html', entry=entry, tasks=tasks)

if __name__ == '__main__':
    app.run(debug=True)