"""Daily 19:05 cron job — finalize variant choice, build next batch, notify group chat."""
import asyncio
import csv
import json
import os
import re
import sqlite3
from datetime import datetime
from pathlib import Path
from typing import Dict, List, Set, Tuple

BASE_DIR   = Path(__file__).parent.parent
VARIANTS_FILE = BASE_DIR / "data/processed/aima_pending_variants.json"
STATE_FILE    = BASE_DIR / "data/processed/aima_bot_state.json"
DB            = BASE_DIR / "data/processed/aima_conversion_shadow.sqlite"
BATCH_SIZE    = int(os.environ.get("AIMA_NEXT_BATCH_SIZE", "20"))
CHAT_ID       = -5178178519
MARKETER      = "@anatolii_malinovskyi"

DEFAULT_ENV_FILES = [
    BASE_DIR / ".env",
    Path(r"C:\python_scripts\top_1_telegram_signals\.env"),
]

CONTACTS_FIELDS = [
    "pilot_id", "pilot_name", "lead_id", "dataset", "segment", "variant",
    "hypothesis", "gate1_text", "phone", "first_name", "last_name",
    "email", "shop_name", "shop_domain", "registered_at", "last_activity_at", "worker-do-not-send",
    "pre_contact_quality", "contamination_status", "notes",
]

LOG_FIELDS = [
    "pilot_id", "lead_id", "dataset", "hypothesis", "phone",
    "worker-do-not-send", "added_to_telegram", "telegram_contact_name",
    "gate1_message_sent", "replied", "reply_text", "next_gate", "notes",
    "sent_at_utc", "recipient_id", "message_len", "message_preview",
    "chat_started", "replied_at_utc", "reply_category",
    "pre_contact_quality", "contamination_status",
    "qualified_product_reply", "next_step_accepted", "negative_stop", "reply_class",
]


TG_SESSION        = BASE_DIR / "data/processed/telegram/aima_support_session"
SCREEN_BATCH_SIZE = 20  # contacts per ImportContactsRequest call
IMPORT_TIMEOUT_SECONDS = 30.0
ENTITY_TIMEOUT_SECONDS = 3.0


def normalize_phone(value):
    # type: (str) -> str
    return "".join(ch for ch in str(value or "") if ch.isdigit() or ch == "+")


async def screen_candidates_for_telegram(candidates, needed):
    # type: (list, int) -> list
    """Return first `needed` candidates that have a Telegram account.

    Imports candidates in batches via ImportContactsRequest and keeps those
    for which Telegram returns a User object.  Stops as soon as `needed`
    confirmed contacts are accumulated.
    """
    load_env()
    api_id   = os.environ.get("TG_API_ID")
    api_hash = os.environ.get("TG_API_HASH")
    if not api_id or not api_hash:
        raise RuntimeError("TG_API_ID/TG_API_HASH missing from .env")

    from telethon import TelegramClient
    from telethon.tl.functions.contacts import ImportContactsRequest
    from telethon.tl.types import InputPhoneContact

    client = TelegramClient(str(TG_SESSION), int(api_id), api_hash)
    await client.connect()
    if not await client.is_user_authorized():
        raise RuntimeError("Session not authorized — run aima_telegram_login.py first")

    found = []  # type: list
    total_scanned = 0

    for offset in range(0, len(candidates), SCREEN_BATCH_SIZE):
        if len(found) >= needed:
            break
        chunk = candidates[offset:offset + SCREEN_BATCH_SIZE]
        contacts_req = [
            InputPhoneContact(
                client_id=offset + i,
                phone=r["phone"],
                first_name=r.get("first_name") or r["phone"][-4:],
                last_name=r.get("last_name") or "",
            )
            for i, r in enumerate(chunk)
        ]
        try:
            result = await asyncio.wait_for(
                client(ImportContactsRequest(contacts_req)),
                timeout=IMPORT_TIMEOUT_SECONDS,
            )
        except Exception as exc:
            print("[screen] import batch failed offset={} size={}: {}".format(
                offset, len(chunk), type(exc).__name__))
            total_scanned += len(chunk)
            continue
        users_by_phone = {
            normalize_phone(getattr(u, "phone", "") or "").lstrip("+"): u
            for u in result.users
        }
        # retry_contacts: user exists but has privacy settings.
        # With Telegram Premium on sender account, get_entity may now resolve them.
        # Use asyncio.wait_for with 5s timeout to prevent hangs.
        retry_client_ids = {rc for rc in getattr(result, "retry_contacts", [])}
        if retry_client_ids:
            cid_map = {offset + i: r for i, r in enumerate(chunk)}

            async def resolve_retry(cid):
                if cid not in cid_map:
                    return None, None
                r = cid_map[cid]
                try:
                    user = await asyncio.wait_for(
                        client.get_entity(r["phone"]),
                        timeout=ENTITY_TIMEOUT_SECONDS,
                    )
                    return normalize_phone(r["phone"]).lstrip("+"), user
                except Exception:
                    return None, None

            retry_results = await asyncio.gather(
                *(resolve_retry(cid) for cid in retry_client_ids),
                return_exceptions=False,
            )
            for phone_key, user in retry_results:
                if phone_key and user:
                    users_by_phone[phone_key] = user
        for r in chunk:
            user = users_by_phone.get(normalize_phone(r["phone"]).lstrip("+"))
            if user:
                r = dict(r)
                r["recipient_id"] = str(getattr(user, "id", ""))
                r["telegram_contact_name"] = " ".join(filter(None, [
                    getattr(user, "first_name", "") or "",
                    getattr(user, "last_name", "") or "",
                ])) or getattr(user, "username", "") or r["recipient_id"]
                found.append(r)
                if len(found) >= needed:
                    break
        total_scanned += len(chunk)
        retry_n = len(getattr(result, "retry_contacts", []))
        print("[screen] scanned={} tg_in_batch={} retry={} total_found={}".format(
            total_scanned, len(result.users), retry_n, len(found)))

    await client.disconnect()
    print("[screen] done: {} з Telegram знайдено (перевірено {} кандидатів)".format(
        len(found), total_scanned))
    return found[:needed]


