"""Parse DOM.RF stat_series XLSX files into long-form observations.

Reads:  data/raw/domrf_full/stat_series/xlsx/*.xlsx (4 standard-format files)
Writes: domrf_stat_series_indicators, domrf_stat_series_observations
        (created by 35_schema_domrf_stats.sql).

Standard format (files 01_01..01_04):
    Sheet '*_00' = index (skipped — derive from individual sheets)
    Sheets '*_NN' = one indicator × all territories × all months
        row 1: section title (long sentence)
        row 2: indicator name ("Количество МКД...")
        row 4: header — col 0 = "Регион", col 1..N = month dates
        row 5+: territory + numeric values

Files with bespoke format (01_05_*, 01_06_*) are skipped — left in
domrf_xlsx_files BYTEA archive for ad-hoc analysis.
"""
import os, re, sys, subprocess, datetime
import openpyxl

HERE = os.path.dirname(os.path.abspath(__file__))
ROOT = os.path.join(HERE, '..', 'raw', 'domrf_full', 'stat_series', 'xlsx')
PG_HOST = os.environ.get('PG_HOST', 'host.docker.internal')
PG_PORT = os.environ.get('PG_PORT', '15432')
PG_USER = os.environ.get('PG_USER', 'gendesign')
PG_DB   = os.environ.get('PG_DB',   'gendesign')
PG_PASS = os.environ.get('PGPASSWORD', '2J2SBPMKuS998fiwhtQqDhMI')
SNAP    = os.environ.get('SNAPSHOT_DATE', datetime.date.today().isoformat())

# Sheet name pattern: NN_NN_NN (e.g. '01_01_01', '01_02_05'). The '*_00' index
# is skipped (last group != '00').
SHEET_RE = re.compile(r'^(\d\d)_(\d\d)_(\d\d)$')

# Standard-format files only. Others have bespoke layouts.
STANDARD_FILES = {
    '01_01_stockvariablesexsales.xlsx',
    '01_01_stockvariableskrt.xlsx',
    '01_02_flowvariables.xlsx',
    '01_03_sales.xlsx',
    '01_04_flowvariables.xlsx',
    '01_04_stockvariables.xlsx',
    '01_06_stockvariables_izs.xlsx',
}


def psql(sql, capture=True):
    cmd = ['docker', 'run', '--rm', '-i', '-e', f'PGPASSWORD={PG_PASS}',
           'postgres:16-alpine', 'psql',
           '-h', PG_HOST, '-p', PG_PORT, '-U', PG_USER, '-d', PG_DB,
           '-v', 'ON_ERROR_STOP=1', '--quiet']
    res = subprocess.run(cmd, input=sql, capture_output=capture, text=True, encoding='utf-8')
    if res.returncode != 0:
        print('psql FAILED:', (res.stderr or '')[-2000:], file=sys.stderr)
        raise SystemExit(res.returncode)
    return res.stdout


def esc(s):
    if s is None:
        return 'NULL'
    return "'" + str(s).replace("'", "''") + "'"


def num(v):
    if v is None or v == '':
        return 'NULL'
    try:
        f = float(v)
        # Strip trailing .0 noise but keep precision
        if f != f:  # NaN
            return 'NULL'
        return repr(f)
    except (ValueError, TypeError):
        return 'NULL'


def to_date(v):
    if v is None:
        return None
    if isinstance(v, datetime.datetime):
        return v.date().isoformat()
    if isinstance(v, datetime.date):
        return v.isoformat()
    s = str(v)
    # Sometimes the cell is a string like '2024-03-01'
    m = re.match(r'(\d{4})-(\d{2})-(\d{2})', s)
    if m:
        return f'{m.group(1)}-{m.group(2)}-{m.group(3)}'
    return None


def parse_sheet(ws, file_name, sheet_code):
    """Return (indicator_name, [(territory, obs_date, value)...]). Stream rows."""
    rows = list(ws.iter_rows(values_only=True))
    if len(rows) < 5:
        return None, []
    # Row 1 (index 0) — section title; row 2 — indicator; rows 4..end — data
    indicator_name = (rows[1][0] or '').strip() if rows[1] and rows[1][0] else ''
    header = rows[3]
    if not header or len(header) < 2:
        return indicator_name, []
    # Skip first 'Регион' column; remaining are month dates
    date_cols = []
    for col_idx, cell in enumerate(header[1:], start=1):
        d = to_date(cell)
        if d:
            date_cols.append((col_idx, d))
    obs = []
    for r in rows[4:]:
        if not r or r[0] is None:
            continue
        territory = str(r[0]).strip()
        if not territory or territory.lower() == 'перейти в начало':
            continue
        for col_idx, obs_date in date_cols:
            if col_idx >= len(r):
                continue
            val = r[col_idx]
            if val is None or val == '':
                continue
            obs.append((territory, obs_date, val))
    return indicator_name, obs


