#!/usr/bin/env python3
"""
apply_budget_limits.py

Inspects the existing limit configured for the current month on each budget
to determine its period (weekly, monthly, custom) and limit amount. It then backfills 
that exact same limit and period length going back N years (default 3).

USAGE:

    export FIREFLY_URL="http://192.168.1.239:8012"
    export FIREFLY_TOKEN="..."

    python3 apply_budget_limits.py --list                 # list current budgets and their limits/periods
    python3 apply_budget_limits.py                       # dry run preview of backfill
    python3 apply_budget_limits.py --apply                # create missing limits
    python3 apply_budget_limits.py --apply --overwrite    # overwrite existing past limits
"""

import argparse
import calendar
import datetime
import os
import sys
import time

try:
    import requests
except ImportError:
    sys.exit("This script needs the 'requests' package: pip install requests")


def clean(value):
    if value is None:
        return ""
    return "".join(ch for ch in str(value).strip().strip('"').strip("'") if ch.isprintable())


class FireflyClient:
    def __init__(self, base_url, token, timeout=30):
        base_url = clean(base_url)
        token = clean(token)
        if not token:
            sys.exit("No Firefly API token provided. Check --token / FIREFLY_TOKEN.")
        if not base_url:
            sys.exit("No Firefly base URL provided. Check --url / FIREFLY_URL.")

        self.base_url = base_url.rstrip("/")
        self.session = requests.Session()
        self.session.headers.update({
            "Authorization": f"Bearer {token}",
            "Accept": "application/json",
            "Content-Type": "application/json",
        })
        self.timeout = timeout

    def _url(self, path):
        return f"{self.base_url}/api/v1{path}"

    def _diagnose_and_raise(self, resp):
        try:
            body = resp.text[:1000]
        except Exception:
            body = "<no body>"
        raise RuntimeError(
            f"{resp.status_code} {resp.reason} for {resp.request.method} {resp.url}\n"
            f"Response body: {body}"
        )

    def _get_all(self, path, params=None):
        items = []
        url = self._url(path)
        first = True
        while url:
            resp = self.session.get(url, params=params if first else None, timeout=self.timeout)
            if not resp.ok:
                self._diagnose_and_raise(resp)
            payload = resp.json()
            items.extend(payload.get("data", []))
            url = payload.get("links", {}).get("next")
            first = False
        return items

    def about(self):
        resp = self.session.get(self._url("/about"), timeout=self.timeout)
        if not resp.ok:
            self._diagnose_and_raise(resp)
        return resp.json().get("data", {})

    def list_budgets(self):
        return self._get_all("/budgets")

    def get_budget_limits(self, budget_id, start, end):
        return self._get_all(f"/budgets/{budget_id}/limits", params={"start": start, "end": end})

    # def create_budget_limit(self, budget_id, start, end, amount, currency_id=None, currency_code=None):
    #     body = {
    #         "budget_id": str(budget_id),
    #         "start": start,
    #         "end": end,
    #         "amount": str(amount),
    #     }
    #     if currency_id:
    #         body["currency_id"] = str(currency_id)
    #     elif currency_code:
    #         body["currency_code"] = currency_code

    #     resp = self.session.post(self._url("/budget-limits"), json=body, timeout=self.timeout)
    #     if not resp.ok:
    #         self._diagnose_and_raise(resp)
    #     return resp.json()["data"]

    def create_budget_limit(self, budget_id, start, end, amount, currency_id=None, currency_code=None):
        body = {
            "start": start,
            "end": end,
            "amount": str(amount),
        }
        if currency_id is not None and str(currency_id).isdigit():
            body["currency_id"] = int(currency_id)  # Must be an integer
        elif currency_code:
            body["currency_code"] = currency_code

        resp = self.session.post(self._url(f"/budgets/{budget_id}/limits"), json=body, timeout=self.timeout)
        if not resp.ok:
            self._diagnose_and_raise(resp)
        return resp.json()["data"]

    def update_budget_limit(self, limit_id, amount):
        resp = self.session.put(self._url(f"/budget-limits/{limit_id}"),
                                 json={"amount": str(amount)}, timeout=self.timeout)
        if not resp.ok:
            self._diagnose_and_raise(resp)
        return resp.json()["data"]


def month_bounds(year, month):
    start = datetime.date(year, month, 1)
    last_day = calendar.monthrange(year, month)[1]
    end = datetime.date(year, month, last_day)
    return start.isoformat(), end.isoformat()


def detect_period_type(limit_start_str, limit_end_str):
    """Detects whether a limit date range represents a monthly or weekly period."""
    s = datetime.date.fromisoformat(limit_start_str[:10])
    e = datetime.date.fromisoformat(limit_end_str[:10])
    days_diff = (e - s).days + 1

    if 6 <= days_diff <= 8:
        return "weekly"
    elif 27 <= days_diff <= 32:
        return "monthly"
    return f"custom ({days_diff} days)"


def generate_past_periods(today, years_back, period_type):
    """Generates (start_str, end_str) tuples based on auto-detected period type."""
    start_history = datetime.date(today.year - years_back, today.month, 1)
    _, last_day = calendar.monthrange(today.year, today.month)
    end_history = datetime.date(today.year, today.month, last_day)

    periods = []

    if period_type == "weekly":
        # Align to Monday–Sunday weeks
        curr_start = start_history - datetime.timedelta(days=start_history.weekday())
        while curr_start <= end_history:
            curr_end = curr_start + datetime.timedelta(days=6)
            periods.append((curr_start.isoformat(), curr_end.isoformat()))
            curr_start += datetime.timedelta(days=7)

    else:  # monthly / default fallback
        curr_y, curr_m = start_history.year, start_history.month
        while (curr_y, curr_m) <= (today.year, today.month):
            s, e = month_bounds(curr_y, curr_m)
            periods.append((s, e))
            curr_m += 1
            if curr_m > 12:
                curr_m = 1
                curr_y += 1

    return periods


