"""Hospital visit API routes."""
from bson import ObjectId
from flask import Blueprint, jsonify, request
from pydantic import ValidationError

from datetime import UTC, datetime

from app.db import Database
from app.models.hospital_visit import HospitalVisit
from app.repositories import find_user_by_id, hospital_visit_repository
from app.routes.access import ensure_can_access_patient_data, ensure_can_create_for_patient
from app.routes.auth import auth_required, require_roles

hospital_visits_bp = Blueprint('hospital_visits_api', __name__)


def _populate_hospital_visits_batch(visit_dicts: list[dict]) -> list[dict]:
    """Populate multiple hospital visits in bulk (max 5 queries total)."""
    if not visit_dicts:
        return visit_dicts

    db = Database.get_db()

    # 1) Collect all visit IDs and explicit diagnosis IDs across all visits
    visit_ids = []
    all_explicit_diag_ids = []
    for v in visit_dicts:
        vid = v.get('_id')
        if vid:
            visit_ids.append(ObjectId(str(vid)))
        for did in v.get('diagnoses', []):
            try:
                all_explicit_diag_ids.append(ObjectId(str(did)))
            except Exception:
                pass

    # 2) Batch-fetch all diagnoses: by explicit ID or by hospitalVisit back-ref
    diag_query = {'$or': []}
    if all_explicit_diag_ids:
        diag_query['$or'].append({'_id': {'$in': all_explicit_diag_ids}})
    if visit_ids:
        diag_query['$or'].append({'hospitalVisit': {'$in': visit_ids}})

    all_diagnoses = list(db.get_collection('diagnoses').find(diag_query)) if diag_query['$or'] else []

    # 3) Batch-fetch all diseases for the ICD codes we found
    icd_codes = set()
    for d in all_diagnoses:
        code = d.get('icdCode', '')
        if code:
            icd_codes.add(code)
            icd_codes.add(code.split('.')[0])

    all_diseases = {
        doc['icdCode']: doc
        for doc in db.get_collection('diseases').find({'icdCode': {'$in': list(icd_codes)}})
    } if icd_codes else {}

    # 4) Batch-fetch prescriptions via diagnosis IDs (Prescription.diagnosisId → Diagnosis._id)
    #    diagnosisId may be stored as ObjectId or string — query both forms
    all_diag_oids = [d['_id'] for d in all_diagnoses]
    all_diag_strings = [str(d['_id']) for d in all_diagnoses]
    all_prescriptions = list(
        db.get_collection('prescriptions').find({
            '$or': [
                {'diagnosisId': {'$in': all_diag_oids}},
                {'diagnosisId': {'$in': all_diag_strings}},
            ]
        })
    ) if all_diagnoses else []

    # 5) Batch-fetch all medications referenced by those prescriptions
    med_ids = set()
    for p in all_prescriptions:
        mid = p.get('medicationId')
        if mid:
            med_ids.add(ObjectId(str(mid)))

    all_medications = {
        str(doc['_id']): doc
        for doc in db.get_collection('medications').find({'_id': {'$in': list(med_ids)}})
    } if med_ids else {}

    # --- Index diagnoses by ID and by visit ---
    diag_by_id = {str(d['_id']): d for d in all_diagnoses}
    diag_by_visit = {}
    for d in all_diagnoses:
        hv = d.get('hospitalVisit')
        if hv:
            diag_by_visit.setdefault(str(hv), []).append(d)

    # --- Index prescriptions by diagnosis ID (normalize to string) ---
    pres_by_diag = {}
    for p in all_prescriptions:
        did = p.get('diagnosisId')
        if did:
            pres_by_diag.setdefault(str(did), []).append(p)

    # --- Assemble populated data per visit ---
    for v in visit_dicts:
        vid = str(v.get('_id', ''))

        # Diagnoses: explicit refs + hospitalVisit back-refs, deduplicated
        seen = set()
        visit_diags = []
        for did in v.get('diagnoses', []):
            d = diag_by_id.get(str(did))
            if d and str(d['_id']) not in seen:
                seen.add(str(d['_id']))
                visit_diags.append(d)
        for d in diag_by_visit.get(vid, []):
            if str(d['_id']) not in seen:
                seen.add(str(d['_id']))
                visit_diags.append(d)

        populated_diagnoses = []
        for d in visit_diags:
            entry = {
                '_id': str(d['_id']),
                'icdCode': d.get('icdCode', ''),
                'notes': d.get('notes', ''),
                'practitioner': d.get('practitioner', ''),
                'diagnosisClinic': d.get('diagnosisClinic', ''),
                'severity': d.get('severity', ''),
            }
            icd_code = d.get('icdCode', '')
            disease = all_diseases.get(icd_code) or all_diseases.get(icd_code.split('.')[0])
            if disease:
                entry['diseaseName'] = disease.get('name', '')
                entry['bodyRegion'] = disease.get('bodyRegion', '')
            populated_diagnoses.append(entry)
        v['populatedDiagnoses'] = populated_diagnoses

        # Prescriptions — collected through diagnoses (Prescription.diagnosisId → Diagnosis._id)
        populated_prescriptions = []
        seen_pres = set()
        for d in visit_diags:
            for p in pres_by_diag.get(str(d['_id']), []):
                pid = str(p['_id'])
                if pid in seen_pres:
                    continue
                seen_pres.add(pid)
                pres = {
                    '_id': pid,
                    'dosage': p.get('dosage', ''),
                    'frequency': p.get('frequency', ''),
                }
                mid = p.get('medicationId')
                if mid:
                    med = all_medications.get(str(mid))
                    if med:
                        pres['medicationName'] = med.get('genericName', '')
                populated_prescriptions.append(pres)
        v['populatedPrescriptions'] = populated_prescriptions


    return visit_dicts


