# routes/assistant.py
from __future__ import annotations

import re
from datetime import date, datetime
from calendar import monthrange
from typing import Dict, List, Optional, Tuple

from flask import Blueprint, jsonify, request
from flask_login import login_required, current_user

from sqlalchemy import or_, func

from extensions import db
from models import (
    Attendance, Holiday, LeaveBalance,
    PayrollPreviewControl, MonthlyPayroll, Payroll,
    FeedbackResponse,
    AssistantConversation, AssistantMessage,
    LeaveSettings,  # ✅ for FY range
)

assistant_bp = Blueprint("assistant", __name__, url_prefix="/assistant")


# -------------------------
# Helpers
# -------------------------

_MONTH_NAMES = {
    "jan": 1, "january": 1,
    "feb": 2, "february": 2,
    "mar": 3, "march": 3,
    "apr": 4, "april": 4,
    "may": 5,
    "jun": 6, "june": 6,
    "jul": 7, "july": 7,
    "aug": 8, "august": 8,
    "sep": 9, "sept": 9, "september": 9,
    "oct": 10, "october": 10,
    "nov": 11, "november": 11,
    "dec": 12, "december": 12,
}

def _month_key(year: int, month: int) -> str:
    return f"{year:04d}-{month:02d}"

def _safe_int(s: Optional[str], default: Optional[int] = None) -> Optional[int]:
    try:
        return int(s) if s is not None and str(s).strip() != "" else default
    except Exception:
        return default

def _parse_month_year(text: str) -> Tuple[int, int]:
    t = (text or "").strip().lower()

    m = re.search(r"\b(20\d{2})[-/](0?[1-9]|1[0-2])\b", t)
    if m:
        return int(m.group(1)), int(m.group(2))

    m2 = re.search(r"\b([a-z]{3,9})\s+(20\d{2})\b", t)
    if m2:
        mon = _MONTH_NAMES.get(m2.group(1).lower())
        if mon:
            return int(m2.group(2)), mon

    # ✅ use local server time (not UTC) to avoid month flip at midnight
    now = datetime.now()
    return now.year, now.month

def _parse_date(text: str) -> Optional[date]:
    t = (text or "").strip()
    m = re.search(r"\b(20\d{2})-(0?[1-9]|1[0-2])-(0?[1-9]|[12]\d|3[01])\b", t)
    if not m:
        return None
    try:
        return date(int(m.group(1)), int(m.group(2)), int(m.group(3)))
    except Exception:
        return None

def _fmt_time(tm) -> str:
    if not tm:
        return "—"
    try:
        return tm.strftime("%H:%M")
    except Exception:
        return str(tm)

def _card(title: str, url: str, subtitle: str = "") -> Dict[str, str]:
    return {"title": title, "action": "open_url", "url": url, "subtitle": subtitle}


# -------------------------
# FY helpers (Leave is FY-based in your system)
# -------------------------

def _get_leave_year_range() -> Tuple[date, date]:
    """
    Mirrors your leave.py logic:
    - If LeaveSettings.leave_year_start/end present -> use them.
    - Else fallback to Apr–Mar FY containing today.
    """
    settings = LeaveSettings.query.first()
    if settings and getattr(settings, "leave_year_start", None) and getattr(settings, "leave_year_end", None):
        return settings.leave_year_start, settings.leave_year_end

    today = datetime.today().date()
    fy_start = date(today.year if today.month >= 4 else today.year - 1, 4, 1)
    fy_end = date(fy_start.year + 1, 3, 31)
    return fy_start, fy_end

def _current_fy_label() -> int:
    fy_start, _ = _get_leave_year_range()
    return fy_start.year