def main():
    parser = argparse.ArgumentParser(description=__doc__,
                                      formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("--url", default=os.environ.get("FIREFLY_URL"))
    parser.add_argument("--token", default=os.environ.get("FIREFLY_TOKEN"))
    parser.add_argument("--years", type=int, default=3,
                         help="How many years back to apply the limit (default: 3)")
    parser.add_argument("-l", "--list", action="store_true",
                         help="Print defined budget limits for the current month and exit.")
    parser.add_argument("--apply", action="store_true",
                         help="Actually create/update budget limits. Without this, only preview.")
    parser.add_argument("--overwrite", action="store_true",
                         help="Also overwrite limits that already exist in past periods.")
    parser.add_argument("--delay", type=float, default=0.2,
                         help="Seconds to sleep between API calls (default: 0.2)")
    args = parser.parse_args()

    client = FireflyClient(args.url, args.token)

    print(f"Connecting to {client.base_url} ...")
    about = client.about()
    print(f"  Connected OK. Firefly III version: {about.get('version', 'unknown')}\n")

    today = datetime.date.today()
    cur_start, cur_end = month_bounds(today.year, today.month)

    budgets = client.list_budgets()

    # --- LIST MODE (--list / -l) ---
    if args.list:
        print(f"Current Month Budget Summary ({cur_start} to {cur_end}):")
        print("=" * 75)
        print(f"{'Budget Name':<25} | {'Amount':<10} | {'Period':<10} | {'Start Date':<10} | {'End Date':<10}")
        print("-" * 75)
        
        no_limits_count = 0
        for b in budgets:
            budget_id = b["id"]
            budget_name = b["attributes"]["name"]
            cur_limits = client.get_budget_limits(budget_id, cur_start, cur_end)
            
            if not cur_limits:
                print(f"{budget_name:<25} | {'NO LIMIT':<10} | {'N/A':<10} | {'N/A':<10} | {'N/A':<10}")
                no_limits_count += 1
                continue

            for limit in cur_limits:
                amt = limit["attributes"]["amount"]
                start = limit["attributes"]["start"][:10]
                end = limit["attributes"]["end"][:10]
                period = detect_period_type(start, end)
                print(f"{budget_name:<25} | {amt:<10} | {period:<10} | {start:<10} | {end:<10}")

        print("=" * 75)
        print(f"Total Budgets: {len(budgets)} | Budgets without active limit: {no_limits_count}\n")
        return

    # --- NORMAL BACKFILL MODE ---
    print(f"Scanning current month window ({cur_start} .. {cur_end}) for active budget configs...\n")
    print(f"Found {len(budgets)} budget(s)\n")

    to_create = []   # (budget_id, budget_name, start, end, amount, currency_id)
    to_update = []   # (limit_id, budget_id, budget_name, start, end, old_amount, new_amount)
    no_current_limit = []

    for b in budgets:
        budget_id = b["id"]
        budget_name = b["attributes"]["name"]

        cur_limits = client.get_budget_limits(budget_id, cur_start, cur_end)
        if not cur_limits:
            no_current_limit.append(budget_name)
            continue

        source = cur_limits[0]
        amount = source["attributes"]["amount"]
        currency_id = source["attributes"].get("currency_id")

        period_type = detect_period_type(source["attributes"]["start"], source["attributes"]["end"])
        periods = generate_past_periods(today, args.years, period_type)

        range_start = periods[0][0]
        range_end = periods[-1][1]

        existing_range = client.get_budget_limits(budget_id, range_start, range_end)
        existing_by_start = {l["attributes"]["start"][:10]: l for l in existing_range}

        for start, end in periods:
            existing = existing_by_start.get(start)
            if existing is None:
                to_create.append((budget_id, budget_name, start, end, amount, currency_id))
            elif args.overwrite and existing["attributes"]["amount"] != amount:
                to_update.append((existing["id"], budget_id, budget_name, start, end,
                                   existing["attributes"]["amount"], amount))

    print(f"Budgets skipped (no active limit found in current month): {len(no_current_limit)}")
    for name in no_current_limit:
        print(f"    - {name}")

    print(f"\nBudget limits to CREATE: {len(to_create)}")
    print(f"Budget limits to UPDATE (--overwrite): {len(to_update)}")

    if not args.apply:
        print("\n[DRY RUN] No --apply flag given, nothing was sent to the API.")
        print("Sample of planned creates:")
        for row in to_create[:10]:
            _, name, start, end, amount, _ = row
            print(f"    {name}: {start} .. {end} -> {amount}")
        if len(to_create) > 10:
            print(f"    ... and {len(to_create) - 10} more")
        print("\nRe-run with --apply to apply changes.")
        return

    print("\nApplying...")
    created, updated, failed = 0, 0, 0

    for budget_id, name, start, end, amount, currency_id in to_create:
        try:
            client.create_budget_limit(budget_id, start, end, amount, currency_id=currency_id)
            created += 1
        except Exception as e:
            print(f"  [FAIL create] {name} {start}: {e}")
            failed += 1
        time.sleep(args.delay)

    for limit_id, budget_id, name, start, end, old_amount, new_amount in to_update:
        try:
            client.update_budget_limit(limit_id, new_amount)
            updated += 1
        except Exception as e:
            print(f"  [FAIL update] {name} {start}: {e}")
            failed += 1
        time.sleep(args.delay)

    print(f"\nDone. Created={created}, Updated={updated}, Failed={failed}")


if __name__ == "__main__":
    main()