#!/usr/bin/env python3
"""
Import SEAP data from data.gov.ro XLSX files into PostgreSQL.

Strategy:
1. Import "Anunturi de initiere" → builds CUI → (localitate, judet) mapping
2. Import "Achizitii directe" → main volume, resolves location via CUI
3. Import "Contracte" → public tenders with winner info
4. Run locality matching → SIRUTA codes
"""

import os
import sys
import urllib.request
import tempfile
from pathlib import Path
from datetime import datetime

import openpyxl
import psycopg2
from psycopg2.extras import execute_values

DB_URL = os.environ.get('DATABASE_URL',
    'postgresql://architools_user:stictMyFon34!_gonY@10.10.10.166:5432/architools_db')

DATA_DIR = Path(__file__).parent / 'data'
DATA_DIR.mkdir(exist_ok=True)

# ── Download URLs for 2025 ──

URLS_2025 = {
    'anunturi_initiere': [
        ('T1', 'https://data.gov.ro/dataset/e0cf7ffc-1fa0-4ffb-a82f-c83981d81f21/resource/6bcc924b-fdb7-482c-91dc-d57751c58b5c/download/datagov_raport_anunturi-de-initiere-publicate_t1_2025.xlsx'),
        ('T2', 'https://data.gov.ro/dataset/e0cf7ffc-1fa0-4ffb-a82f-c83981d81f21/resource/49940d22-9a5a-41ff-92a6-da7d3ef45800/download/anunturi-de-initiere-publicate-t2-2025.xlsx'),
        ('T3', 'https://data.gov.ro/dataset/e0cf7ffc-1fa0-4ffb-a82f-c83981d81f21/resource/64e18773-e97c-4478-9b3d-3654d58b020f/download/datagov-anunturi-de-initiere-publicate-tiii-2025.xlsx'),
        ('T4', 'https://data.gov.ro/dataset/e0cf7ffc-1fa0-4ffb-a82f-c83981d81f21/resource/92e3dcec-41ff-4771-895b-6ec880a5ad6a/download/anunuri-de-iniiere-publicate-t_iv_2025.xlsx'),
    ],
    'achizitii_directe': [
        ('T1', 'https://data.gov.ro/dataset/e0cf7ffc-1fa0-4ffb-a82f-c83981d81f21/resource/4ea2f0d0-ad5d-440f-af9d-7101bc9e4969/download/datagov_raport-achizitii_directe_t1_2025.xlsx'),
        ('T2', 'https://data.gov.ro/dataset/e0cf7ffc-1fa0-4ffb-a82f-c83981d81f21/resource/8e6fa9e7-62e9-4ec2-bef5-495f3d09eef3/download/achizitii-directe-t2-2025.xlsx'),
        ('T3', 'https://data.gov.ro/dataset/e0cf7ffc-1fa0-4ffb-a82f-c83981d81f21/resource/21cd9887-26ca-418d-ade4-be5d369b4246/download/datagov-achizitii-directe-tiii-2025.xlsx'),
        ('T4', 'https://data.gov.ro/dataset/e0cf7ffc-1fa0-4ffb-a82f-c83981d81f21/resource/370af861-b17f-4807-b69c-9cf3b67df997/download/achiziii-directe-t_iv_2025.xlsx'),
    ],
    'contracte': [
        ('T1', 'https://data.gov.ro/dataset/e0cf7ffc-1fa0-4ffb-a82f-c83981d81f21/resource/7344eeaf-c478-4f87-9669-c1bac3e521a8/download/datagov_raport_contracte_t1_2025.xlsx'),
        ('T2', 'https://data.gov.ro/dataset/e0cf7ffc-1fa0-4ffb-a82f-c83981d81f21/resource/91947695-b315-4d72-b292-84bd57a9c72b/download/contracte-t2-2025.xlsx'),
        ('T3', 'https://data.gov.ro/dataset/e0cf7ffc-1fa0-4ffb-a82f-c83981d81f21/resource/e2ef5a81-59ec-4789-9baa-fd175b217893/download/datagov-contracte-tiii-2025.xlsx'),
        ('T4', 'https://data.gov.ro/dataset/e0cf7ffc-1fa0-4ffb-a82f-c83981d81f21/resource/a1936e88-7fc5-4ffc-af65-6f946e98a005/download/contracte-t_iv_2025.xlsx'),
    ],
}