def _fy_label_from_message(message: str) -> int:
    """
    Parse FY label from user message.

    Rules (simple & predictable):
    - If user says 'fy 2025' -> fy_label=2025
    - If user says '2025-26' -> fy_label=2025
    - If user says just '2026' -> treat as fy_label=2026 (FY starting 2026)
    - If nothing provided -> current FY label
    """
    low = (message or "").lower()

    # "fy 2025"
    m = re.search(r"\bfy\s*(20\d{2})\b", low)
    if m:
        return int(m.group(1))

    # "2025-26" or "2025/26"
    m2 = re.search(r"\b(20\d{2})\s*[-/]\s*(\d{2})\b", low)
    if m2:
        return int(m2.group(1))

    # any year "2026"
    m3 = re.search(r"\b(20\d{2})\b", low)
    if m3:
        return int(m3.group(1))

    return _current_fy_label()


# -------------------------
# Conversation storage
# -------------------------

def _get_or_create_conversation() -> AssistantConversation:
    conv = (
        AssistantConversation.query
        .filter_by(user_id=current_user.id, is_active=True)
        .order_by(AssistantConversation.id.desc())
        .first()
    )
    if conv:
        conv.updated_at = datetime.utcnow()
        db.session.commit()
        return conv

    conv = AssistantConversation(
        user_id=current_user.id,
        title="HR Assistant",
        is_active=True,
        created_at=datetime.utcnow(),
        updated_at=datetime.utcnow(),
    )
    db.session.add(conv)
    db.session.commit()
    return conv

def _save_msg(conv_id: int, role: str, content: str) -> None:
    m = AssistantMessage(
        conversation_id=conv_id,
        role=role,
        content=content or "",
        created_at=datetime.utcnow(),
    )
    db.session.add(m)
    AssistantConversation.query.filter_by(id=conv_id).update({"updated_at": datetime.utcnow()})
    db.session.commit()


# -------------------------
# Tools (read-only)
# -------------------------

def tool_my_leave_balance_from_msg(message: str) -> Tuple[str, List[Dict[str, str]]]:
    fy_label = _fy_label_from_message(message)

    rows = LeaveBalance.query.filter_by(employee_id=current_user.id, year=fy_label).all()
    if not rows:
        reply = (
            f"I couldn't find any leave balance rows for **FY {fy_label}** for your account. "
            "If balances were not allocated yet, HR may need to run allocation for this FY."
        )
        return reply, [
            _card("Open Leave History", "/leave/history"),
            _card("Apply Leave", "/leave/apply"),
            _card("Leave Allocation (HR)", "/leave/allocation", "If you are HR/Admin"),
        ]

    parts = []
    for r in rows:
        lt = getattr(r, "leave_type", None)
        name = getattr(lt, "name", None) or getattr(lt, "code", None) or f"Type #{r.leave_type_id}"
        parts.append(f"- **{name}**: {float(r.balance or 0):g}")

    reply = f"Your leave balances for **FY {fy_label}**:\n" + "\n".join(parts)
    return reply, [_card("Apply Leave", "/leave/apply"), _card("View Leave History", "/leave/history")]

def tool_my_attendance(month: int, year: int) -> Tuple[str, List[Dict[str, str]]]:
    start = date(year, month, 1)
    end = date(year, month, monthrange(year, month)[1])

    rows = (
        Attendance.query
        .filter(Attendance.employee_id == current_user.id)
        .filter(Attendance.date >= start)
        .filter(Attendance.date <= end)
        .order_by(Attendance.date.asc())
        .all()
    )

    if not rows:
        reply = (
            f"I couldn't find attendance records for **{_month_key(year, month)}**. "
            "If this month is not processed yet, HR can run monthly attendance processing."
        )
        return reply, [_card("Open Attendance Calendar", "/attendance/calendar")]

    status_counts: Dict[str, int] = {}
    for r in rows:
        key = (r.status or "").strip() or "unknown"
        status_counts[key] = status_counts.get(key, 0) + 1

    summary = ", ".join([f"{k}: {v}" for k, v in sorted(status_counts.items())])
    reply = (
        f"Your attendance summary for **{_month_key(year, month)}**: {summary}.\n"
        "Open the calendar for full day-wise details."
    )
    return reply, [_card("Open Attendance Calendar", f"/attendance/calendar?month={month:02d}&year={year}")]

