#!/usr/bin/env python3
"""
FPRates - EDF half-hourly usage to CSV

Runs on your own computer. Your EDF login is used only to talk to EDF's own
API (api.edfgb-kraken.energy). Nothing is sent to FPRates or anywhere else.
Needs Python 3.8 or later and nothing else.
"""
import csv, getpass, json, os, sys, urllib.error, urllib.parse, urllib.request
from datetime import datetime, timedelta

BASE = "https://api.edfgb-kraken.energy/v1"
GQL = BASE + "/graphql/"


def post_json(url, payload, token=None):
    headers = {"Content-Type": "application/json"}
    if token:
        headers["Authorization"] = token
    req = urllib.request.Request(url, data=json.dumps(payload).encode(), headers=headers)
    with urllib.request.urlopen(req, timeout=30) as r:
        return json.load(r)


def get_json(url, auth):
    req = urllib.request.Request(url, headers={"Authorization": auth})
    with urllib.request.urlopen(req, timeout=60) as r:
        return json.load(r)


def login(email, password):
    q = ("mutation($e: String!, $p: String!) "
         "{ obtainKrakenToken(input: {email: $e, password: $p}) { token } }")
    res = post_json(GQL, {"query": q, "variables": {"e": email, "p": password}})
    try:
        return res["data"]["obtainKrakenToken"]["token"]
    except (KeyError, TypeError):
        sys.exit("Login failed - check your EDF email and password.")


def find_meters(token):
    q = """query { viewer { accounts { number ... on AccountType { properties {
             electricityMeterPoints { mpan meters { serialNumber } } } } } } }"""
    res = post_json(GQL, {"query": q}, token)
    if res.get("errors"):
        sys.exit(f"Couldn't read your meters: {res['errors']}")
    meters = []
    for acc in res["data"]["viewer"]["accounts"]:
        for prop in acc.get("properties") or []:
            for mp in prop.get("electricityMeterPoints") or []:
                for m in mp.get("meters") or []:
                    if m.get("serialNumber"):
                        meters.append((mp["mpan"], m["serialNumber"]))
    return meters


def fetch(token, mpan, serial, start, end):
    params = urllib.parse.urlencode({
        "period_from": start.isoformat(), "period_to": end.isoformat(),
        "page_size": 25000, "order_by": "period"})
    url = f"{BASE}/electricity-meter-points/{mpan}/meters/{serial}/consumption/?{params}"
    for auth in (token, f"Bearer {token}", f"JWT {token}"):
        try:
            rows, next_url = [], url
            while next_url:
                data = get_json(next_url, auth)
                rows.extend(data["results"])
                next_url = data.get("next")
            return rows
        except urllib.error.HTTPError as e:
            if e.code in (401, 403):
                continue          # try the next way of sending the login token
            if e.code == 404:
                return []         # old or replaced meter
            raise
    sys.exit("EDF refused the usage request (not authorised).")


def ask_date(prompt):
    while True:
        s = input(prompt).strip()
        try:
            return datetime.strptime(s, "%Y-%m-%d")
        except ValueError:
            print("  Please use YYYY-MM-DD, for example 2026-09-01")


def main():
    print(__doc__)
    email = os.environ.get("EDF_EMAIL") or input("EDF email: ").strip()
    password = os.environ.get("EDF_PASS") or getpass.getpass("EDF password (not shown): ")
    d_from = ask_date("From date (YYYY-MM-DD): ")
    d_to = ask_date("To date, inclusive (YYYY-MM-DD): ")
    if d_to < d_from:
        sys.exit("The 'to' date is before the 'from' date.")
    start = d_from.astimezone()                       # local midnight
    end = (d_to + timedelta(days=1)).astimezone()     # midnight after the 'to' date

    token = login(email, password)
    meters = find_meters(token)
    if not meters:
        sys.exit("No electricity meters found on this account.")

    for mpan, serial in meters:
        rows = fetch(token, mpan, serial, start, end)
        by_start = {}
        for r in rows:
            s = datetime.fromisoformat(r["interval_start"].replace("Z", "+00:00"))
            if start <= s < end:
                by_start[s] = r
        rows = [by_start[k] for k in sorted(by_start)]
        label = f"MPAN ...{mpan[-4:]}, meter ...{serial[-4:]}"
        if not rows:
            print(f"{label}: no readings in that period - skipped")
            continue
        fname = f"edf_usage_{mpan[-4:]}_{serial[-4:]}_{d_from:%Y%m%d}_{d_to:%Y%m%d}.csv"
        total = 0.0
        with open(fname, "w", newline="") as f:
            w = csv.writer(f)
            w.writerow(["start", "end", "kWh"])
            for r in rows:
                s = datetime.fromisoformat(r["interval_start"].replace("Z", "+00:00")).astimezone()
                e = datetime.fromisoformat(r["interval_end"].replace("Z", "+00:00")).astimezone()
                w.writerow([s.strftime("%Y-%m-%d %H:%M"), e.strftime("%Y-%m-%d %H:%M"), r["consumption"]])
                total += r["consumption"]
        expected = int((end - start).total_seconds() // 1800)
        print(f"{label}: saved {len(rows)} half-hours ({total:.2f} kWh) to {fname}")
        if len(rows) < expected:
            print(f"  Note: {expected - len(rows)} half-hours missing - EDF may not have received them yet.")


if __name__ == "__main__":
    main()