from flask import Blueprint, render_template, request, redirect, url_for, flash, session, jsonify
from utils.logging.logger_setup import logger
from utils.db.db import mongo, get_next_id
import bcrypt
from functools import wraps
from bson import json_util
import json

assistant_bp = Blueprint('assistant', __name__, template_folder='../../templates', url_prefix='/assistant')

# --- AUTH DECORATORS ---
def assistant_login_required(f):
    """Ensures the user is logged in as an assistant."""
    @wraps(f)
    def decorated_function(*args, **kwargs):
        if session.get('user_type') != 'assistant':
            flash('You must be logged in as an assistant to view this page.', 'warning')
            return redirect(url_for('auth.assistant_login'))
        return f(*args, **kwargs)
    return decorated_function

def assistant_role_required(role_name):
    """Ensures the assistant has the required role."""
    def decorator(f):
        @wraps(f)
        def decorated_function(*args, **kwargs):
            if session.get('role') != role_name:
                flash(f'You require the {role_name} role for this page.', 'warning')
                return redirect(url_for('assistant.dashboard'))
            return f(*args, **kwargs)
        return decorated_function
    return decorator

# --- HELPER FUNCTIONS ---
def find_patient(patient_id):
    """Finds a patient by their integer ID or legacy sql_id."""
    try:
        pid = int(patient_id)
        return mongo.db.patient.find_one({"$or": [{"id": pid}, {"sql_id": pid}]})
    except (ValueError, TypeError):
        return None

# --- CORE ROUTES ---
@assistant_bp.route('/dashboard')
@assistant_login_required
def dashboard():
    """Redirects assistant to their specific dashboard based on role."""
    try:
        role = session.get('role')
        if role == 'account_creation':
            return redirect(url_for('assistant.account_creation_dashboard'))
        elif role == 'data_entry':
            return redirect(url_for('assistant.data_entry_dashboard'))
        else:
            flash('Your assistant role is not defined. Please contact an administrator.', 'error')
            return redirect(url_for('admin.logout'))
    except Exception as e:
        logger.error(f"Error in assistant dashboard redirect: {e}")
        return redirect(url_for('admin.logout'))

@assistant_bp.route('/dashboard/account_creation', methods=['GET', 'POST'])
@assistant_login_required
@assistant_role_required('account_creation')
def account_creation_dashboard():
    """
    Handles new patient account creation form.
    
    Verifies no duplicate email exists before creating a new patient record.
    """
    if request.method == 'POST':
        try:
            name = request.form['name']
            email = request.form['email']
            password = request.form['password']
            phone = request.form.get('phone')
            
            if mongo.db.patient.find_one({"email": email}):
                flash(f'A patient with the email {email} already exists.', 'error')
            else:
                hashed_password = bcrypt.hashpw(password.encode('utf-8'), bcrypt.gensalt()).decode('utf-8')
                new_id = get_next_id('patient')
                new_patient = {
                    "id": new_id, "sql_id": new_id, "name": name, "email": email, 
                    "password": hashed_password, "phone": phone
                }
                mongo.db.patient.insert_one(new_patient)
                flash(f'Patient "{name}" (ID: {new_id}) created successfully!', 'success')
                logger.info(f"Patient account created by {session.get('user_name')}: {email}")
                return redirect(url_for('assistant.account_creation_dashboard'))
        except Exception as e:
            logger.error(f"Error creating patient account: {e}")
            flash('An error occurred while creating the account.', 'error')
            
    return render_template('assistant/assistant_account_creation.html')

@assistant_bp.route('/dashboard/data_entry', methods=['GET'])
@assistant_login_required
@assistant_role_required('data_entry')
def data_entry_dashboard():
    """
    The main data entry interface.
    
    Allows searching for patients and loading their related medical data.
    """
    try:
        query = request.args.get('q', '').strip()
        selected_patient_id = request.args.get('patient_id')
        
        patients = []
        # Search logic
        if query:
            criteria = [{"name": {"$regex": query, "$options": "i"}},
                        {"email": {"$regex": query, "$options": "i"}},
                        {"phone": {"$regex": query, "$options": "i"}}]
            if query.isdigit():
                criteria.extend([{"id": int(query)}, {"sql_id": int(query)}])
            patients = list(mongo.db.patient.find({"$or": criteria}))
        else:
            # Initial load
            patients = list(mongo.db.patient.find().sort("name", 1).limit(50))

        # Fetch details
        selected_patient = None
        related_data = {"medical_history": [], "consultations": [], "test_reports": [], 
                        "treatments": [], "billings": [], "insurance": None}
        
        if selected_patient_id:
            selected_patient = find_patient(selected_patient_id)
        elif len(patients) == 1:
             # Auto-select single result
            selected_patient = patients[0]
            selected_patient_id = selected_patient.get('id')

        if selected_patient:
            resolved_id = selected_patient.get('id')
            
            # Update Session for RAG/Chat Context
            session['active_patient_id'] = resolved_id
            
            # Fetch related data
            related_data["medical_history"] = list(mongo.db.medicalhistory.find({"patient_id": resolved_id}))
            related_data["consultations"] = list(mongo.db.consultation.find({"patient_id": resolved_id}))
            related_data["test_reports"] = list(mongo.db.testreport.find({"patient_id": resolved_id}))
            related_data["treatments"] = list(mongo.db.treatment.find({"patient_id": resolved_id}))
            related_data["billings"] = list(mongo.db.billing.find({"patient_id": resolved_id}))
            related_data["insurance"] = mongo.db.insurance.find_one({"patient_id": resolved_id})

        return render_template(
            'assistant/assistant_data_entry.html', 
            patients=patients, 
            search_query=query,
            selected_patient=selected_patient,
            **related_data
        )
    except Exception as e:
        logger.error(f"Error loading data entry dashboard: {e}")
        flash("An error occurred loading the dashboard.", "error")
        return render_template('assistant/assistant_data_entry.html', patients=[], search_query=query, selected_patient=None)

