from flask import Blueprint, jsonify, request, session, current_app, Response
from extensions import db
from itsdangerous import URLSafeSerializer
from models.notepad import Notepad
from models.note import Note
from models.collaborator import Collaborator
from models.user import User
from models.favorite import Favorite
from models.notepad_key import NotepadKey
from sqlalchemy import func
import csv
from io import StringIO

api_bp = Blueprint('api', __name__, url_prefix='/api')

def get_share_token(item_type, item_id):
    s = URLSafeSerializer(current_app.config['SECRET_KEY'], salt='share-link')
    return s.dumps({'type': item_type, 'id': item_id})

def validate_share_token(token, item_type, item_id):
    s = URLSafeSerializer(current_app.config['SECRET_KEY'], salt='share-link')
    try:
        data = s.loads(token)
        return data.get('type') == item_type and data.get('id') == item_id
    except:
        return False

def check_pad_read_permission(notepad, user_id, token=None):
    is_shared_by_link = token and validate_share_token(token, 'notepad', notepad.id)
    if notepad.is_public or is_shared_by_link:
        return True
    if not user_id:
        return False
    if notepad.user_id == user_id:
        return True
    # Check for collaborator (any permission level grants read)
    if Collaborator.query.filter_by(user_id=user_id, notepad_id=notepad.id).first():
        return True
    return False

def check_pad_write_permission(notepad, user_id):
    if not user_id:
        return False
    user = User.query.get(user_id)
    if user and user.is_admin:
        return True
    if notepad.is_public:
        return True
    if notepad.user_id == user_id:
        return True
    # Check for collaborator with 'write' permission
    collab = Collaborator.query.filter_by(user_id=user_id, notepad_id=notepad.id, permission='write').first()
    return collab is not None

def check_note_write_permission(note, notepad, user_id):
    if not user_id:
        return False
    user = User.query.get(user_id)
    if user and user.is_admin:
        return True
    if note.user_id and note.user_id == user_id:
        return True
    if notepad.user_id and notepad.user_id == user_id:
        return True
    collab = Collaborator.query.filter_by(user_id=user_id, notepad_id=notepad.id, permission='write').first()
    return collab is not None

@api_bp.route('/users/lookup', methods=['GET'])
def lookup_user():
    user_id = session.get('user_id')
    if not user_id:
        return jsonify({'error': 'Unauthorized'}), 401
    email = request.args.get('email')
    if not email:
        return jsonify({'error': 'Email is required'}), 400
    user = User.query.filter_by(email=email).first()
    if not user:
        return jsonify({'error': 'User not found'}), 404
    return jsonify({
        'id': user.id,
        'email': user.email,
        'first_name': user.first_name,
        'last_name': user.last_name,
        'public_key': user.public_key
    })

@api_bp.route('/notepads', methods=['GET'])
def get_notepads():
    user_id = session.get('user_id')
    filter_type = request.args.get('type')
    
    query = Notepad.query
    
    if filter_type == 'favorites':
        query = Notepad.query.join(Favorite).filter(Favorite.user_id == user_id) if user_id else Notepad.query.filter_by(id=-1)
    elif filter_type == 'mine':
        query = query.filter_by(user_id=user_id) if user_id else Notepad.query.filter_by(id=-1)
    elif filter_type == 'shared':
        query = Notepad.query.join(Collaborator).filter(Collaborator.user_id == user_id) if user_id else Notepad.query.filter_by(id=-1)
    elif filter_type == 'public':
        query = query.filter_by(is_public=True)
    else:
        # Default behavior (union) if no type specified
        if user_id:
            query = query.filter((Notepad.is_public == True) | (Notepad.user_id == user_id))
        else:
            query = query.filter_by(is_public=True)
        
    pads = query.order_by(Notepad.name).all()
    
    # Optimize note count query (O(1) database hit instead of O(N))
    note_counts = dict(db.session.query(Note.notepad_id, func.count(Note.id)).filter(Note.is_hidden == False).group_by(Note.notepad_id).all())
    
    # Optimize write permissions check and keys
    collab_write_pad_ids = set()
    user_keys = {}
    if user_id:
        collab_write_pad_ids = {c.notepad_id for c in Collaborator.query.filter_by(user_id=user_id, permission='write').all()}
        user_keys = {k.notepad_id: k.encrypted_key for k in NotepadKey.query.filter_by(user_id=user_id).all()}

    results = []
    for pad in pads:
        note_count = note_counts.get(pad.id, 0)
        p_dict = pad.to_dict(note_count=note_count, encrypted_key=user_keys.get(pad.id))
        if user_id:
            if pad.user_id == user_id:
                p_dict['share_token'] = get_share_token('notepad', pad.id)
            p_dict['can_write'] = pad.is_public or (pad.user_id == user_id) or (pad.id in collab_write_pad_ids)
        else:
            p_dict['can_write'] = False
        results.append(p_dict)
    return jsonify(results)