def _populate_hospital_visit(visit_dict: dict) -> dict:
    """Populate a single hospital visit. Delegates to batch for consistency."""
    return _populate_hospital_visits_batch([visit_dict])[0]


@hospital_visits_bp.route('/hospital-visits', methods=['POST'])
@auth_required
def create_hospital_visit_endpoint():
    """Create a hospital visit.
    ---
    tags:
        - Hospital Visits
    security:
        - BearerAuth: []
    requestBody:
        required: true
        content:
            application/json:
                schema:
                    $ref: '#/components/schemas/HospitalVisitCreateRequest'
    responses:
        201:
            description: Visit created
            content:
                application/json:
                    schema:
                        $ref: '#/components/schemas/HospitalVisitResponse'
        400:
            description: Validation or payload error
            content:
                application/json:
                    schema:
                        $ref: '#/components/schemas/ErrorResponse'
        401:
            description: Unauthorized
            content:
                application/json:
                    schema:
                        $ref: '#/components/schemas/ErrorResponse'
        403:
            description: Forbidden
            content:
                application/json:
                    schema:
                        $ref: '#/components/schemas/ErrorResponse'
    """
    payload = request.get_json(silent=True) or {}
    patient_id = str(payload.get('patientId', ''))
    create_err = ensure_can_create_for_patient(patient_id)
    if create_err:
        return create_err

    try:
        hospital_visit = HospitalVisit.model_validate(payload)
    except ValidationError as exc:
        return jsonify({'error': 'ValidationError', 'detail': str(exc)}), 400

    created = hospital_visit_repository.create(hospital_visit)
    return jsonify(created.model_dump(by_alias=True, mode='json')), 201


@hospital_visits_bp.route('/hospital-visits/<hospital_visit_id>', methods=['GET'])
@auth_required
def get_hospital_visit_by_id_endpoint(hospital_visit_id: str):
    """Get a hospital visit by ID.
    ---
    tags:
        - Hospital Visits
    security:
        - BearerAuth: []
    parameters:
        - in: path
          name: hospital_visit_id
          required: true
          schema:
            type: string
    responses:
        200:
            description: Visit found
            content:
                application/json:
                    schema:
                        $ref: '#/components/schemas/HospitalVisitResponse'
        401:
            description: Unauthorized
            content:
                application/json:
                    schema:
                        $ref: '#/components/schemas/ErrorResponse'
        403:
            description: Forbidden
            content:
                application/json:
                    schema:
                        $ref: '#/components/schemas/ErrorResponse'
        404:
            description: Visit not found
            content:
                application/json:
                    schema:
                        $ref: '#/components/schemas/ErrorResponse'
    """
    hospital_visit = hospital_visit_repository.find_by_id(hospital_visit_id)
    if not hospital_visit:
        return jsonify({'error': 'NotFound'}), 404

    _, err = ensure_can_access_patient_data(str(hospital_visit.patientId))
    if err:
        return err

    result = _populate_hospital_visit(hospital_visit.model_dump(by_alias=True, mode='json'))
    return jsonify(result)


@hospital_visits_bp.route('/hospital-visits', methods=['GET'])
@auth_required
def list_hospital_visits_by_patient_endpoint():
    """List hospital visits by patient ID query parameter.
    ---
    tags:
        - Hospital Visits
    security:
        - BearerAuth: []
    parameters:
        - in: query
          name: patientId
          required: true
          schema:
            type: string
    responses:
        200:
            description: Visit list
            content:
                application/json:
                    schema:
                        $ref: '#/components/schemas/HospitalVisitResponseList'
        400:
            description: Missing or invalid patientId
            content:
                application/json:
                    schema:
                        $ref: '#/components/schemas/ErrorResponse'
        401:
            description: Unauthorized
            content:
                application/json:
                    schema:
                        $ref: '#/components/schemas/ErrorResponse'
        403:
            description: Forbidden
            content:
                application/json:
                    schema:
                        $ref: '#/components/schemas/ErrorResponse'
    """
    patient_id = request.args.get('patientId')
    if not patient_id:
        return jsonify({'error': 'MissingQueryParam', 'param': 'patientId'}), 400

    _, err = ensure_can_access_patient_data(patient_id)
    if err:
        return err

    try:
        visits = hospital_visit_repository.find_by_patient_id(patient_id)
    except ValueError as exc:
        return jsonify({'error': 'InvalidPatientId', 'detail': str(exc)}), 400

    visit_dicts = [v.model_dump(by_alias=True, mode='json') for v in visits]
    return jsonify(_populate_hospital_visits_batch(visit_dicts))