URLS_2026 = {
    'anunturi_initiere': [
        ('T1', 'https://data.gov.ro/dataset/e8d22de4-f8ce-4b42-8561-98d2481ddef9/resource/5720192a-6c1a-4f40-bccc-9c12bc6a2a8f/download/raport-anunturi-de-initiere-publicate-ti-2026.xlsx'),
    ],
    'achizitii_directe': [
        ('T1', 'https://data.gov.ro/dataset/e8d22de4-f8ce-4b42-8561-98d2481ddef9/resource/5bcff70e-7541-4e7f-86e2-f21b54807e26/download/raport-achizitii-directe-ti-2026.xlsx'),
    ],
    'contracte': [
        ('T1', 'https://data.gov.ro/dataset/e8d22de4-f8ce-4b42-8561-98d2481ddef9/resource/f78b2b07-48aa-442e-b7e3-4b39f45a0b5b/download/raport-contracte-ti-2026.xlsx'),
    ],
}


def download(url, label):
    """Download file if not cached."""
    fname = url.split('/')[-1]
    fpath = DATA_DIR / fname
    if fpath.exists() and fpath.stat().st_size > 1000:
        print(f"  [cached] {label}: {fname} ({fpath.stat().st_size // 1024}KB)")
        return fpath
    print(f"  [download] {label}: {fname}...")
    urllib.request.urlretrieve(url, fpath)
    print(f"  [done] {fpath.stat().st_size // (1024*1024)}MB")
    return fpath


def find_columns(headers, *names):
    """Find column index by trying multiple possible names."""
    header_map = {}
    for i, h in enumerate(headers):
        if h:
            header_map[str(h).strip().upper()] = i
    for name in names:
        if name.upper() in header_map:
            return header_map[name.upper()]
    return None


def read_xlsx_rows(fpath, max_rows=None):
    """Read XLSX file in read-only mode, yield (headers, rows)."""
    wb = openpyxl.load_workbook(fpath, read_only=True, data_only=True)
    ws = wb.active
    rows = ws.iter_rows(values_only=True)
    headers = [str(h).strip() if h else '' for h in next(rows)]
    count = 0
    for row in rows:
        if max_rows and count >= max_rows:
            break
        yield headers, row
        count += 1
    wb.close()


def get_conn():
    return psycopg2.connect(DB_URL)


# ── Step 1: Import anunturi initiere → CUI location mapping ──

def import_anunturi_initiere(conn, urls):
    """Extract CUI → (localitate, judet) from announcement files."""
    cur = conn.cursor()

    # Create temp mapping table
    cur.execute("""
        CREATE TABLE IF NOT EXISTS seap.cui_location (
            cui TEXT PRIMARY KEY,
            name TEXT,
            city TEXT,
            county TEXT,
            updated_at TIMESTAMPTZ DEFAULT now()
        )
    """)
    conn.commit()

    total = 0
    for label, url in urls:
        fpath = download(url, f"Anunturi initiere {label}")
        batch = []

        for headers, row in read_xlsx_rows(fpath):
            cui_idx = find_columns(headers, 'CUI', 'CUI_AC', 'Cui')
            name_idx = find_columns(headers, 'Autoritate contractanta', 'DENUMIRE_AC',
                                    'Autoritate Contractanta', 'autoritate contractanta')
            city_idx = find_columns(headers, 'Localitate', 'LOCALITATE', 'localitate')
            county_idx = find_columns(headers, 'Judet', 'JUDET', 'judet', 'Județ')

            if cui_idx is None or city_idx is None:
                print(f"  [skip] Missing columns in {label}. Headers: {headers[:10]}")
                break

            cui = str(row[cui_idx]).strip() if row[cui_idx] else None
            name = str(row[name_idx]).strip() if name_idx and row[name_idx] else None
            city = str(row[city_idx]).strip() if row[city_idx] else None
            county = str(row[county_idx]).strip() if county_idx and row[county_idx] else None

            if cui and city:
                # Clean CUI
                cui = cui.replace('RO', '').replace('ro', '').strip()
                batch.append((cui, name, city, county))

            if len(batch) >= 5000:
                _insert_cui_batch(cur, batch)
                total += len(batch)
                batch = []

        if batch:
            _insert_cui_batch(cur, batch)
            total += len(batch)

        conn.commit()
        print(f"  [imported] {label}: {total} CUI mappings total")

    return total


