"""Staff scheduling — greedy shift assignment from demand signal + staff availability.

Assigns staff to demand slots using a cost-minimizing greedy algorithm:
    1. Sort staff by hourly rate (cheapest first)
    2. For each slot sorted by demand (highest first), assign cheapest available
       staff until coverage is met, respecting max_hours constraints
    3. Report coverage gaps where assigned hours < needed hours

Backends:
    - builtin: Greedy cost-minimizing assignment (no dependencies)
    - ortools: Constraint-based scheduling via Google OR-Tools (optimal)

Usage via provider:
    from platoon.learning.providers import get_provider
    sched = get_provider("scheduling")
    result = sched.schedule(demand_by_slot, staff, shift_hours=8)

Direct usage:
    from platoon.learning.scheduling import schedule_staff
    result = schedule_staff(demand_by_slot, staff)
"""

from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional

from platoon.learning.providers import ModelProvider, register_provider


# ---------------------------------------------------------------------------
# Result types
# ---------------------------------------------------------------------------


@dataclass
class ShiftAssignment:
    """A single staff-to-slot assignment."""

    staff_id: str
    staff_name: str
    day: str  # day label (e.g. "Monday", "2025-03-01")
    slot_label: str  # slot identifier (e.g. "Monday_morning")
    hours: float
    cost: float


@dataclass
class ScheduleResult:
    """Output of staff scheduling."""

    assignments: List[ShiftAssignment]
    total_cost: float
    total_hours: float
    coverage_by_slot: Dict[str, float]  # slot -> hours assigned
    demand_by_slot: Dict[str, float]  # slot -> hours needed
    coverage_gaps: List[Dict[str, Any]]  # [{slot, needed, assigned, gap}]
    schedule_grid: Dict[str, List[str]]  # day -> [staff names]


# ---------------------------------------------------------------------------
# Pure-Python implementation
# ---------------------------------------------------------------------------


def _parse_skills(staff_dict: Dict[str, Any]) -> set:
    """Extract skills from a staff dict as a lowercase set."""
    skills = staff_dict.get("skills", [])
    if isinstance(skills, str):
        skills = [s.strip() for s in skills.split(",") if s.strip()]
    return {s.lower() for s in skills}