def load_env():
    for path in DEFAULT_ENV_FILES:
        if not path.exists():
            continue
        for raw in path.read_text(encoding="utf-8", errors="ignore").splitlines():
            line = raw.strip()
            if not line or line.startswith("#") or "=" not in line:
                continue
            k, v = line.split("=", 1)
            os.environ.setdefault(k.strip(), v.strip().strip('"').strip("'"))


def collect_used():
    # type: () -> Tuple[Set[str], Set[str], Set[str]]
    sent_leads = set()  # type: Set[str]
    sent_phones = set()  # type: Set[str]
    blocked_phones = set()  # type: Set[str]
    proc = BASE_DIR / "data/processed"
    for p in sorted(proc.glob("*telegram_log.csv")):
        with p.open(encoding="utf-8-sig") as f:
            for row in csv.DictReader(f):
                lid = str(row.get("lead_id", "") or "").strip()
                ph  = str(row.get("phone",   "") or "").strip()
                if str(row.get("worker-do-not-send", "") or "").strip().casefold() in {"1", "yes", "true", "так"}:
                    if ph:
                        blocked_phones.add(ph)
                if str(row.get("gate1_message_sent", "") or "").strip().casefold() == "yes":
                    if lid:
                        sent_leads.add(lid)
                    if ph:
                        sent_phones.add(ph)
    ledger = proc / "aima_telegram_message_ledger.csv"
    if ledger.exists():
        with ledger.open(encoding="utf-8-sig") as f:
            for row in csv.DictReader(f):
                if str(row.get("status", "") or "").strip().casefold() == "sent":
                    lid = str(row.get("lead_id", "") or "").strip()
                    ph = str(row.get("phone", "") or "").strip()
                    if lid:
                        sent_leads.add(lid)
                    if ph:
                        sent_phones.add(ph)
    return sent_leads, sent_phones, blocked_phones


def next_batch_key():
    # type: () -> str
    nums = []
    for p in (BASE_DIR / "data/processed").glob("aima_far_v*_contacts.csv"):
        m = re.search(r"v(\d+)_contacts", p.name)
        if m:
            nums.append(int(m.group(1)))
    return "v{}".format(max(nums) + 1) if nums else "v7"


def next_pilot_id_start():
    # type: () -> int
    max_id = 120
    for p in (BASE_DIR / "data/processed").glob("*_contacts.csv"):
        with p.open(encoding="utf-8-sig") as f:
            for row in csv.DictReader(f):
                try:
                    max_id = max(max_id, int(row.get("pilot_id", 0) or 0))
                except (ValueError, TypeError):
                    pass
    return max_id + 1