def _insert_cui_batch(cur, batch):
    execute_values(cur, """
        INSERT INTO seap.cui_location (cui, name, city, county)
        VALUES %s
        ON CONFLICT (cui) DO UPDATE SET
            name = COALESCE(EXCLUDED.name, seap.cui_location.name),
            city = COALESCE(EXCLUDED.city, seap.cui_location.city),
            county = COALESCE(EXCLUDED.county, seap.cui_location.county),
            updated_at = now()
    """, batch)


# ── Step 2: Import achizitii directe ──

def import_achizitii_directe(conn, urls):
    """Import direct acquisitions from XLSX."""
    cur = conn.cursor()
    total = 0
    skipped = 0

    for label, url in urls:
        fpath = download(url, f"Achizitii directe {label}")
        batch = []
        file_rows = 0

        for headers, row in read_xlsx_rows(fpath):
            # Find column indices
            nr_idx = find_columns(headers, 'NUMAR_ACHIZITIE_DIRECTA', 'Numar achizitie directa')
            date_pub_idx = find_columns(headers, 'DATA_PUBLICARE_ACHIZITIE', 'Data publicare achizitie', 'Data publicare')
            date_attr_idx = find_columns(headers, 'DATA_ATRIBUIRE_ACHIZITIE', 'Data atribuire achizitie', 'Data finalizare')
            state_idx = find_columns(headers, 'STARE_ACHIZITIE', 'Stare achizitie')
            auth_name_idx = find_columns(headers, 'DENUMIRE_AC', 'Denumire AC', 'Autoritate contractanta')
            auth_cui_idx = find_columns(headers, 'CUI_AC', 'CUI AC', 'Cui AC',
                                        'CUI autoritate contractanta', 'CUI AUTORITATE CONTRACTANTA')
            name_idx = find_columns(headers, 'DENUMIRE_ACHIZITIE', 'Denumire achizitie')
            cpv_code_idx = find_columns(headers, 'COD_CPV', 'Cod CPV')
            cpv_name_idx = find_columns(headers, 'DENUMIRE_CPV', 'Denumire CPV')
            est_val_idx = find_columns(headers, 'VALOARE_ESTIMATA_RON', 'Valoare estimata (RON)')
            attr_val_idx = find_columns(headers, 'VALOARE_ATRIBUITA_RON', 'Valoare atribuita (RON)',
                                        'Valoare achizitie (RON)', 'VALOARE_ACHIZITIE_RON')
            supplier_idx = find_columns(headers, 'OFERTANT', 'Ofertant castigator')
            supplier_cui_idx = find_columns(headers, 'CUI_OFERTANT', 'CUI ofertant', 'Cui ofertant',
                                            'CUI ofertant castigator', 'CUI OFERTANT CASTIGATOR')

            if nr_idx is None:
                print(f"  [skip] Can't find DA number column. Headers: {headers[:15]}")
                break

            da_nr = str(row[nr_idx]).strip() if row[nr_idx] else None
            if not da_nr:
                continue

            def safe_float(idx):
                if idx is None: return None
                v = row[idx]
                if v is None: return None
                try: return float(str(v).replace(',', '.').replace(' ', ''))
                except: return None

            def safe_str(idx):
                if idx is None: return None
                return str(row[idx]).strip() if row[idx] else None

            def safe_date(idx):
                if idx is None: return None
                v = row[idx]
                if isinstance(v, datetime): return v
                if v is None: return None
                try: return datetime.fromisoformat(str(v).replace('/', '-'))
                except: return None

            auth_cui = safe_str(auth_cui_idx)
            if auth_cui:
                auth_cui = auth_cui.replace('RO', '').replace('ro', '').strip()

            sup_cui = safe_str(supplier_cui_idx)
            if sup_cui:
                sup_cui = sup_cui.replace('RO', '').replace('ro', '').strip()

            batch.append((
                da_nr,                          # unique_code
                safe_str(name_idx),             # name
                safe_str(cpv_code_idx),         # cpv_code
                safe_str(cpv_name_idx),         # cpv_name
                safe_date(date_pub_idx),        # publication_date
                safe_date(date_attr_idx),       # finalization_date
                safe_float(est_val_idx),        # estimated_value
                safe_float(attr_val_idx),       # closing_value
                safe_str(state_idx),            # state_text
                safe_str(auth_name_idx),        # authority_name (temporary)
                auth_cui,                       # authority_cui
                safe_str(supplier_idx),         # supplier_name (temporary)
                sup_cui,                        # supplier_cui
            ))
            file_rows += 1

            if len(batch) >= 5000:
                inserted = _insert_da_batch(cur, batch)
                total += inserted
                skipped += len(batch) - inserted
                batch = []
                print(f"  [{label}] {file_rows} rows processed, {total} inserted, {skipped} skipped...")

        if batch:
            inserted = _insert_da_batch(cur, batch)
            total += inserted
            skipped += len(batch) - inserted

        conn.commit()
        print(f"  [done] {label}: {file_rows} rows, {total} total inserted")

    return total