@hospital_visits_bp.route('/hospital-visits/by-practitioner', methods=['GET'])
@auth_required
@require_roles('practitioner', 'admin')
def list_hospital_visits_by_practitioner_endpoint():
    """List hospital visits across all patients shared with the authenticated practitioner.
    ---
    tags:
        - Hospital Visits
    security:
        - BearerAuth: []
    parameters:
        - in: query
          name: startDate
          schema:
            type: string
            format: date
          description: Filter visits on or after this date (ISO format)
        - in: query
          name: endDate
          schema:
            type: string
            format: date
          description: Filter visits on or before this date (ISO format)
    responses:
        200:
            description: Visit list across practitioner's patients
            content:
                application/json:
                    schema:
                        $ref: '#/components/schemas/HospitalVisitResponseList'
        401:
            description: Unauthorized
        403:
            description: Forbidden
        404:
            description: Practitioner not found
    """
    from flask import g
    practitioner_id = str((g.get('auth_user') or {}).get('sub', ''))
    practitioner = find_user_by_id(practitioner_id)
    if not practitioner:
        return jsonify({'error': 'NotFound', 'detail': 'Practitioner not found'}), 404

    patient_ids = [str(pid) for pid in (practitioner.sharingInfo or [])]
    if not patient_ids:
        return jsonify([])

    start_date = None
    end_date = None
    raw_start = request.args.get('startDate')
    raw_end = request.args.get('endDate')
    if raw_start:
        start_date = datetime.fromisoformat(raw_start).replace(tzinfo=UTC)
    if raw_end:
        end_date = datetime.fromisoformat(raw_end).replace(hour=23, minute=59, second=59, tzinfo=UTC)

    visits = hospital_visit_repository.find_by_patient_ids(patient_ids, start_date, end_date, doctor_name=practitioner.dispName)
    visit_dicts = [v.model_dump(by_alias=True, mode='json') for v in visits]
    return jsonify(_populate_hospital_visits_batch(visit_dicts))


@hospital_visits_bp.route('/hospital-visits/<hospital_visit_id>', methods=['PATCH'])
@auth_required
def update_hospital_visit_endpoint(hospital_visit_id: str):
    """Update a hospital visit record.
    ---
    tags:
        - Hospital Visits
    security:
        - BearerAuth: []
    parameters:
        - in: path
          name: hospital_visit_id
          required: true
          schema:
            type: string
    requestBody:
        required: true
        content:
            application/json:
                schema:
                    $ref: '#/components/schemas/HospitalVisitUpdateRequest'
    responses:
        200:
            description: Visit updated
            content:
                application/json:
                    schema:
                        $ref: '#/components/schemas/HospitalVisitResponse'
        400:
            description: No updatable fields
            content:
                application/json:
                    schema:
                        $ref: '#/components/schemas/ErrorResponse'
        401:
            description: Unauthorized
            content:
                application/json:
                    schema:
                        $ref: '#/components/schemas/ErrorResponse'
        403:
            description: Forbidden
            content:
                application/json:
                    schema:
                        $ref: '#/components/schemas/ErrorResponse'
        404:
            description: Visit not found
            content:
                application/json:
                    schema:
                        $ref: '#/components/schemas/ErrorResponse'
    """
    existing = hospital_visit_repository.find_by_id(hospital_visit_id)
    if not existing:
        return jsonify({'error': 'NotFound'}), 404

    _, err = ensure_can_access_patient_data(str(existing.patientId))
    if err:
        return err

    updates = request.get_json(silent=True) or {}
    forbidden = {'_id', 'id', 'created_at', 'updated_at'}
    filtered_updates = {k: v for k, v in updates.items() if k not in forbidden}
    if not filtered_updates:
        return jsonify({'error': 'NoUpdatableFields'}), 400

    updated = hospital_visit_repository.update(hospital_visit_id, filtered_updates)
    if not updated:
        return jsonify({'error': 'NotFound'}), 404


    return jsonify(updated.model_dump(by_alias=True, mode='json'))


@hospital_visits_bp.route('/hospital-visits/<hospital_visit_id>', methods=['DELETE'])
@auth_required
def delete_hospital_visit_endpoint(hospital_visit_id: str):
    """Delete a hospital visit record.
    ---
    tags:
        - Hospital Visits
    security:
        - BearerAuth: []
    parameters:
        - in: path
          name: hospital_visit_id
          required: true
          schema:
            type: string
    responses:
        204:
            description: Visit deleted
        401:
            description: Unauthorized
        403:
            description: Forbidden
        404:
            description: Visit not found
    """
    existing = hospital_visit_repository.find_by_id(hospital_visit_id)
    if not existing:
        return jsonify({'error': 'NotFound'}), 404

    _, err = ensure_can_access_patient_data(str(existing.patientId))
    if err:
        return err

    success = hospital_visit_repository.delete(hospital_visit_id)
    if not success:
        return jsonify({'error': 'NotFound'}), 404

    return '', 204