def build_next_batch(batch_key, variant_texts, pilot_id_start):
    # type: (str, Dict[str, str], int) -> Tuple[Path, Path]
    sent_leads, sent_phones, blocked_phones = collect_used()

    conn = sqlite3.connect(str(DB))
    conn.row_factory = sqlite3.Row
    cur = conn.cursor()
    cur.execute(
        """SELECT dataset, lead_id, first_name, last_name, phone, email, shop_name, shop_domain,
                  registered_at, last_activity_at, row_index
           FROM aima_imported_contacts
           WHERE phone IS NOT NULL AND phone != ''
           ORDER BY dataset, row_index DESC"""
    )
    all_rows = [dict(row) for row in cur.fetchall()]
    conn.close()

    candidates = []
    seen_phones = set()  # type: Set[str]
    for row in all_rows:
        phone = normalize_phone(row.get("phone", ""))
        row["phone"] = phone
        has_entered_data = bool(row.get("first_name") or row.get("last_name") or row.get("email"))
        has_shop_data = bool(row.get("shop_name") or row.get("shop_domain"))
        if row["dataset"] == "users_without_shop_202605191510" and has_entered_data:
            row["segment"] = "REGISTERED_WITH_DATA_NO_SHOP"
            row["warm_rank"] = "1"
        elif row["dataset"] == "lost_shops_202605191501" and has_shop_data:
            row["segment"] = "LOST_SHOP_WITH_STORE_DATA"
            row["warm_rank"] = "2"
        else:
            continue
        if not phone or row["lead_id"] in sent_leads or phone in sent_phones:
            continue
        if phone in blocked_phones or phone in seen_phones:
            continue
        if len("".join(ch for ch in phone if ch.isdigit())) != 12:
            continue
        candidates.append(row)
        seen_phones.add(phone)

    if not candidates:
        raise RuntimeError("No unprocessed contacts available in dataset")

    candidates.sort(
        key=lambda row: (
            int(row["warm_rank"]),
            row.get("last_activity_at") or row.get("registered_at") or "",
        ),
        reverse=False,
    )
    candidates.sort(key=lambda row: row.get("last_activity_at") or row.get("registered_at") or "", reverse=True)
    candidates.sort(key=lambda row: int(row["warm_rank"]))

    # Pre-screen: scan candidates in batches, keep only those with Telegram accounts
    selected = asyncio.get_event_loop().run_until_complete(
        screen_candidates_for_telegram(candidates, BATCH_SIZE)
    )
    if len(selected) < BATCH_SIZE:
        raise RuntimeError(
            "Not enough Telegram-ready contacts: {}/{} found (scanned {} candidates); batch not activated".format(
                len(selected), BATCH_SIZE, len(candidates)
            )
        )
    v_keys = list(variant_texts.keys())
    pilot_name = "manual_registered_warm_{}".format(batch_key)

    contacts_csv = BASE_DIR / "data/processed/aima_registered_warm_{}_contacts.csv".format(batch_key)
    log_csv      = BASE_DIR / "data/processed/aima_registered_warm_{}_telegram_log.csv".format(batch_key)

    contact_rows = []
    log_rows = []

    for i, row in enumerate(selected):
        pilot_id = pilot_id_start + i
        variant  = v_keys[i % len(v_keys)]
        gate1_text = variant_texts[variant]

        contact_rows.append({
            "pilot_id": pilot_id,
            "pilot_name": pilot_name,
            "lead_id": row["lead_id"],
            "dataset": row["dataset"],
            "segment": row["segment"],
            "variant": variant,
            "hypothesis": variant,
            "gate1_text": gate1_text,
            "phone": row["phone"],
            "first_name": row.get("first_name") or "",
            "last_name": row.get("last_name") or "",
            "email": row.get("email") or "",
            "shop_name": row.get("shop_name") or "",
            "shop_domain": row.get("shop_domain") or "",
            "registered_at": row.get("registered_at") or "",
            "last_activity_at": row.get("last_activity_at") or "",
            "worker-do-not-send": "",
            "pre_contact_quality": "warm_registered_data",
            "contamination_status": "telegram_preflight_ok",
            "notes": "prepared {}; Telegram resolved; recipient_id={}; not sent".format(
                batch_key, row.get("recipient_id", "")),
        })

        log_row = {k: "" for k in LOG_FIELDS}
        log_row.update({
            "pilot_id": pilot_id,
            "lead_id": row["lead_id"],
            "dataset": row["dataset"],
            "hypothesis": variant,
            "phone": row["phone"],
            "worker-do-not-send": "",
            "added_to_telegram": "yes",
            "telegram_contact_name": row.get("telegram_contact_name", ""),
            "recipient_id": row.get("recipient_id", ""),
            "pre_contact_quality": "warm_registered_data",
            "contamination_status": "telegram_preflight_ok",
            "notes": "prepared {}; Telegram resolved; not sent".format(batch_key),
        })
        log_rows.append(log_row)

    with contacts_csv.open("w", encoding="utf-8", newline="") as f:
        writer = csv.DictWriter(f, fieldnames=CONTACTS_FIELDS)
        writer.writeheader()
        writer.writerows(contact_rows)

    with log_csv.open("w", encoding="utf-8", newline="") as f:
        writer = csv.DictWriter(f, fieldnames=LOG_FIELDS)
        writer.writeheader()
        writer.writerows(log_rows)

    return contacts_csv, log_csv