def _insert_da_batch(cur, batch):
    """Insert DA batch using bulk insert."""
    if not batch:
        return 0

    values = []
    for row in batch:
        (unique_code, name, cpv_code, cpv_name, pub_date, fin_date,
         est_val, close_val, state_text, auth_name, auth_cui,
         sup_name, sup_cui) = row
        values.append((unique_code, name, cpv_code, cpv_name, pub_date, fin_date,
                       est_val, close_val, state_text, auth_cui, sup_cui))

    try:
        execute_values(cur, """
            INSERT INTO seap.direct_acquisitions
                (id, unique_code, name, cpv_code, cpv_name,
                 publication_date, finalization_date,
                 estimated_value, closing_value, state_text,
                 authority_cui, supplier_cui)
            SELECT nextval('seap.da_import_seq'),
                   d.unique_code, d.name, d.cpv_code, d.cpv_name,
                   d.pub_date::timestamptz, d.fin_date::timestamptz,
                   d.est_val::numeric, d.close_val::numeric, d.state_text,
                   d.auth_cui, d.sup_cui
            FROM (VALUES %s) AS d(
                unique_code, name, cpv_code, cpv_name,
                pub_date, fin_date, est_val, close_val, state_text,
                auth_cui, sup_cui)
            WHERE NOT EXISTS (
                SELECT 1 FROM seap.direct_acquisitions da WHERE da.unique_code = d.unique_code
            )
        """, values, template="(%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)")
        inserted = cur.rowcount
        cur.connection.commit()
        return inserted
    except Exception as e:
        cur.connection.rollback()
        print(f"  [error] DA batch: {e}")
        return 0


# ── Step 3: Import contracte ──