@api_bp.route('/notepads/<int:pad_id>', methods=['GET'])
def get_notepad(pad_id):
    notepad = Notepad.query.get_or_404(pad_id)
    user_id = session.get('user_id')
    
    token = request.args.get('token')
    
    if not check_pad_read_permission(notepad, user_id, token):
        return jsonify({'error': 'Unauthorized'}), 403
        
    enc_key = None
    if user_id and not notepad.is_public:
        nk = NotepadKey.query.filter_by(notepad_id=pad_id, user_id=user_id).first()
        if nk:
            enc_key = nk.encrypted_key

    p_dict = notepad.to_dict(encrypted_key=enc_key)
    p_dict['can_write'] = check_pad_write_permission(notepad, user_id)
    if user_id and notepad.user_id == user_id:
        p_dict['share_token'] = get_share_token('notepad', pad_id)
    return jsonify(p_dict)

@api_bp.route('/notepads', methods=['POST'])
def create_notepad():
    user_id = session.get('user_id')
    if not user_id:
        return jsonify({'error': 'Unauthorized'}), 401
        
    data = request.get_json()
    is_public = data.get('is_public', False)
    new_pad = Notepad(
        name=data['name'],
        is_public=is_public,
        user_id=user_id,
        color=data.get('color', '#3388ff')
    )
    db.session.add(new_pad)
    db.session.flush()

    encrypted_key = data.get('encrypted_key')
    if not is_public and encrypted_key:
        notepad_key = NotepadKey(
            notepad_id=new_pad.id,
            user_id=user_id,
            encrypted_key=encrypted_key
        )
        db.session.add(notepad_key)

    db.session.commit()
    return jsonify(new_pad.to_dict(encrypted_key=encrypted_key)), 201

@api_bp.route('/notepads/<int:pad_id>/notes', methods=['GET'])
def get_notes(pad_id):
    notepad = Notepad.query.get_or_404(pad_id)
    user_id = session.get('user_id')
    
    token = request.args.get('token')
    
    if not check_pad_read_permission(notepad, user_id, token):
        return jsonify({'error': 'Unauthorized'}), 403

    notes = notepad.notes.filter_by(is_hidden=False).all()
    results = []
    for note in notes:
        n_dict = note.to_dict()
        if user_id and notepad.user_id == user_id:
            n_dict['share_token'] = get_share_token('note', note.id)
        results.append(n_dict)
    return jsonify(results)

@api_bp.route('/notes/<int:note_id>', methods=['GET'])
def get_note(note_id):
    note = Note.query.get_or_404(note_id)
    notepad = Notepad.query.get(note.notepad_id)
    
    user_id = session.get('user_id')
    
    token = request.args.get('token')
    is_shared = token and validate_share_token(token, 'note', note_id)

    if not notepad.is_public and (not user_id or notepad.user_id != user_id) and not is_shared:
        return jsonify({'error': 'Unauthorized'}), 403
        
    n_dict = note.to_dict()
    n_dict['can_write'] = check_note_write_permission(note, notepad, user_id)
    
    pad_dict = notepad.to_dict()
    # Redact name if private and not owner
    if not notepad.is_public and (not user_id or notepad.user_id != user_id):
        pad_dict['name'] = "Private Notepad"
        
    n_dict['notepad'] = pad_dict
    n_dict['owner_id'] = notepad.user_id
    n_dict['is_public'] = notepad.is_public
    
    if user_id and notepad.user_id == user_id:
        n_dict['share_token'] = get_share_token('note', note_id)
        
    return jsonify(n_dict)
