"""
Medications API - medication scheduling, logging, and adherence tracking
"""

import json
import os
import uuid
from datetime import datetime, date, timedelta, timezone

import flask
import jwt
from psycopg2.extras import Json
import core.auth as auth
import core.postgres as postgres
import core.tz as tz


def _get_user_uuid(token):
    try:
        payload = jwt.decode(token, os.getenv("JWT_SECRET"), algorithms=["HS256"])
        return payload.get("sub")
    except (jwt.ExpiredSignatureError, jwt.InvalidTokenError):
        return None


def _auth(request):
    """Extract and verify token. Returns user_uuid or None."""
    header = request.headers.get("Authorization", "")
    if not header.startswith("Bearer "):
        return None
    token = header[7:]
    user_uuid = _get_user_uuid(token)
    if not user_uuid or not auth.verifyLoginToken(token, userUUID=user_uuid):
        return None
    return user_uuid


_JSON_COLS = {"times", "days_of_week"}


def _wrap_json(data):
    """Wrap list/dict values in Json() for psycopg2 when targeting JSON columns."""
    return {k: Json(v) if k in _JSON_COLS and isinstance(v, (list, dict)) else v for k, v in data.items()}


def _filter_times_in_range(times, start, end):
    """Return times (HH:MM strings) that fall within [start, end] inclusive."""
    return [t for t in times if start <= t <= end]


def _is_med_due_today(med, today, current_day):
    """Check if a medication is scheduled for today based on its frequency."""
    freq = med.get("frequency", "daily")
    if freq in ("daily", "twice_daily"):
        return True
    if freq == "as_needed":
        return True  # always show PRN, but with no scheduled times
    if freq == "specific_days":
        days = med.get("days_of_week", [])
        return current_day in days
    if freq == "every_n_days":
        start = med.get("start_date")
        interval = med.get("interval_days")
        if start and interval:
            start_d = start if isinstance(start, date) else datetime.strptime(str(start), "%Y-%m-%d").date()
            days_since = (today - start_d).days
            return days_since >= 0 and days_since % interval == 0
        return False
    return True


def _compute_next_dose_date(med):
    """Compute the next dose date for every_n_days medications."""
    interval = med.get("interval_days")
    if not interval:
        return None
    return (tz.user_today() + timedelta(days=interval)).isoformat()


def _count_expected_doses(med, period_start, days):
    """Count expected doses in a period based on frequency and schedule."""
    freq = med.get("frequency", "daily")
    times_per_day = len(med.get("times", [])) or 1

    if freq == "as_needed":
        return 0  # PRN has no expected doses
    if freq in ("daily", "twice_daily"):
        return days * times_per_day
    if freq == "specific_days":
        dow = med.get("days_of_week", [])
        count = 0
        for d in range(days):
            check_date = period_start + timedelta(days=d)
            if check_date.strftime("%a").lower() in dow:
                count += times_per_day
        return count
    if freq == "every_n_days":
        interval = med.get("interval_days", 1)
        start = med.get("start_date")
        if not start:
            return 0
        start_d = start if isinstance(start, date) else datetime.strptime(str(start), "%Y-%m-%d").date()
        count = 0
        for d in range(days):
            check_date = period_start + timedelta(days=d)
            diff = (check_date - start_d).days
            if diff >= 0 and diff % interval == 0:
                count += times_per_day
        return count
    return days * times_per_day


def _log_local_date(created_at, user_tz):
    """Convert a DB created_at (naive UTC datetime) to a local date string YYYY-MM-DD."""
    if created_at is None:
        return ""
    if isinstance(created_at, datetime):
        # Treat naive datetimes as UTC
        if created_at.tzinfo is None:
            created_at = created_at.replace(tzinfo=timezone.utc)
        return created_at.astimezone(user_tz).date().isoformat()
    # Fallback: already a string
    return str(created_at)[:10]


def _count_logs_in_period(logs, period_start_str, action, user_tz=None):
    """Count logs of a given action where created_at (local date) >= period_start."""
    return sum(
        1 for log in logs
        if log.get("action") == action
        and (_log_local_date(log.get("created_at"), user_tz) if user_tz else str(log.get("created_at", ""))[:10]) >= period_start_str
    )