def import_contracte(conn, urls):
    """Import contracts (public tenders) from XLSX."""
    cur = conn.cursor()
    total = 0

    for label, url in urls:
        fpath = download(url, f"Contracte {label}")
        batch = []
        file_rows = 0

        for headers, row in read_xlsx_rows(fpath):
            auth_idx = find_columns(headers, 'Autoritate contractanta', 'AUTORITATE_CONTRACTANTA')
            cui_idx = find_columns(headers, 'CUI', 'CUI_AC')
            cpv_code_idx = find_columns(headers, 'Cod CPV', 'COD_CPV')
            cpv_name_idx = find_columns(headers, 'Denumire CPV', 'DENUMIRE_CPV')
            notice_no_idx = find_columns(headers, 'Numar anunt atribuire', 'NUMAR_ANUNT_ATRIBUIRE')
            pub_date_idx = find_columns(headers, 'Data publicare', 'DATA_PUBLICARE')
            contract_date_idx = find_columns(headers, 'Data contract', 'DATA_CONTRACT')
            contract_no_idx = find_columns(headers, 'Numar contract', 'NUMAR_CONTRACT')
            value_idx = find_columns(headers, 'Valoare contract (RON)', 'VALOARE_CONTRACT_RON',
                                     'Valoare contract(RON)')
            winner_idx = find_columns(headers, 'Ofertant', 'OFERTANT', 'Ofertant castigator')
            winner_cui_idx = find_columns(headers, 'CUI ofertant', 'CUI_OFERTANT')
            winner_city_idx = find_columns(headers, 'Oras', 'ORAS', 'oras')
            proc_type_idx = find_columns(headers, 'Tip procedura', 'TIP_PROCEDURA')
            contract_type_idx = find_columns(headers, 'Tip contract', 'TIP_CONTRACT')

            if notice_no_idx is None and contract_no_idx is None:
                print(f"  [skip] Can't find notice/contract columns. Headers: {headers[:15]}")
                break

            def safe_str(idx):
                if idx is None: return None
                return str(row[idx]).strip() if row[idx] else None

            def safe_float(idx):
                if idx is None: return None
                v = row[idx]
                if v is None: return None
                try: return float(str(v).replace(',', '.').replace(' ', ''))
                except: return None

            def safe_date(idx):
                if idx is None: return None
                v = row[idx]
                if isinstance(v, datetime): return v
                if v is None: return None
                try: return datetime.fromisoformat(str(v).replace('/', '-'))
                except: return None

            notice_no = safe_str(notice_no_idx) or safe_str(contract_no_idx)
            if not notice_no:
                continue

            auth_cui = safe_str(cui_idx)
            if auth_cui:
                auth_cui = auth_cui.replace('RO', '').replace('ro', '').strip()

            batch.append((
                notice_no,
                safe_str(auth_idx),
                auth_cui,
                safe_str(cpv_code_idx),
                safe_str(cpv_name_idx),
                safe_float(value_idx),
                safe_date(pub_date_idx),
                safe_date(contract_date_idx),
                safe_str(proc_type_idx),
                safe_str(contract_type_idx),
                safe_str(winner_idx),
                safe_str(winner_cui_idx),
                safe_str(winner_city_idx),
            ))
            file_rows += 1

            if len(batch) >= 5000:
                inserted = _insert_contract_batch(cur, batch)
                total += inserted
                batch = []
                print(f"  [{label}] {file_rows} rows, {total} inserted...")

        if batch:
            inserted = _insert_contract_batch(cur, batch)
            total += inserted

        conn.commit()
        print(f"  [done] {label}: {file_rows} rows, {total} total inserted")

    return total


def _insert_contract_batch(cur, batch):
    if not batch:
        return 0

    inserted = 0
    for row in batch:
        (notice_no, auth_name, auth_cui, cpv_code, cpv_name, value,
         pub_date, contract_date, proc_type, contract_type,
         winner_name, winner_cui, winner_city) = row

        try:
            cur.execute("""
                INSERT INTO seap.public_notices
                    (id, notice_no, contract_title, cpv_code, cpv_name,
                     contract_value, publication_date, state_date,
                     procedure_type_text, contract_type_text, state_text,
                     authority_cui, authority_city, authority_county)
                VALUES (
                    nextval('seap.pn_import_seq'),
                    %s, %s, %s, %s, %s, %s, %s, %s, %s,
                    'Importat data.gov.ro', %s, NULL, NULL
                )
                RETURNING id
            """, (notice_no, auth_name, cpv_code, cpv_name, value,
                  pub_date, contract_date, proc_type, contract_type, auth_cui))

            result = cur.fetchone()
            if result:
                notice_id = result[0]
                if winner_name:
                    cur.execute("""
                        INSERT INTO seap.notice_contracts
                            (notice_id, winner_name, winner_fiscal, winner_city, contract_value, contract_date)
                        VALUES (%s, %s, %s, %s, %s, %s)
                    """, (notice_id, winner_name, winner_cui, winner_city, value, contract_date))
                inserted += 1
                cur.connection.commit()
        except Exception as e:
            cur.connection.rollback()
            continue

    return inserted