@api_bp.route('/notes/mine', methods=['GET'])
def get_my_all_notes():
    user_id = session.get('user_id')
    if not user_id:
        return jsonify({'error': 'Unauthorized'}), 401
        
    notes = Note.query.join(Notepad).filter(Notepad.user_id == user_id, Note.is_hidden == False).all()
    
    # Eagerly load note counts and user keys to avoid N+1 queries
    note_counts = dict(db.session.query(Note.notepad_id, func.count(Note.id)).filter(Note.is_hidden == False).group_by(Note.notepad_id).all())
    user_keys = {k.notepad_id: k.encrypted_key for k in NotepadKey.query.filter_by(user_id=user_id).all()}
    
    results = []
    for note in notes:
        n_dict = note.to_dict()
        n_dict['notepad'] = note.notepad.to_dict(note_count=note_counts.get(note.notepad_id, 0), encrypted_key=user_keys.get(note.notepad_id))
        n_dict['owner_id'] = user_id
        n_dict['is_public'] = note.notepad.is_public
        n_dict['share_token'] = get_share_token('note', note.id)
        results.append(n_dict)
    return jsonify(results)

@api_bp.route('/notes/default', methods=['GET'])
def get_default_notes():
    user_id = session.get('user_id')
    user_keys = {}
    if user_id:
        # Get owned notepads + favorite notepads + shared notepads
        owned_pad_ids = [p.id for p in Notepad.query.filter_by(user_id=user_id).all()]
        fav_pad_ids = [f.notepad_id for f in Favorite.query.filter_by(user_id=user_id).all()]
        all_visible_pad_ids = list(set(owned_pad_ids + fav_pad_ids))
        
        notes = Note.query.filter(Note.notepad_id.in_(all_visible_pad_ids), Note.is_hidden == False).all()
        user_keys = {k.notepad_id: k.encrypted_key for k in NotepadKey.query.filter_by(user_id=user_id).all()}
    else:
        # Logged out: All public notes
        notes = Note.query.join(Notepad).filter(Notepad.is_public == True, Note.is_hidden == False).all()
        
    note_counts = dict(db.session.query(Note.notepad_id, func.count(Note.id)).filter(Note.is_hidden == False).group_by(Note.notepad_id).all())
    
    results = []
    for note in notes:
        n_dict = note.to_dict()
        n_dict['notepad'] = note.notepad.to_dict(note_count=note_counts.get(note.notepad_id, 0), encrypted_key=user_keys.get(note.notepad_id))
        n_dict['owner_id'] = note.notepad.user_id
        n_dict['is_public'] = note.notepad.is_public
        if user_id and note.notepad.user_id == user_id:
            n_dict['share_token'] = get_share_token('note', note.id)
        results.append(n_dict)
    return jsonify(results)

@api_bp.route('/notes', methods=['POST'])
def create_note():
    user_id = session.get('user_id')
    if not user_id:
        return jsonify({'error': 'Unauthorized'}), 401

    data = request.get_json()
    notepad_id = data.get('notepad_id')
    if not notepad_id:
        return jsonify({'error': 'notepad_id is required'}), 400

    notepad = Notepad.query.get_or_404(notepad_id)

    if not check_pad_write_permission(notepad, user_id):
        return jsonify({'error': 'Forbidden'}), 403

    is_encrypted = not notepad.is_public
    if is_encrypted:
        new_note = Note(
            title=data.get('title', '🔒 Encrypted Note'),
            content=data['content'],
            lat=0.0,
            lng=0.0,
            is_encrypted=True,
            notepad_id=notepad_id,
            user_id=user_id
        )
    else:
        new_note = Note(
            title=data.get('title', 'Untitled'),
            content=data['content'],
            lat=float(data.get('lat', 0.0)),
            lng=float(data.get('lng', 0.0)),
            is_encrypted=False,
            notepad_id=notepad_id,
            user_id=user_id
        )
    db.session.add(new_note)
    db.session.commit()
    return jsonify(new_note.to_dict()), 201