def tool_explain_attendance(d: date) -> Tuple[str, List[Dict[str, str]]]:
    r = Attendance.query.filter_by(employee_id=current_user.id, date=d).first()

    if not r:
        reply = (
            f"I couldn't find an attendance entry for **{d.isoformat()}**. "
            "It may not have been processed/synced."
        )
        return reply, [_card("Open Attendance Calendar", "/attendance/calendar")]

    reply = (
        f"Attendance on **{d.isoformat()}**:\n"
        f"- **Status**: {r.status or '—'}\n"
        f"- **Check-in**: {_fmt_time(r.check_in_time)}\n"
        f"- **Check-out**: {_fmt_time(r.check_out_time)}\n"
        f"- **Working hours**: {float(r.working_hours or 0):g}\n"
        f"- **Source**: {r.source or '—'}\n"
        f"- **Remarks**: {r.remarks or '—'}"
    )
    return reply, [
        _card("Open Attendance Calendar", f"/attendance/calendar?focus={d.isoformat()}"),
        _card("View Leave History", "/leave/history"),
    ]

def tool_holidays(month: int, year: int) -> Tuple[str, List[Dict[str, str]]]:
    start = date(year, month, 1)
    end = date(year, month, monthrange(year, month)[1])

    user_loc = (getattr(current_user, "location", None) or "").strip()

    q = Holiday.query.filter(Holiday.date >= start, Holiday.date <= end)

    # ✅ Your system uses location='ALL' (and sometimes None/blank). Keep matches broad.
    if user_loc:
        ul = user_loc.upper()
        q = q.filter(
            or_(
                Holiday.location == None,  # noqa
                Holiday.location == "",
                func.upper(Holiday.location) == "ALL",
                func.upper(Holiday.location) == ul,
            )
        )

    rows = q.order_by(Holiday.date.asc()).all()
    if not rows:
        return f"No holidays found for **{_month_key(year, month)}**{(' ('+user_loc+')') if user_loc else ''}.", []

    lines = []
    for h in rows[:25]:
        typ = f" ({h.type})" if getattr(h, "type", None) else ""
        lines.append(f"- {h.date.strftime('%d %b %Y')}: **{h.name}**{typ}")

    reply = f"Holidays for **{_month_key(year, month)}**{(' ('+user_loc+')') if user_loc else ''}:\n" + "\n".join(lines)
    return reply, []

def tool_payroll_status(month: int, year: int) -> Tuple[str, List[Dict[str, str]]]:
    mk = _month_key(year, month)

    preview = PayrollPreviewControl.query.filter_by(month=mk).first()
    preview_enabled = bool(getattr(preview, "preview_enabled", False)) if preview else False

    frozen_count = (
        MonthlyPayroll.query
        .filter_by(employee_id=current_user.id, month=month, year=year)
        .count()
    )
    if frozen_count == 0:
        frozen_count = Payroll.query.filter_by(employee_id=current_user.id, month=mk).count()
    frozen = frozen_count > 0

    feedback_done = FeedbackResponse.query.filter_by(submitted_by_id=current_user.id, month=mk).first() is not None

    reply = (
        f"Payroll status for **{mk}**:\n"
        f"- Preview enabled: **{'Yes' if preview_enabled else 'No'}**\n"
        f"- Frozen for you: **{'Yes' if frozen else 'No'}**\n"
        f"- Feedback submitted: **{'Yes' if feedback_done else 'No'}**"
    )

    cards: List[Dict[str, str]] = [_card("My Salary Slips", "/my-salary-slips", "View available salary slips")]
    if not feedback_done:
        cards.append(_card("Submit Monthly Feedback", f"/feedback/form?month={mk}&next=/my-salary-slips", "Required to unlock this month"))
    return reply, cards