def write_send_script(batch_key):
    # type: (str) -> None
    """Generate send script for next registered-warm batch."""
    src_path = BASE_DIR / "src/aima_batch_send_{}.py".format(batch_key)
    template_path = BASE_DIR / "src/aima_batch_send_v19.py"
    if not template_path.exists():
        template_path = BASE_DIR / "src/aima_batch_send_v6.py"
    if not template_path.exists():
        print("[warn] send script template not found, cannot generate {}".format(src_path.name))
        return
    src = template_path.read_text(encoding="utf-8")
    src = re.sub(r"aima_(?:far|registered_warm)_v\d+_contacts\.csv",
                 "aima_registered_warm_{}_contacts.csv".format(batch_key), src)
    src = re.sub(r"aima_(?:far|registered_warm)_v\d+_telegram_log\.csv",
                 "aima_registered_warm_{}_telegram_log.csv".format(batch_key), src)
    src = src.replace("batch_send_v6", "batch_send_{}".format(batch_key))
    src = src.replace("batch_send_v19", "batch_send_{}".format(batch_key))
    src = src.replace("batch_send_v17", "batch_send_{}".format(batch_key))
    src = re.sub(r"batch_send_v\d+_registered_warm", "batch_send_{}_registered_warm".format(batch_key), src)
    src = re.sub(r"AIMA (?:FAR|registered-with-data warm|v\d+ warm) [^\\n]*pilot",
                 "AIMA registered warm {} pilot".format(batch_key), src)
    src_path.write_text(src, encoding="utf-8")
    print("[script] wrote {}".format(src_path.name))


def send_telegram(token, chat_id, text):
    import requests
    r = requests.post(
        "https://api.telegram.org/bot{}/sendMessage".format(token),
        json={"chat_id": chat_id, "text": text, "parse_mode": "HTML"},
        timeout=30,
    )
    return r.json()


