#!/usr/bin/env python3
"""Reproduce the Netflix "What We Watched" hours analysis from scratch.

Downloads all seven engagement reports (H1 2023 - H1 2026), joins show seasons
into one row per show, sums hours across all reports, and writes netflix_hours.csv
with: title, type, hours_viewed, seasons (the " - "-separated season titles that
were joined), n_reports, first_release, available_globally, runtime_min (total
content runtime; blank for titles only in the runtime-less H1 2023 report), and
per-half hours columns h_2023h1 ... h_2026h1.

Requires: openpyxl  (pip install openpyxl)
"""
import csv
import os
import re
import urllib.request

import openpyxl

REPORTS = {
    'What_We_Watched_2023Jan-Jun.xlsx': 'https://assets.ctfassets.net/4cd45et68cgf/1HyknFM84ISQpeua6TjM7A/97a0a393098937a8f29c9d29c48dbfa8/What_We_Watched_A_Netflix_Engagement_Report_2023Jan-Jun.xlsx',
    'What_We_Watched_2023Jul-Dec.xlsx': 'https://assets.ctfassets.net/4cd45et68cgf/inuAnzotdsAEgbInGLzH5/1be323ba419b2af3a96bffa29acc31a3/What_We_Watched_A_Netflix_Engagement_Report_2023Jul-Dec.xlsx',
    'What_We_Watched_2024Jan-Jun.xlsx': 'https://assets.ctfassets.net/4cd45et68cgf/2PoZlfdc46dH2gQvI8eUzI/9db5840720c47acfcf7b89ffe2402860/What_We_Watched_A_Netflix_Engagement_Report_2024Jan-Jun.xlsx',
    'What_We_Watched_2024Jul-Dec.xlsx': 'https://assets.ctfassets.net/4cd45et68cgf/6XSmoEjBjVMPRtYybT9d1E/8c0b0b2645b8712d5597b0bdbe0d64e2/What_We_Watched_A_Netflix_Engagement_Report_2024Jul-Dec.xlsx',
    'What_We_Watched_2025Jan-Jun.xlsx': 'https://assets.ctfassets.net/4cd45et68cgf/mplcXj5ulHDfbCPCr0f0I/5dbb6ec09f03df89706476e380e9b8bd/What_We_Watched_A_Netflix_Engagement_Report_2025Jan-Jun.xlsx',
    'What_We_Watched_2025Jul-Dec.xlsx': 'https://assets.ctfassets.net/4cd45et68cgf/2vdDPGLKA0XX2cjF2APJFn/7f1c367b39ed73a6a588751d3c5d0252/What_We_Watched_A_Netflix_Engagement_Report_2025Jul-Dec__6_.xlsx',
    'What_We_Watched_2026Jan-Jun.xlsx': 'https://assets.ctfassets.net/4cd45et68cgf/40WGcHJa9vRua31kU6Gbz5/43b15c7fbd6924392fc80e885839e39f/Netflix-s_What_We_Watched_Report_2026Jan-Jun__1_.xlsx',
}

ROMAN = re.compile(r'\s+(II|III|IV|V|VI|VII|VIII|IX|X)$')
SUFFIX = re.compile(
    r'\s*[:\-–]\s*(Limited Series|Season|Series|Part|Volume|Vol\.?|Chapter|Book|'
    r'Collection|Edition|Saison|Temporada|Staffel|Cour)\b.*$', re.I)
OTHER = re.compile(r'^Other (Shows|Movies|TV|Films?)$', re.I)


def show_key(title):
    """Collapse a season title to its base show name.
    'Squid Game: Season 3 // ...' -> 'Squid Game'; 'Stranger Things 5' -> 'Stranger Things'."""
    s = title.split('//')[0].strip()
    s = SUFFIX.sub('', s)
    s = ROMAN.sub('', s)
    s = re.sub(r'\s+(\d{1,2})$', lambda m: '' if int(m.group(1)) <= 25 else m.group(0), s)
    return s.strip()


def movie_key(title):
    return title.split('//')[0].strip()


def season_like(title):
    s = title.split('//')[0].strip()
    return bool(SUFFIX.search(s) or ROMAN.search(s))


HALVES = ['2023Jan-Jun', '2023Jul-Dec', '2024Jan-Jun', '2024Jul-Dec',
          '2025Jan-Jun', '2025Jul-Dec', '2026Jan-Jun']