def tool_help() -> Tuple[str, List[Dict[str, str]]]:
    fy_now = _current_fy_label()
    reply = (
        "I can help you with:\n"
        f"- **Leave balance (FY-based)** (e.g., 'show my leave balance FY {fy_now}' or 'leave balance {fy_now}')\n"
        "- **Attendance summary** (e.g., 'attendance Jan 2026' or 'attendance 2026-01')\n"
        "- **Explain a day** (e.g., 'why LWP on 2026-01-12')\n"
        "- **Holidays** (e.g., 'holidays Feb 2026')\n"
        "- **Payroll status** (e.g., 'payroll status 2026-01')\n\n"
        "Note: v1 is read-only; actions will open the correct page."
    )
    return reply, [
        _card("Apply Leave", "/leave/apply"),
        _card("Attendance Calendar", "/attendance/calendar"),
        _card("My Salary Slips", "/my-salary-slips"),
    ]


# -------------------------
# Intent router (rule-based)
# -------------------------

def route_intent(message: str) -> Tuple[str, List[Dict[str, str]]]:
    msg = (message or "").strip()
    low = msg.lower()

    if any(k in low for k in ["help", "what can you do", "options", "commands"]):
        return tool_help()

    # ✅ Explain day: if user asks why/explain/reason + provides a date, explain it
    if any(k in low for k in ["why", "explain", "reason"]):
        d = _parse_date(low)
        if d:
            return tool_explain_attendance(d)

    # ✅ FY-based leave balance
    if "leave" in low and any(k in low for k in ["balance", "balances", "left", "remaining"]):
        return tool_my_leave_balance_from_msg(msg)

    if "attendance" in low or "punch" in low or "check-in" in low or "check in" in low:
        y, m = _parse_month_year(low)
        return tool_my_attendance(m, y)

    if "holiday" in low or "holidays" in low:
        y, m = _parse_month_year(low)
        return tool_holidays(m, y)

    if "payroll" in low or "salary" in low:
        y, m = _parse_month_year(low)
        return tool_payroll_status(m, y)

    if "apply leave" in low or ("leave" in low and "apply" in low):
        return ("Sure — opening the leave application page.", [_card("Apply Leave", "/leave/apply")])

    if "leave history" in low or ("leave" in low and "history" in low):
        return ("Opening your leave history.", [_card("Leave History", "/leave/history")])

    return (
        "I can help with leave balance (FY), attendance, holidays, and payroll status. Type **help** to see examples.",
        [_card("Help", "#", "Type: help")],
    )


# -------------------------
# Routes
# -------------------------

@assistant_bp.get("/history")
@login_required
def history():
    conv = _get_or_create_conversation()
    limit = _safe_int(request.args.get("limit"), 50) or 50
    limit = min(max(limit, 1), 200)

    msgs = (
        AssistantMessage.query
        .filter_by(conversation_id=conv.id)
        .order_by(AssistantMessage.id.asc())
        .limit(limit)
        .all()
    )

    return jsonify({
        "conversation_id": conv.id,
        "messages": [{"role": m.role, "content": m.content, "created_at": m.created_at.isoformat()} for m in msgs]
    })

@assistant_bp.post("/clear")
@login_required
def clear():
    conv = _get_or_create_conversation()
    AssistantMessage.query.filter_by(conversation_id=conv.id).delete()
    conv.updated_at = datetime.utcnow()
    db.session.commit()
    return jsonify({"ok": True})

@assistant_bp.post("/chat")
@login_required
def chat():
    data = request.get_json(silent=True) or {}
    message = (data.get("message") or "").strip()

    if not message:
        return jsonify({"reply": "Type a message to get started.", "cards": []})

    conv = _get_or_create_conversation()

    _save_msg(conv.id, "user", message)

    reply, cards = route_intent(message)

    _save_msg(conv.id, "assistant", reply)

    return jsonify({"reply": reply, "cards": cards})