def main():
    load_env()
    token = os.environ.get("AIMA_BOT_TOKEN")
    if not token:
        raise SystemExit("AIMA_BOT_TOKEN missing")

    if not VARIANTS_FILE.exists():
        print("[warn] {} not found — run aima_agent_prepare.py first".format(VARIANTS_FILE.name))
        return

    data = json.loads(VARIANTS_FILE.read_text(encoding="utf-8"))
    status   = data.get("status", "pending_choice")
    variants = data.get("variants", [])  # type: List[dict]
    chosen   = data.get("chosen", [])    # type: List[str]
    locked_label = data.get("locked_label") or "A"

    if not variants:
        print("[warn] no variants in file")
        return

    available_labels = [v["label"] for v in variants]

    # Resolve chosen labels. A is a locked champion/control; user selects only B/C/D.
    if status == "chosen":
        competitor = next((lb for lb in chosen if lb != locked_label and lb in available_labels), None)
        if competitor:
            chosen_labels = [locked_label, competitor]
        else:
            chosen_labels = [locked_label]
        auto_picked = False
    else:
        # Auto-pick locked A and the best ranked competitor after A.
        if locked_label in available_labels:
            competitor = next((lb for lb in available_labels if lb != locked_label), None)
            chosen_labels = [locked_label, competitor] if competitor else [locked_label]
        elif len(variants) >= 3:
            chosen_labels = [variants[0]["label"], variants[2]["label"]]
        else:
            chosen_labels = [v["label"] for v in variants[:2]]
        auto_picked = True

    chosen_variants = [v for v in variants if v["label"] in chosen_labels]
    if len(chosen_variants) < 2:
        fallback_competitor = next((v for v in variants if v["label"] != locked_label), None)
        locked_variant = next((v for v in variants if v["label"] == locked_label), None)
        chosen_variants = [v for v in [locked_variant, fallback_competitor] if v]
    if len(chosen_variants) < 2:
        print("[warn] cannot resolve two variants from labels {}".format(chosen_labels))
        return

    print("[pick] {}: {}".format("auto" if auto_picked else "user",
                                  [v["label"] for v in chosen_variants]))

    # Build variant_texts dict keyed by hypothesis name
    variant_texts = {}  # type: Dict[str, str]
    for v in chosen_variants:
        key = "FAR_{}_agent_gen".format(v["label"])
        variant_texts[key] = v["text"]

    batch_key      = next_batch_key()
    pilot_id_start = next_pilot_id_start()

    # Build batch CSVs
    try:
        contacts_csv, log_csv = build_next_batch(batch_key, variant_texts, pilot_id_start)
        print("[batch] {} -> {}".format(batch_key, contacts_csv.name))
        import csv as _csv
        with open(str(contacts_csv), encoding="utf-8-sig") as _f:
            n_contacts = sum(1 for _ in _csv.DictReader(_f))
    except RuntimeError as exc:
        print("[error] {}".format(exc))
        send_telegram(token, CHAT_ID,
                      "⚠️ AIMA агент: не вдалося побудувати батч {}: {}".format(batch_key, exc))
        return

    # Generate send script
    write_send_script(batch_key)

    # Update bot state
    state = {}
    if STATE_FILE.exists():
        state = json.loads(STATE_FILE.read_text(encoding="utf-8"))

    send_script = str(BASE_DIR / "src/aima_batch_send_{}.py".format(batch_key))
    gdrive_url  = (state.get("proposals", {}).get("v6") or {}).get("gdrive_url", "")
    state.setdefault("proposals", {})[batch_key] = {
        "name": "REGISTERED WARM {} — {} Telegram-ready контактів (агент)".format(batch_key, n_contacts),
        "contacts_csv": str(contacts_csv),
        "log_csv":      str(log_csv),
        "send_script":  send_script,
        "gdrive_url":   gdrive_url,
    }
    state["active_batch"] = batch_key
    state["pending"]      = batch_key
    # Only reset daily if previous cycle is fully done (sent/postponed/rejected).
    # If status=approved or ping_sent, the current day's approval must be preserved.
    prev_daily_status = state.get("daily", {}).get("status", "idle")
    if prev_daily_status not in ("approved", "ping_sent"):
        state["daily"] = {
            "date": "", "status": "idle", "morning_msg_id": None,
            "send_triggered": False, "sent_count": None,
            "approved_by": None, "approved_at": None, "sent_at": None,
            "reminder_sent": False,
        }
    STATE_FILE.write_text(json.dumps(state, ensure_ascii=False, indent=2), encoding="utf-8")
    print("[state] active_batch={} (daily preserved={})".format(batch_key, prev_daily_status))

    # Finalize variants file
    data["status"]           = "finalized"
    data["finalized_labels"] = [v["label"] for v in chosen_variants]
    data["finalized_at"]     = datetime.now().isoformat()
    VARIANTS_FILE.write_text(json.dumps(data, ensure_ascii=False, indent=2), encoding="utf-8")

    # Send group message
    auto_note = " (автовибір)" if auto_picked else ""
    chosen_info = "\n".join(
        "  <b>{}:</b> {}".format(v["label"], v["text"][:160])
        for v in chosen_variants
    )
    msg = (
        "\U0001f514 {}, підготовка наступної розсилки!\n\n"
        "Батч: <b>REGISTERED WARM {}</b> | Telegram-ready контактів: <b>{}</b>\n"
        "Гіпотези{}:\n{}\n\n"
        "Завтра о <b>08:00</b> бот нагадає для підтвердження.\n"
        "/status — поточний стан"
    ).format(os.environ.get("AIMA_MARKETER", MARKETER), batch_key, n_contacts, auto_note, chosen_info)

    result = send_telegram(token, CHAT_ID, msg)
    if result.get("ok"):
        print("[sent] group proposal for {}".format(batch_key))
    else:
        print("[warn] telegram error: {}".format(result))


if __name__ == "__main__":
    main()