# ── Step 4: Resolve CUI → location and update SIRUTA ──

def resolve_locations(conn):
    """Map CUI to SIRUTA using cui_location table, update DAs and notices."""
    cur = conn.cursor()

    # Step A: Match cui_location entries to SIRUTA
    cur.execute("""
        ALTER TABLE seap.cui_location ADD COLUMN IF NOT EXISTS siruta TEXT
    """)
    conn.commit()

    # Exact match
    cur.execute("""
        UPDATE seap.cui_location cl
        SET siruta = u.siruta
        FROM public."GisUat" u
        WHERE cl.siruta IS NULL
          AND cl.city IS NOT NULL AND cl.county IS NOT NULL
          AND seap.normalize_locality(u.name) = seap.normalize_locality(cl.city)
          AND seap.normalize_locality(u.county) = seap.normalize_locality(cl.county)
    """)
    exact = cur.rowcount
    print(f"  [exact match] {exact} CUIs matched to SIRUTA")

    # Fuzzy match
    cur.execute("""
        UPDATE seap.cui_location cl
        SET siruta = sub.siruta
        FROM (
            SELECT DISTINCT ON (cl2.cui)
                cl2.cui,
                u.siruta,
                similarity(seap.normalize_locality(u.name), seap.normalize_locality(cl2.city)) AS score
            FROM seap.cui_location cl2
            JOIN public."GisUat" u
                ON seap.normalize_locality(u.county) = seap.normalize_locality(cl2.county)
            WHERE cl2.siruta IS NULL AND cl2.city IS NOT NULL AND cl2.county IS NOT NULL
              AND similarity(seap.normalize_locality(u.name), seap.normalize_locality(cl2.city)) > 0.3
            ORDER BY cl2.cui, score DESC
        ) sub
        WHERE cl.cui = sub.cui
    """)
    fuzzy = cur.rowcount
    print(f"  [fuzzy match] {fuzzy} more CUIs matched")
    conn.commit()

    # Step B: Update DAs — set authority_siruta via CUI lookup
    # First, add authority_siruta column to DA if not exists
    cur.execute("""
        ALTER TABLE seap.direct_acquisitions ADD COLUMN IF NOT EXISTS authority_siruta TEXT
    """)
    conn.commit()

    cur.execute("""
        UPDATE seap.direct_acquisitions da
        SET authority_siruta = cl.siruta
        FROM seap.cui_location cl
        WHERE da.authority_cui = cl.cui
          AND da.authority_siruta IS NULL
          AND cl.siruta IS NOT NULL
    """)
    da_matched = cur.rowcount
    print(f"  [DA location] {da_matched} acquisitions linked to SIRUTA")

    # Step C: Update notices — set authority_siruta via CUI
    cur.execute("""
        UPDATE seap.public_notices pn
        SET authority_siruta = cl.siruta
        FROM seap.cui_location cl
        WHERE pn.authority_cui = cl.cui
          AND pn.authority_siruta IS NULL
          AND cl.siruta IS NOT NULL
    """)
    pn_matched = cur.rowcount
    print(f"  [Notice location] {pn_matched} notices linked to SIRUTA")
    conn.commit()

    # Step D: Rebuild materialized view to use CUI-based matching
    print("  [refresh] Dropping and recreating materialized view...")
    cur.execute("DROP MATERIALIZED VIEW IF EXISTS seap.uat_procurement_stats")
    cur.execute("""
        CREATE MATERIALIZED VIEW seap.uat_procurement_stats AS
        SELECT
          u.siruta,
          u.name AS uat_name,
          u.county,
          COALESCE(da_s.da_count, 0)::bigint AS da_count,
          COALESCE(da_s.da_total_value, 0)::numeric AS da_total_value,
          COALESCE(pn_s.notice_count, 0)::bigint AS notice_count,
          COALESCE(pn_s.notice_total_value, 0)::numeric AS notice_total_value,
          (COALESCE(da_s.da_count, 0) + COALESCE(pn_s.notice_count, 0))::bigint AS total_contracts,
          (COALESCE(da_s.da_total_value, 0) + COALESCE(pn_s.notice_total_value, 0))::numeric AS total_value
        FROM public."GisUat" u
        LEFT JOIN (
          SELECT authority_siruta AS siruta,
                 COUNT(*) AS da_count,
                 SUM(closing_value) AS da_total_value
          FROM seap.direct_acquisitions
          WHERE authority_siruta IS NOT NULL
          GROUP BY authority_siruta
        ) da_s ON da_s.siruta = u.siruta
        LEFT JOIN (
          SELECT authority_siruta AS siruta,
                 COUNT(*) AS notice_count,
                 SUM(contract_value) AS notice_total_value
          FROM seap.public_notices
          WHERE authority_siruta IS NOT NULL
          GROUP BY authority_siruta
        ) pn_s ON pn_s.siruta = u.siruta
    """)
    cur.execute("CREATE UNIQUE INDEX idx_ups_siruta ON seap.uat_procurement_stats(siruta)")
    conn.commit()
    print("  [done] Materialized view rebuilt")