@assistant_bp.route('/save_complete_record/<patient_id>', methods=['POST'])
@assistant_login_required
@assistant_role_required('data_entry')
def save_complete_record(patient_id):
    """
    Handles the batch 'Save Complete Record' POST request.
    
    Updates profile, insurance, and inserts new medical records in a single transaction-like flow.
    """
    try:
        patient = find_patient(patient_id)
        if not patient:
            return jsonify({"success": False, "message": "Patient not found"}), 404
            
        data = request.get_json()
        if not data:
            return jsonify({"success": False, "message": "No data provided"}), 400

        resolved_id = patient.get('id')
        actions_taken = []

        # --- Update Single-Object Fields ---
        if 'profile' in data and data['profile']:
            mongo.db.patient.update_one({"id": resolved_id}, {"$set": data['profile']})
            actions_taken.append("Profile Updated")

        if 'insurance' in data and data['insurance']:
            mongo.db.insurance.update_one(
                {"patient_id": resolved_id},
                {"$set": data['insurance']},
                upsert=True
            )
            actions_taken.append("Insurance Updated")

        # --- Insert Items from Lists ---
        list_map = {
            'summary': (mongo.db.medicalhistory, 'medicalhistory'),
            'consultation': (mongo.db.consultation, 'consultation'),
            'test': (mongo.db.testreport, 'testreport'),
            'treatment': (mongo.db.treatment, 'treatment'),
            'billing': (mongo.db.billing, 'billing')
        }

        for key, (collection, id_name) in list_map.items():
            if key in data and data[key]:
                items_to_insert = []
                for item in data[key]:
                    new_id = get_next_id(id_name)
                    item.update({'id': new_id, 'sql_id': new_id, 'patient_id': resolved_id})
                    items_to_insert.append(item)
                
                if items_to_insert:
                    collection.insert_many(items_to_insert)
                    actions_taken.append(f"{len(items_to_insert)} {key.capitalize()} record(s) added")

        if not actions_taken:
            return jsonify({"success": False, "message": "No new data was staged to save."}), 400

        msg = f"Success! {', '.join(actions_taken)}."
        flash(msg, 'success')
        logger.info(f"Complete record saved for patient {resolved_id} by {session.get('user_name')}")
        
        return jsonify({
            "success": True, 
            "message": msg,
            "redirect": url_for('assistant.data_entry_dashboard', patient_id=resolved_id)
        })

    except Exception as e:
        logger.error(f"Error in save_complete_record for patient {patient_id}: {e}")
        return jsonify({"success": False, "message": f"Server error: {str(e)}"}), 500


@assistant_bp.route('/get_patient_data/<patient_id>', methods=['GET'])
@assistant_login_required
def get_patient_data(patient_id):
    """Fetches all data for a specific patient as JSON."""
    try:
        selected_patient = find_patient(patient_id)
        if not selected_patient:
            return jsonify({"error": "Patient not found"}), 404

        resolved_id = selected_patient.get('id')
        
        # Update Session for RAG/Chat Context
        session['active_patient_id'] = resolved_id
        
        related_data = {
            "user": selected_patient, # Profile data
            "medical_history": list(mongo.db.medicalhistory.find({"patient_id": resolved_id})),
            "consultations": list(mongo.db.consultation.find({"patient_id": resolved_id})),
            "test_reports": list(mongo.db.testreport.find({"patient_id": resolved_id})),
            "treatments": list(mongo.db.treatment.find({"patient_id": resolved_id})),
            "billings": list(mongo.db.billing.find({"patient_id": resolved_id})),
            "insurance": mongo.db.insurance.find_one({"patient_id": resolved_id})
        }

        # Use bson.json_util to handle ObjectId and datetime serialization
        return json.loads(json_util.dumps(related_data))
    except Exception as e:
        logger.error(f"Error in get_patient_data: {e}")
        return jsonify({"error": "Failed to fetch patient data"}), 500