def main():
    if not os.path.isdir(ROOT):
        print('xlsx dir missing:', ROOT)
        return
    schema = open(os.path.join(HERE, '35_schema_domrf_stats.sql'),
                  encoding='utf-8').read()
    psql(schema)
    print('schema applied')

    indicator_rows = []  # (file_name, sheet_code, indicator_name, n_periods, n_territories)
    obs_rows = []        # (file_name, sheet_code, territory, obs_date, value)
    parsed_sheets = 0
    skipped_files = []

    for fname in sorted(os.listdir(ROOT)):
        if not fname.endswith('.xlsx'):
            continue
        if fname not in STANDARD_FILES:
            skipped_files.append(fname)
            continue
        path = os.path.join(ROOT, fname)
        print(f'  parsing {fname}')
        wb = openpyxl.load_workbook(path, read_only=True, data_only=True)
        for sn in wb.sheetnames:
            m = SHEET_RE.match(sn)
            if not m or m.group(3) == '00':
                continue
            ws = wb[sn]
            indicator_name, obs = parse_sheet(ws, fname, sn)
            if not obs:
                continue
            territories = {t for t, _, _ in obs}
            periods = {d for _, d, _ in obs}
            indicator_rows.append((fname, sn, indicator_name, len(periods), len(territories)))
            for t, d, v in obs:
                obs_rows.append((fname, sn, t, d, v))
            parsed_sheets += 1
        wb.close()

    print(f'parsed {parsed_sheets} sheets; skipped (bespoke format): {skipped_files}')
    print(f'  indicators: {len(indicator_rows)}')
    print(f'  observations: {len(obs_rows)}')

    # ── INDICATORS upsert ──
    if indicator_rows:
        sql = ("INSERT INTO domrf_stat_series_indicators "
               "(file_name, sheet_code, indicator_name, n_periods, n_territories) VALUES "
               + ','.join(
                    f"({esc(f)},{esc(c)},{esc(n)},{p},{t})"
                    for f, c, n, p, t in indicator_rows)
               + " ON CONFLICT (file_name, sheet_code) DO UPDATE SET "
                 "indicator_name=EXCLUDED.indicator_name, "
                 "n_periods=EXCLUDED.n_periods, "
                 "n_territories=EXCLUDED.n_territories;")
        psql(sql)
        print(f'  indicators upserted: {len(indicator_rows)}')

    # ── OBSERVATIONS upsert (chunked — ~810K rows). Dedupe within batch
    # because the same (file, sheet, territory, obs_date) can appear twice
    # in the source XLSX (rarely — duplicate header dates) and ON CONFLICT
    # rejects in-batch duplicates.
    seen = set()
    deduped = []
    for f, c, t, d, v in obs_rows:
        key = (f, c, t, d)
        if key in seen:
            continue
        seen.add(key)
        deduped.append((f, c, t, d, v))
    if len(deduped) != len(obs_rows):
        print(f'  deduped {len(obs_rows) - len(deduped)} duplicate observations')
    obs_rows = deduped

    BATCH = 2000
    total = 0
    for i in range(0, len(obs_rows), BATCH):
        chunk = obs_rows[i:i + BATCH]
        values = ','.join(
            f"({esc(SNAP)},{esc(f)},{esc(c)},{esc(t)},{esc(d)},{num(v)})"
            for f, c, t, d, v in chunk
        )
        sql = ("INSERT INTO domrf_stat_series_observations "
               "(snapshot_date, file_name, sheet_code, territory, obs_date, value) VALUES "
               + values
               + " ON CONFLICT (snapshot_date, file_name, sheet_code, territory, obs_date) "
                 "DO UPDATE SET value=EXCLUDED.value;")
        psql(sql)
        total += len(chunk)
        if total % 20000 == 0 or total == len(obs_rows):
            print(f'    obs upserted: {total}/{len(obs_rows)}')


if __name__ == '__main__':
    main()