def runtime_min(rt):
    if isinstance(rt, str) and ':' in rt:
        h, m = rt.split(':')[:2]
        try:
            return int(h) * 60 + int(m)
        except ValueError:
            pass
    return None


def rows_of(path, sheet):
    wb = openpyxl.load_workbook(path, read_only=True)
    ws = wb[sheet]
    out, started = [], False
    for row in ws.iter_rows(values_only=True):
        if not started:
            if row[1] == 'Title':
                started = True
            continue
        title, avail, released, hours = row[1], row[2], row[3], row[4]
        rt = runtime_min(row[5]) if len(row) > 5 else None
        if title and isinstance(hours, (int, float)) and not OTHER.match(title.strip()):
            out.append((title.strip(), avail == 'Yes',
                        str(released)[:10] if released else None, hours, rt))
    wb.close()
    return out


def main():
    outdir = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'reports')
    os.makedirs(outdir, exist_ok=True)
    for name, url in REPORTS.items():
        path = os.path.join(outdir, name)
        if not os.path.exists(path):
            print('downloading', name)
            urllib.request.urlretrieve(url, path)

    shows, movies = {}, {}

    def add(store, key, is_global, released, hours, member, half_idx, rt):
        e = store.setdefault(key, {'h': 0, 'g': True, 'd': None, 'm': [], 'p': 0,
                                   'hh': [0] * 7, 'rt': {}})
        e['h'] += hours
        e['hh'][half_idx] += hours
        e['p'] += 1
        if not is_global:
            e['g'] = False
        if released and (e['d'] is None or released < e['d']):
            e['d'] = released
        mkey = member or key
        if member and member not in e['m']:
            e['m'].append(member)
        if rt is not None:
            e['rt'][mkey] = max(e['rt'].get(mkey, 0), rt)

    # H2 2023 onward: separate show/movie sheets
    for name in sorted(REPORTS):
        if '2023Jan-Jun' in name:
            continue
        hi = next(i for i, hv in enumerate(HALVES) if hv in name)
        path = os.path.join(outdir, name)
        wb = openpyxl.load_workbook(path, read_only=True)
        sheets = wb.sheetnames
        wb.close()
        for t, g, d, h, rt in rows_of(path, 'Shows' if 'Shows' in sheets else 'TV'):
            add(shows, show_key(t), g, d, h, t.split('//')[0].strip(), hi, rt)
        for t, g, d, h, rt in rows_of(path, 'Movies' if 'Movies' in sheets else 'Film'):
            add(movies, movie_key(t), g, d, h, None, hi, rt)

    # H1 2023: one mixed sheet; classify each row (known movies beat the
    # trailing-digit rule so movie sequels like "Bad Neighbours 2" stay movies)
    path = os.path.join(outdir, 'What_We_Watched_2023Jan-Jun.xlsx')
    for t, g, d, h, rt in rows_of(path, 'Engagement'):
        k, base = show_key(t), movie_key(t)
        if base in movies:
            add(movies, base, g, d, h, None, 0, rt)
        elif season_like(t) or k in shows:
            add(shows, k, g, d, h, base, 0, rt)
        else:
            add(movies, base, g, d, h, None, 0, rt)

    out = os.path.join(os.path.dirname(outdir), 'netflix_hours.csv')
    rows = ([('show', k, e) for k, e in shows.items()] +
            [('movie', k, e) for k, e in movies.items()])
    rows.sort(key=lambda r: -r[2]['h'])
    half_cols = ['h_' + hv.replace('Jan-Jun', 'h1').replace('Jul-Dec', 'h2')[:7].lower()
                 for hv in HALVES]
    with open(out, 'w', newline='') as f:
        w = csv.writer(f)
        w.writerow(['title', 'type', 'hours_viewed', 'seasons', 'n_reports',
                    'first_release', 'available_globally', 'runtime_min'] + half_cols)
        for typ, k, e in rows:
            total_rt = sum(e['rt'].values()) if e['rt'] else ''
            w.writerow([k, typ, int(e['h']), ' - '.join(sorted(e['m'])),
                        e['p'], e['d'] or '', 'yes' if e['g'] else 'no', total_rt]
                       + [int(x) for x in e['hh']])

    total = sum(e['h'] for _, _, e in rows)
    print(f'wrote {out}: {len(rows):,} titles '
          f'({len(shows):,} shows, {len(movies):,} movies), '
          f'{total/1e9:.1f}B hours total')


if __name__ == '__main__':
    main()