@api_bp.route('/notes/<int:note_id>', methods=['PUT'])
def update_note(note_id):
    user_id = session.get('user_id')
    if not user_id:
        return jsonify({'error': 'Unauthorized'}), 401
    note = Note.query.get_or_404(note_id)
    notepad = Notepad.query.get(note.notepad_id)
    if not check_note_write_permission(note, notepad, user_id):
        return jsonify({'error': 'Forbidden'}), 403
    data = request.get_json()
    note.title = data.get('title', note.title)
    note.content = data.get('content', note.content)
    if not note.is_encrypted:
        if 'lat' in data:
            note.lat = float(data['lat'])
        if 'lng' in data:
            note.lng = float(data['lng'])
    db.session.commit()
    return jsonify(note.to_dict())

@api_bp.route('/notepads/<int:pad_id>', methods=['PUT'])
def update_notepad(pad_id):
    user_id = session.get('user_id')
    if not user_id:
        return jsonify({'error': 'Unauthorized'}), 401
    
    notepad = Notepad.query.get_or_404(pad_id)
    # Only owner can edit notepad metadata
    if notepad.user_id != user_id:
        return jsonify({'error': 'Forbidden'}), 403
        
    data = request.get_json()
    notepad.name = data.get('name', notepad.name)
    notepad.is_public = data.get('is_public', notepad.is_public)
    notepad.color = data.get('color', notepad.color)
    db.session.commit()
    return jsonify(notepad.to_dict())

@api_bp.route('/notes/<int:note_id>/report', methods=['POST'])
def report_note(note_id):
    note = Note.query.get_or_404(note_id)
    if not note.notepad.is_public:
        return jsonify({'error': 'Can only report public notes'}), 400
    note.is_hidden = True
    db.session.commit()
    return jsonify({'message': 'Note reported'}), 200

@api_bp.route('/notepads/<int:pad_id>/favorite', methods=['POST'])
def favorite_notepad(pad_id):
    user_id = session.get('user_id')
    if not user_id:
        return jsonify({'error': 'Unauthorized'}), 401

    existing_favorite = Favorite.query.filter_by(user_id=user_id, notepad_id=pad_id).first()
    if existing_favorite:
        return jsonify({'message': 'Notepad already favorited'}), 200

    new_favorite = Favorite(user_id=user_id, notepad_id=pad_id)
    db.session.add(new_favorite)
    db.session.commit()
    return jsonify({'message': 'Notepad favorited'}), 201

@api_bp.route('/notepads/<int:pad_id>/unfavorite', methods=['POST'])
def unfavorite_notepad(pad_id):
    user_id = session.get('user_id')
    favorite = Favorite.query.filter_by(user_id=user_id, notepad_id=pad_id).first()
    if favorite:
        db.session.delete(favorite)
        db.session.commit()
    return jsonify({'message': 'Notepad unfavorited'}), 200

@api_bp.route('/notepads/<int:pad_id>/collaborators', methods=['GET'])
def get_collaborators(pad_id):
    user_id = session.get('user_id')
    notepad = Notepad.query.get_or_404(pad_id)

    # Only owner can see collaborators
    if notepad.user_id != user_id:
        return jsonify({'error': 'Forbidden'}), 403

    collaborators = Collaborator.query.join(User).filter(Collaborator.notepad_id == pad_id).all()
    
    results = []
    for collab in collaborators:
        results.append({
            'user_id': collab.user.id,
            'email': collab.user.email,
            'permission': collab.permission
        })
    return jsonify(results)