def schedule_staff(
    demand_by_slot: Dict[str, float],
    staff: List[Dict[str, Any]],
    shift_hours: float = 8.0,
    min_coverage: float = 1.0,
    required_skills_by_slot: Optional[Dict[str, List[str]]] = None,
) -> ScheduleResult:
    """Assign staff to demand slots using greedy cost minimization.

    Args:
        demand_by_slot: {slot_label: hours_needed}. Slot labels should encode
            day info (e.g. "Monday", "Tuesday_morning").
        staff: List of staff dicts with keys: staff_id, name, hourly_rate,
            max_hours_per_week, available_days (list of day names).
        shift_hours: Hours per shift assignment (default 8).
        min_coverage: Minimum fraction of demand to cover per slot (default 1.0 = 100%).
        required_skills_by_slot: Optional {slot_label: [skill1, skill2, ...]}.
            When set, only staff with at least one matching skill (case-insensitive)
            can be assigned to that slot. None = no skill filtering (backward compatible).

    Returns:
        ScheduleResult with assignments, costs, and coverage analysis.
    """
    # Sort staff by hourly rate (cheapest first)
    sorted_staff = sorted(staff, key=lambda s: float(s.get("hourly_rate", 0)))

    # Track hours assigned to each staff member
    staff_hours: Dict[str, float] = {s["staff_id"]: 0.0 for s in staff}
    max_hours: Dict[str, float] = {
        s["staff_id"]: float(s.get("max_hours_per_week", 40))
        for s in staff
    }
    staff_available_days: Dict[str, set] = {}
    for s in staff:
        days = s.get("available_days", [])
        if isinstance(days, str):
            days = [d.strip() for d in days.split(",")]
        staff_available_days[s["staff_id"]] = set(d.lower() for d in days)

    staff_lookup = {s["staff_id"]: s for s in staff}

    # Sort slots by demand descending (prioritize busiest slots)
    sorted_slots = sorted(demand_by_slot.items(), key=lambda x: x[1], reverse=True)

    assignments: List[ShiftAssignment] = []
    coverage: Dict[str, float] = {slot: 0.0 for slot in demand_by_slot}
    schedule_grid: Dict[str, List[str]] = {}

    for slot_label, hours_needed in sorted_slots:
        target = hours_needed * min_coverage
        # Extract day from slot label (first part before underscore, or whole label)
        day = slot_label.split("_")[0] if "_" in slot_label else slot_label

        for s in sorted_staff:
            if coverage[slot_label] >= target:
                break

            sid = s["staff_id"]
            rate = float(s.get("hourly_rate", 0))

            # Check availability
            if staff_available_days.get(sid):
                if day.lower() not in staff_available_days[sid]:
                    continue

            # Check skill requirements
            if required_skills_by_slot:
                slot_skills = required_skills_by_slot.get(slot_label, [])
                if slot_skills:
                    staff_skills = _parse_skills(s)
                    if not any(sk.lower() in staff_skills for sk in slot_skills):
                        continue

            # Check max hours
            remaining = max_hours[sid] - staff_hours[sid]
            if remaining < shift_hours:
                continue

            # Assign
            hours = min(shift_hours, target - coverage[slot_label])
            cost = hours * rate

            assignments.append(ShiftAssignment(
                staff_id=sid,
                staff_name=s.get("name", sid),
                day=day,
                slot_label=slot_label,
                hours=hours,
                cost=round(cost, 2),
            ))

            coverage[slot_label] += hours
            staff_hours[sid] += hours

            if day not in schedule_grid:
                schedule_grid[day] = []
            schedule_grid[day].append(s.get("name", sid))

    # Compute gaps
    gaps = []
    for slot_label, needed in demand_by_slot.items():
        assigned = coverage[slot_label]
        if assigned < needed * min_coverage:
            gaps.append({
                "slot": slot_label,
                "needed": round(needed, 1),
                "assigned": round(assigned, 1),
                "gap": round(needed * min_coverage - assigned, 1),
            })

    total_cost = sum(a.cost for a in assignments)
    total_hours = sum(a.hours for a in assignments)

    return ScheduleResult(
        assignments=assignments,
        total_cost=round(total_cost, 2),
        total_hours=round(total_hours, 1),
        coverage_by_slot={k: round(v, 1) for k, v in coverage.items()},
        demand_by_slot=demand_by_slot,
        coverage_gaps=gaps,
        schedule_grid=schedule_grid,
    )


# ---------------------------------------------------------------------------
# Built-in provider
# ---------------------------------------------------------------------------


@register_provider
class BuiltinSchedulingProvider(ModelProvider):
    """Greedy cost-minimizing staff scheduling — no dependencies."""

    name = "builtin_scheduling"
    domain = "scheduling"
    backend = "builtin"

    def schedule(
        self,
        demand_by_slot: Dict[str, float],
        staff: List[Dict[str, Any]],
        shift_hours: float = 8.0,
        min_coverage: float = 1.0,
        required_skills_by_slot: Optional[Dict[str, List[str]]] = None,
    ) -> ScheduleResult:
        """Assign staff to demand slots."""
        return schedule_staff(
            demand_by_slot, staff, shift_hours=shift_hours,
            min_coverage=min_coverage, required_skills_by_slot=required_skills_by_slot,
        )


# ---------------------------------------------------------------------------
# Library stub
# ---------------------------------------------------------------------------


@register_provider
class ORToolsSchedulingProvider(ModelProvider):
    """Constraint-based scheduling via Google OR-Tools.

    TODO: Implement schedule() using ortools.sat.python.cp_model for optimal
    shift assignment with hard/soft constraints (min rest, max consecutive, etc.).
    Requires: ortools
    """

    name = "ortools_scheduling"
    domain = "scheduling"
    backend = "ortools"

    @classmethod
    def is_available(cls) -> bool:
        try:
            import ortools  # noqa: F401
            return True
        except ImportError:
            return False

    def schedule(
        self,
        demand_by_slot: Dict[str, float],
        staff: List[Dict[str, Any]],
        shift_hours: float = 8.0,
        min_coverage: float = 1.0,
        required_skills_by_slot: Optional[Dict[str, List[str]]] = None,
    ) -> ScheduleResult:
        """Assign staff to demand slots using constraint programming."""
        raise NotImplementedError("ortools scheduling provider not yet implemented")