def register(app):

    # ── Medications CRUD ──────────────────────────────────────────

    @app.route("/api/medications", methods=["GET"])
    def api_listMedications():
        """List all medications for the logged-in user."""
        user_uuid = _auth(flask.request)
        if not user_uuid:
            return flask.jsonify({"error": "unauthorized"}), 401
        meds = postgres.select("medications", where={"user_uuid": user_uuid}, order_by="name")
        return flask.jsonify(meds), 200

    def _time_str_to_minutes(time_str):
        """Convert 'HH:MM' to minutes since midnight."""
        parts = time_str.split(":")
        return int(parts[0]) * 60 + int(parts[1])

    def _get_routine_duration_minutes(routine_id):
        """Get total duration of a routine from its steps."""
        steps = postgres.select("routine_steps", where={"routine_id": routine_id})
        total = sum(s.get("duration_minutes", 0) or 0 for s in steps)
        return max(total, 1)

    def _check_med_schedule_conflicts(user_uuid, new_times, new_days=None, exclude_med_id=None):
        """Check if the proposed medication schedule conflicts with existing routines or medications.
        Returns (has_conflict, conflict_message) tuple.
        """
        if not new_times:
            return False, None

        # Check conflicts with routines (duration-aware)
        user_routines = postgres.select("routines", {"user_uuid": user_uuid})
        for r in user_routines:
            sched = postgres.select_one("routine_schedules", {"routine_id": r["id"]})
            if not sched or not sched.get("time"):
                continue
            routine_days = sched.get("days", [])
            if isinstance(routine_days, str):
                routine_days = json.loads(routine_days)
            if new_days and not any(d in routine_days for d in new_days):
                continue
            routine_start = _time_str_to_minutes(sched["time"])
            routine_dur = _get_routine_duration_minutes(r["id"])
            for t in new_times:
                med_start = _time_str_to_minutes(t)
                # Med falls within routine time range
                if routine_start <= med_start < routine_start + routine_dur:
                    return True, f"Time conflicts with routine: {r.get('name', 'Unnamed routine')}"

        # Check conflicts with other medications
        user_meds = postgres.select("medications", {"user_uuid": user_uuid, "active": True})
        for med in user_meds:
            if med["id"] == exclude_med_id:
                continue
            med_times = med.get("times", [])
            if isinstance(med_times, str):
                med_times = json.loads(med_times)
            if any(t in med_times for t in new_times):
                return True, f"Time conflicts with medication: {med.get('name', 'Unnamed medication')}"

        return False, None

    @app.route("/api/medications", methods=["POST"])
    def api_addMedication():
        """Add a medication. Body: {name, dosage, unit, frequency, times?, days_of_week?, interval_days?, start_date?, notes?}"""
        user_uuid = _auth(flask.request)
        if not user_uuid:
            return flask.jsonify({"error": "unauthorized"}), 401
        data = flask.request.get_json()
        if not data:
            return flask.jsonify({"error": "missing body"}), 400
        required = ["name", "dosage", "unit", "frequency"]
        missing = [f for f in required if not data.get(f)]
        if missing:
            return flask.jsonify({"error": f"missing required fields: {', '.join(missing)}"}), 400

        # Validate every_n_days required fields
        if data.get("frequency") == "every_n_days":
            if not data.get("start_date") or not data.get("interval_days"):
                return flask.jsonify({"error": "every_n_days frequency requires both start_date and interval_days"}), 400

        # Check for schedule conflicts
        new_times = data.get("times", [])
        new_days = data.get("days_of_week", [])
        has_conflict, conflict_msg = _check_med_schedule_conflicts(
            user_uuid, new_times, new_days
        )
        if has_conflict:
            return flask.jsonify({"error": conflict_msg}), 409

        row = {
            "id": str(uuid.uuid4()),
            "user_uuid": user_uuid,
            "name": data["name"],
            "dosage": data["dosage"],
            "unit": data["unit"],
            "frequency": data["frequency"],
            "times": data.get("times", []),
            "days_of_week": data.get("days_of_week", []),
            "interval_days": data.get("interval_days"),
            "start_date": data.get("start_date"),
            "notes": data.get("notes"),
            "active": True,
        }

        # Compute next_dose_date for interval meds
        if data.get("frequency") == "every_n_days" and data.get("start_date") and data.get("interval_days"):
            start = datetime.strptime(data["start_date"], "%Y-%m-%d").date()
            today = tz.user_today()
            if start > today:
                row["next_dose_date"] = data["start_date"]
            else:
                days_since = (today - start).days
                remainder = days_since % data["interval_days"]
                if remainder == 0:
                    row["next_dose_date"] = today.isoformat()
                else:
                    row["next_dose_date"] = (today + timedelta(days=data["interval_days"] - remainder)).isoformat()

        med = postgres.insert("medications", _wrap_json(row))
        return flask.jsonify(med), 201

    @app.route("/api/medications/<med_id>", methods=["GET"])
    def api_getMedication(med_id):
        """Get a single medication with its schedule."""
        user_uuid = _auth(flask.request)
        if not user_uuid:
            return flask.jsonify({"error": "unauthorized"}), 401
        med = postgres.select_one("medications", {"id": med_id, "user_uuid": user_uuid})
        if not med:
            return flask.jsonify({"error": "not found"}), 404
        return flask.jsonify(med), 200

    @app.route("/api/medications/<med_id>", methods=["PUT"])
    def api_updateMedication(med_id):
        """Update medication details."""
        user_uuid = _auth(flask.request)
        if not user_uuid:
            return flask.jsonify({"error": "unauthorized"}), 401
        data = flask.request.get_json()
        if not data:
            return flask.jsonify({"error": "missing body"}), 400
        existing = postgres.select_one("medications", {"id": med_id, "user_uuid": user_uuid})
        if not existing:
            return flask.jsonify({"error": "not found"}), 404
        allowed = [
            "name", "dosage", "unit", "frequency", "times", "notes", "active",
            "days_of_week", "interval_days", "start_date", "next_dose_date",
        ]

        # Check for schedule conflicts if times are being updated
        if "times" in data:
            new_times = data.get("times", [])
            new_days = data.get("days_of_week") or existing.get("days_of_week", [])
            has_conflict, conflict_msg = _check_med_schedule_conflicts(
                user_uuid, new_times, new_days, exclude_med_id=med_id
            )
            if has_conflict:
                return flask.jsonify({"error": conflict_msg}), 409

        updates = {k: v for k, v in data.items() if k in allowed}
        if not updates:
            return flask.jsonify({"error": "no valid fields to update"}), 400
        result = postgres.update("medications", _wrap_json(updates), {"id": med_id, "user_uuid": user_uuid})
        return flask.jsonify(result[0] if result else {}), 200

    @app.route("/api/medications/<med_id>", methods=["DELETE"])
    def api_deleteMedication(med_id):
        """Delete a medication and its logs."""
        user_uuid = _auth(flask.request)
        if not user_uuid:
            return flask.jsonify({"error": "unauthorized"}), 401
        existing = postgres.select_one("medications", {"id": med_id, "user_uuid": user_uuid})
        if not existing:
            return flask.jsonify({"error": "not found"}), 404
        postgres.delete("med_logs", {"medication_id": med_id})
        postgres.delete("medications", {"id": med_id, "user_uuid": user_uuid})
        return flask.jsonify({"deleted": True}), 200

    # ── Medication Logging (take / skip / snooze) ─────────────────

    @app.route("/api/medications/<med_id>/take", methods=["POST"])
    def api_takeMedication(med_id):
        """Log that a dose was taken. Body: {scheduled_time?, notes?}"""
        user_uuid = _auth(flask.request)
        if not user_uuid:
            return flask.jsonify({"error": "unauthorized"}), 401
        med = postgres.select_one("medications", {"id": med_id, "user_uuid": user_uuid})
        if not med:
            return flask.jsonify({"error": "not found"}), 404
        data = flask.request.get_json() or {}
        log_entry = {
            "id": str(uuid.uuid4()),
            "medication_id": med_id,
            "user_uuid": user_uuid,
            "action": "taken",
            "scheduled_time": data.get("scheduled_time"),
            "notes": data.get("notes"),
        }
        log = postgres.insert("med_logs", log_entry)
        # Advance next_dose_date for interval meds
        if med.get("frequency") == "every_n_days" and med.get("interval_days"):
            next_date = _compute_next_dose_date(med)
            if next_date:
                postgres.update("medications", {"next_dose_date": next_date}, {"id": med_id})
        return flask.jsonify(log), 201

    @app.route("/api/medications/<med_id>/skip", methods=["POST"])
    def api_skipMedication(med_id):
        """Log a skipped dose. Body: {scheduled_time?, reason?}"""
        user_uuid = _auth(flask.request)
        if not user_uuid:
            return flask.jsonify({"error": "unauthorized"}), 401
        med = postgres.select_one("medications", {"id": med_id, "user_uuid": user_uuid})
        if not med:
            return flask.jsonify({"error": "not found"}), 404
        data = flask.request.get_json() or {}
        log_entry = {
            "id": str(uuid.uuid4()),
            "medication_id": med_id,
            "user_uuid": user_uuid,
            "action": "skipped",
            "scheduled_time": data.get("scheduled_time"),
            "notes": data.get("reason"),
        }
        log = postgres.insert("med_logs", log_entry)
        # Advance next_dose_date for interval meds
        if med.get("frequency") == "every_n_days" and med.get("interval_days"):
            next_date = _compute_next_dose_date(med)
            if next_date:
                postgres.update("medications", {"next_dose_date": next_date}, {"id": med_id})
        return flask.jsonify(log), 201

    @app.route("/api/medications/<med_id>/snooze", methods=["POST"])
    def api_snoozeMedication(med_id):
        """Snooze a reminder. Body: {minutes: 15}"""
        user_uuid = _auth(flask.request)
        if not user_uuid:
            return flask.jsonify({"error": "unauthorized"}), 401
        med = postgres.select_one("medications", {"id": med_id, "user_uuid": user_uuid})
        if not med:
            return flask.jsonify({"error": "not found"}), 404
        data = flask.request.get_json() or {}
        minutes = data.get("minutes", 15)
        return flask.jsonify({"snoozed_until_minutes": minutes}), 200

    # ── Medication Log / History ──────────────────────────────────

    @app.route("/api/medications/<med_id>/log", methods=["GET"])
    def api_getMedLog(med_id):
        """Get dose log for a medication. Query: ?days=30"""
        user_uuid = _auth(flask.request)
        if not user_uuid:
            return flask.jsonify({"error": "unauthorized"}), 401
        med = postgres.select_one("medications", {"id": med_id, "user_uuid": user_uuid})
        if not med:
            return flask.jsonify({"error": "not found"}), 404
        days = flask.request.args.get("days", 30, type=int)
        logs = postgres.select(
            "med_logs",
            where={"medication_id": med_id},
            order_by="created_at DESC",
            limit=days * 10,
        )
        return flask.jsonify(logs), 200

    @app.route("/api/medications/today", methods=["GET"])
    def api_todaysMeds():
        """Get today's medication schedule with taken/pending status.

        Includes cross-midnight lookahead:
        - After 22:00: includes next-day meds scheduled 00:00-02:00 (is_next_day)
        - Before 02:00: includes previous-day meds scheduled 22:00-23:59 (is_previous_day)
        """
        user_uuid = _auth(flask.request)
        if not user_uuid:
            return flask.jsonify({"error": "unauthorized"}), 401

        meds = postgres.select("medications", where={"user_uuid": user_uuid, "active": True})
        now = tz.user_now()
        user_tz = now.tzinfo
        today = now.date()
        today_str = today.isoformat()
        current_day = now.strftime("%a").lower()  # "mon","tue", etc.
        current_hour = now.hour

        result = []
        seen_med_ids = set()

        # Main pass: today's meds
        for med in meds:
            if not _is_med_due_today(med, today, current_day):
                continue

            freq = med.get("frequency", "daily")
            is_prn = freq == "as_needed"

            all_logs = postgres.select(
                "med_logs",
                where={"medication_id": med["id"]},
            )
            today_taken = [
                log.get("scheduled_time") or ""
                for log in all_logs
                if log.get("action") == "taken"
                and _log_local_date(log.get("created_at"), user_tz) == today_str
            ]
            today_skipped = [
                log.get("scheduled_time") or ""
                for log in all_logs
                if log.get("action") == "skipped"
                and _log_local_date(log.get("created_at"), user_tz) == today_str
            ]

            result.append({
                "medication": med,
                "scheduled_times": [] if is_prn else med.get("times", []),
                "taken_times": today_taken,
                "skipped_times": today_skipped,
                "is_prn": is_prn,
            })
            seen_med_ids.add(med["id"])

        # Late night pass (22:00+): include next-day meds scheduled 00:00-02:00
        if current_hour >= 22:
            tomorrow = today + timedelta(days=1)
            tomorrow_str = tomorrow.isoformat()
            tomorrow_day = tomorrow.strftime("%a").lower()
            for med in meds:
                if med["id"] in seen_med_ids:
                    continue
                if not _is_med_due_today(med, tomorrow, tomorrow_day):
                    continue
                freq = med.get("frequency", "daily")
                if freq == "as_needed":
                    continue
                times = med.get("times", [])
                early_times = _filter_times_in_range(times, "00:00", "02:00")
                if not early_times:
                    continue

                all_logs = postgres.select(
                    "med_logs",
                    where={"medication_id": med["id"]},
                )
                tomorrow_taken = [
                    log.get("scheduled_time") or ""
                    for log in all_logs
                    if log.get("action") == "taken"
                    and _log_local_date(log.get("created_at"), user_tz) == tomorrow_str
                ]
                tomorrow_skipped = [
                    log.get("scheduled_time") or ""
                    for log in all_logs
                    if log.get("action") == "skipped"
                    and _log_local_date(log.get("created_at"), user_tz) == tomorrow_str
                ]

                result.append({
                    "medication": med,
                    "scheduled_times": early_times,
                    "taken_times": tomorrow_taken,
                    "skipped_times": tomorrow_skipped,
                    "is_prn": False,
                    "is_next_day": True,
                })
                seen_med_ids.add(med["id"])

        # Early morning pass (<02:00): include previous-day meds scheduled 22:00-23:59
        if current_hour < 2:
            yesterday = today - timedelta(days=1)
            yesterday_str = yesterday.isoformat()
            yesterday_day = yesterday.strftime("%a").lower()
            for med in meds:
                if med["id"] in seen_med_ids:
                    continue
                if not _is_med_due_today(med, yesterday, yesterday_day):
                    continue
                freq = med.get("frequency", "daily")
                if freq == "as_needed":
                    continue
                times = med.get("times", [])
                late_times = _filter_times_in_range(times, "22:00", "23:59")
                if not late_times:
                    continue

                all_logs = postgres.select(
                    "med_logs",
                    where={"medication_id": med["id"]},
                )
                yesterday_taken = [
                    log.get("scheduled_time") or ""
                    for log in all_logs
                    if log.get("action") == "taken"
                    and _log_local_date(log.get("created_at"), user_tz) == yesterday_str
                ]
                yesterday_skipped = [
                    log.get("scheduled_time") or ""
                    for log in all_logs
                    if log.get("action") == "skipped"
                    and _log_local_date(log.get("created_at"), user_tz) == yesterday_str
                ]

                result.append({
                    "medication": med,
                    "scheduled_times": late_times,
                    "taken_times": yesterday_taken,
                    "skipped_times": yesterday_skipped,
                    "is_prn": False,
                    "is_previous_day": True,
                })
                seen_med_ids.add(med["id"])

        return flask.jsonify(result), 200

    # ── Adherence Stats ───────────────────────────────────────────

    @app.route("/api/medications/adherence", methods=["GET"])
    def api_adherenceStats():
        """Get adherence stats across all meds. Query: ?days=30"""
        user_uuid = _auth(flask.request)
        if not user_uuid:
            return flask.jsonify({"error": "unauthorized"}), 401
        num_days = flask.request.args.get("days", 30, type=int)
        meds = postgres.select("medications", where={"user_uuid": user_uuid, "active": True})
        now = tz.user_now()
        user_tz = now.tzinfo
        today = now.date()
        period_start = today - timedelta(days=num_days)
        period_start_str = period_start.isoformat()

        result = []
        for med in meds:
            freq = med.get("frequency", "daily")
            is_prn = freq == "as_needed"
            expected = _count_expected_doses(med, period_start, num_days)

            logs = postgres.select("med_logs", where={"medication_id": med["id"]})
            taken = _count_logs_in_period(logs, period_start_str, "taken", user_tz)
            skipped = _count_logs_in_period(logs, period_start_str, "skipped", user_tz)

            if is_prn:
                adherence_pct = None
            elif expected > 0:
                adherence_pct = round(min(taken / expected * 100, 100), 1)
            else:
                adherence_pct = 0

            result.append({
                "medication_id": med["id"],
                "name": med["name"],
                "taken": taken,
                "skipped": skipped,
                "expected": expected,
                "adherence_percent": adherence_pct,
                "is_prn": is_prn,
            })
        return flask.jsonify(result), 200

    @app.route("/api/medications/<med_id>/adherence", methods=["GET"])
    def api_medAdherence(med_id):
        """Get adherence stats for a single medication. Query: ?days=30"""
        user_uuid = _auth(flask.request)
        if not user_uuid:
            return flask.jsonify({"error": "unauthorized"}), 401
        med = postgres.select_one("medications", {"id": med_id, "user_uuid": user_uuid})
        if not med:
            return flask.jsonify({"error": "not found"}), 404
        num_days = flask.request.args.get("days", 30, type=int)
        now = tz.user_now()
        user_tz = now.tzinfo
        today = now.date()
        period_start = today - timedelta(days=num_days)
        period_start_str = period_start.isoformat()

        freq = med.get("frequency", "daily")
        is_prn = freq == "as_needed"
        expected = _count_expected_doses(med, period_start, num_days)

        logs = postgres.select("med_logs", where={"medication_id": med_id})
        taken = _count_logs_in_period(logs, period_start_str, "taken", user_tz)
        skipped = _count_logs_in_period(logs, period_start_str, "skipped", user_tz)

        if is_prn:
            adherence_pct = None
        elif expected > 0:
            adherence_pct = round(min(taken / expected * 100, 100), 1)
        else:
            adherence_pct = 0

        return flask.jsonify({
            "medication_id": med_id,
            "name": med["name"],
            "taken": taken,
            "skipped": skipped,
            "expected": expected,
            "adherence_percent": adherence_pct,
            "is_prn": is_prn,
        }), 200

    # ── Refills ───────────────────────────────────────────────────

    @app.route("/api/medications/<med_id>/refill", methods=["PUT"])
    def api_setRefill(med_id):
        """Set refill info. Body: {quantity_remaining, refill_date?, pharmacy_notes?}"""
        user_uuid = _auth(flask.request)
        if not user_uuid:
            return flask.jsonify({"error": "unauthorized"}), 401
        med = postgres.select_one("medications", {"id": med_id, "user_uuid": user_uuid})
        if not med:
            return flask.jsonify({"error": "not found"}), 404
        data = flask.request.get_json()
        if not data:
            return flask.jsonify({"error": "missing body"}), 400
        allowed = ["quantity_remaining", "refill_date", "pharmacy_notes"]
        updates = {k: v for k, v in data.items() if k in allowed}
        if not updates:
            return flask.jsonify({"error": "no valid fields to update"}), 400
        result = postgres.update("medications", updates, {"id": med_id, "user_uuid": user_uuid})
        return flask.jsonify(result[0] if result else {}), 200

    @app.route("/api/medications/refills-due", methods=["GET"])
    def api_refillsDue():
        """Get medications that need refills soon. Query: ?days_ahead=7"""
        user_uuid = _auth(flask.request)
        if not user_uuid:
            return flask.jsonify({"error": "unauthorized"}), 401
        days_ahead = flask.request.args.get("days_ahead", 7, type=int)
        cutoff = (tz.user_today() + timedelta(days=days_ahead)).isoformat()
        meds = postgres.select(
            "medications",
            where={"user_uuid": user_uuid},
        )
        due = []
        for med in meds:
            qty = med.get("quantity_remaining")
            refill_date = med.get("refill_date")
            if qty is not None and qty <= 7:
                due.append(med)
            elif refill_date and str(refill_date) <= cutoff:
                due.append(med)
        return flask.jsonify(due), 200