@api_bp.route('/notepads/<int:pad_id>/collaborators', methods=['POST'])
def add_collaborator(pad_id):
    user_id = session.get('user_id')
    notepad = Notepad.query.get_or_404(pad_id)

    if notepad.user_id != user_id:
        return jsonify({'error': 'Forbidden'}), 403

    data = request.get_json()
    email = data.get('email')
    permission = data.get('permission')

    if not email or not permission or permission not in ['read', 'write']:
        return jsonify({'error': 'Email and permission (read/write) are required'}), 400

    collaborator_user = User.query.filter_by(email=email).first()
    if not collaborator_user:
        return jsonify({'error': 'User not found'}), 404
        
    if collaborator_user.id == user_id:
        return jsonify({'error': 'Cannot add yourself as a collaborator'}), 400

    existing_collab = Collaborator.query.filter_by(notepad_id=pad_id, user_id=collaborator_user.id).first()
    if existing_collab:
        existing_collab.permission = permission
    else:
        new_collab = Collaborator(notepad_id=pad_id, user_id=collaborator_user.id, permission=permission)
        db.session.add(new_collab)

    encrypted_key = data.get('encrypted_key')
    if not notepad.is_public and encrypted_key:
        existing_key = NotepadKey.query.filter_by(notepad_id=pad_id, user_id=collaborator_user.id).first()
        if existing_key:
            existing_key.encrypted_key = encrypted_key
        else:
            new_key = NotepadKey(notepad_id=pad_id, user_id=collaborator_user.id, encrypted_key=encrypted_key)
            db.session.add(new_key)
    
    db.session.commit()
    return jsonify({'message': 'Collaborator added/updated successfully'}), 200

@api_bp.route('/notepads/<int:pad_id>/collaborators/<int:collaborator_user_id>', methods=['DELETE'])
def remove_collaborator(pad_id, collaborator_user_id):
    user_id = session.get('user_id')
    notepad = Notepad.query.get_or_404(pad_id)

    if notepad.user_id != user_id:
        return jsonify({'error': 'Forbidden'}), 403

    collab_to_remove = Collaborator.query.filter_by(notepad_id=pad_id, user_id=collaborator_user_id).first()
    if not collab_to_remove:
        return jsonify({'error': 'Collaborator not found'}), 404

    NotepadKey.query.filter_by(notepad_id=pad_id, user_id=collaborator_user_id).delete()
    db.session.delete(collab_to_remove)
    db.session.commit()
    return jsonify({'message': 'Collaborator removed'}), 200

@api_bp.route('/notepads/<int:pad_id>/export', methods=['GET'])
def export_notepad(pad_id):
    notepad = Notepad.query.get_or_404(pad_id)
    user_id = session.get('user_id')
    
    token = request.args.get('token')
    if not check_pad_read_permission(notepad, user_id, token):
        return jsonify({'error': 'Unauthorized'}), 403

    si = StringIO()
    cw = csv.writer(si)
    cw.writerow(['Note ID', 'Title', 'Content', 'Latitude', 'Longitude', 'Created At'])
    
    for note in notepad.notes.filter_by(is_hidden=False):
        cw.writerow([note.id, note.title, note.content, note.lat, note.lng, note.created_at])
        
    output = si.getvalue()
    return Response(
        output,
        mimetype="text/csv",
        headers={"Content-Disposition": f"attachment;filename=notepad_{pad_id}_export.csv"}
    )

@api_bp.route('/notepads/<int:pad_id>', methods=['DELETE'])
def delete_notepad(pad_id):
    user_id = session.get('user_id')
    if not user_id:
        return jsonify({'error': 'Unauthorized'}), 401
    
    notepad = Notepad.query.get_or_404(pad_id)
    if notepad.user_id != user_id:
        return jsonify({'error': 'Forbidden'}), 403
        
    for note in notepad.notes:
        db.session.delete(note)
    Favorite.query.filter_by(notepad_id=notepad.id).delete()
    NotepadKey.query.filter_by(notepad_id=notepad.id).delete()
    Collaborator.query.filter_by(notepad_id=notepad.id).delete()
    db.session.delete(notepad)
    db.session.commit()
    return jsonify({'message': 'Notepad deleted'}), 200

@api_bp.route('/notes/<int:note_id>', methods=['DELETE'])
def delete_note(note_id):
    user_id = session.get('user_id')
    if not user_id:
        return jsonify({'error': 'Unauthorized'}), 401
        
    note = Note.query.get_or_404(note_id)
    notepad = Notepad.query.get(note.notepad_id)
    
    if not check_note_write_permission(note, notepad, user_id):
        return jsonify({'error': 'Forbidden'}), 403
        
    db.session.delete(note)
    db.session.commit()
    return jsonify({'message': 'Note deleted'}), 200