"""Загрузчик кварталов dataset_СДЕЛКИ через SSH-tunnel на localhost:15432.

Эквивалент 02_load_all_quarters.sh, но через psycopg вместо docker — удобнее
запускать с Windows / любой машины с уже открытым SSH-tunnel.

Idempotent: пропускает кварталы, уже залитые в rosreestr_deals.

Usage:
    cd backend
    uv run python ../data/sql/02b_load_via_tunnel.py 2024Q1 2024Q2
    uv run python ../data/sql/02b_load_via_tunnel.py            # все из JOBS
"""

from __future__ import annotations

import os
import sys
import time
from pathlib import Path

import psycopg

PG_DSN = os.environ.get(
    "TUNNEL_DSN",
    "postgresql://gendesign:2J2SBPMKuS998fiwhtQqDhMI@localhost:15432/gendesign",
)

RAW = Path(__file__).resolve().parent.parent / "raw"

# (source_quarter, period_start, csv_filename, delimiter)
JOBS: list[tuple[str, str, str, str]] = [
    ("2024Q1", "2024-01-01", "dataset_СДЕЛКИ_r-r_01-92_y_2024_q_1.csv", ";"),
    ("2024Q2", "2024-04-01", "dataset_СДЕЛКИ_r-r_01-92_y_2024_q_2.csv", ";"),
    # Q3+ уже залиты, но оставляем как fallback / re-run target.
    ("2024Q3", "2024-07-01", "dataset_СДЕЛКИ_r_all_q_3.csv", ";"),
    ("2024Q4", "2024-10-01", "dataset_СДЕЛКИ_r-r_01-92_y_2024_q_4.csv", "~"),
    ("2025Q1", "2025-01-01", "dataset_СДЕЛКИ_r-r_01-92_y_2025_q_1.csv", "~"),
    ("2025Q2", "2025-04-01", "dataset_СДЕЛКИ_r-r_01-92_y_2025_q_2.csv", "~"),
    ("2025Q3", "2025-07-01", "dataset_СДЕЛКИ_r-r_01-92_y_2025_q_3.csv", "~"),
    ("2025Q4", "2025-10-01", "dataset_СДЕЛКИ_r-r_01-92_y_2025_q_4.csv", "~"),
    ("2026Q1", "2026-01-01", "dataset_СДЕЛКИ_r-r_01-92_y_2026_q_1.csv", "~"),
]


COPY_SQL = """
    COPY rosreestr_deals_staging
    FROM STDIN (FORMAT csv, HEADER true, DELIMITER %r, QUOTE '"', NULL '')
"""

INSERT_SQL = """
INSERT INTO rosreestr_deals (
    source_quarter, period_start_date,
    region_code, okato, district, city, quarter_cad_number, street,
    realestate_type_code, wall_material_code, purpose_code,
    year_build, floor, area,
    doc_type, deal_price, currency, deal_count
)
SELECT
    %s, period_start_date,
    region_code, NULLIF(okato, '-'), NULLIF(district, ''), NULLIF(city, ''),
    quarter_cad_number, NULLIF(street, ''),
    NULLIF(realestate_type_code, ''),
    NULLIF(wall_material_code, ''),
    NULLIF(purpose_code, ''),
    CASE WHEN year_build ~ '^[12][0-9]{3}$' THEN year_build::INT ELSE NULL END,
    NULLIF(floor, ''),
    CASE WHEN area >  0 AND area <  1e8 THEN area ELSE NULL END,
    NULLIF(doc_type, ''),
    CASE WHEN deal_price > 0 AND deal_price < 1e15 THEN deal_price ELSE NULL END,
    COALESCE(NULLIF(currency, ''), 'рубль'),
    COALESCE(deal_count, 1)
FROM rosreestr_deals_staging
WHERE region_code IS NOT NULL
  AND quarter_cad_number IS NOT NULL
  AND quarter_cad_number <> ''
  AND doc_type IS NOT NULL
  AND doc_type <> ''
  AND period_start_date IS NOT NULL
"""


def main() -> int:
    targets = sys.argv[1:] or [j[0] for j in JOBS]
    selected = [j for j in JOBS if j[0] in targets]
    if not selected:
        print(f"Нет matching кварталов в JOBS. Targets: {targets}")
        return 1

    conn = psycopg.connect(PG_DSN, connect_timeout=15)
    try:
        for sq, period, fname, sep in selected:
            with conn.cursor() as cur:
                cur.execute(
                    "SELECT COUNT(*) FROM rosreestr_deals WHERE source_quarter = %s",
                    (sq,),
                )
                already = cur.fetchone()[0]
            if already:
                print(f"=== {sq} SKIP — уже {already} строк (DELETE FROM rosreestr_deals WHERE source_quarter='{sq}' чтобы перезалить)")
                continue
            csv_path = RAW / fname
            if not csv_path.exists():
                print(f"=== {sq} SKIP — нет файла {csv_path}")
                continue
            print(f"=== {sq}  ({fname}, sep='{sep}', {csv_path.stat().st_size/1e6:.1f} MB) ===")
            t0 = time.time()
            with conn.cursor() as cur:
                cur.execute("TRUNCATE rosreestr_deals_staging")
                copy_sql = (
                    f"COPY rosreestr_deals_staging FROM STDIN "
                    f"(FORMAT csv, HEADER true, DELIMITER '{sep}', QUOTE '\"', NULL '')"
                )
                with cur.copy(copy_sql) as copy:
                    with open(csv_path, "rb") as f:
                        while True:
                            chunk = f.read(8 * 1024 * 1024)
                            if not chunk:
                                break
                            copy.write(chunk)
                t_copy = time.time() - t0
                cur.execute(INSERT_SQL, (sq,))
                inserted = cur.rowcount
            conn.commit()
            t_total = time.time() - t0
            print(f"  COPY {t_copy:.0f}s, INSERT {inserted} rows, total {t_total:.0f}s")

        # Final summary
        with conn.cursor() as cur:
            cur.execute(
                """
                SELECT source_quarter, COUNT(*) AS rows, SUM(deal_count) AS deals
                  FROM rosreestr_deals
                 GROUP BY source_quarter
                 ORDER BY source_quarter
                """
            )
            print("\n=== Итоги по rosreestr_deals ===")
            for sq, rows, deals in cur.fetchall():
                print(f"  {sq}  rows={rows:>9}  deals={deals:>10}")
    finally:
        conn.close()
    return 0


if __name__ == "__main__":
    sys.exit(main())