# ── Main ──

def main():
    mode = sys.argv[1] if len(sys.argv) > 1 else 'all'

    conn = get_conn()

    # Fix: DA table needs auto-increment ID since data.gov has no numeric IDs
    cur = conn.cursor()
    cur.execute("""
        DO $$ BEGIN
            IF NOT EXISTS (SELECT 1 FROM pg_sequences WHERE schemaname = 'seap' AND sequencename = 'da_import_seq') THEN
                CREATE SEQUENCE seap.da_import_seq START WITH 200000000;
            END IF;
        END $$;
    """)
    cur.execute("""
        DO $$ BEGIN
            IF NOT EXISTS (SELECT 1 FROM pg_sequences WHERE schemaname = 'seap' AND sequencename = 'pn_import_seq') THEN
                CREATE SEQUENCE seap.pn_import_seq START WITH 500000000;
            END IF;
        END $$;
    """)
    conn.commit()

    print(f"\n=== data.gov.ro Import — {datetime.now().isoformat()} ===\n")

    years = {'2025': URLS_2025, '2026': URLS_2026}

    if mode in ('anunturi', 'all'):
        print("── Step 1: Anunturi initiere (CUI → location mapping) ──")
        for year, urls in years.items():
            if 'anunturi_initiere' in urls:
                print(f"\n  [{year}]")
                count = import_anunturi_initiere(conn, urls['anunturi_initiere'])
                print(f"  [{year}] Total: {count} CUI mappings\n")

    if mode in ('da', 'all'):
        print("── Step 2: Achizitii directe ──")
        for year, urls in years.items():
            if 'achizitii_directe' in urls:
                print(f"\n  [{year}]")
                count = import_achizitii_directe(conn, urls['achizitii_directe'])
                print(f"  [{year}] Total: {count} direct acquisitions\n")

    if mode in ('contracte', 'all'):
        print("── Step 3: Contracte ──")
        for year, urls in years.items():
            if 'contracte' in urls:
                print(f"\n  [{year}]")
                count = import_contracte(conn, urls['contracte'])
                print(f"  [{year}] Total: {count} contracts\n")

    if mode in ('resolve', 'all'):
        print("── Step 4: Resolve locations → SIRUTA ──")
        resolve_locations(conn)

    # Final stats
    cur = conn.cursor()
    cur.execute("SELECT COUNT(*) FROM seap.direct_acquisitions")
    da = cur.fetchone()[0]
    cur.execute("SELECT COUNT(*) FROM seap.public_notices")
    pn = cur.fetchone()[0]
    cur.execute("SELECT COUNT(*) FROM seap.entities WHERE siruta IS NOT NULL")
    matched = cur.fetchone()[0]
    cur.execute("SELECT COUNT(*) FROM seap.uat_procurement_stats WHERE total_contracts > 0")
    uats = cur.fetchone()[0]

    print(f"\n=== Done ===")
    print(f"  DA: {da}, Notices: {pn}, Matched entities: {matched}, UATs with data: {uats}")

    conn.close()


if __name__ == '__main__':
    main()
