import hashlib
import json
import os
import queue
import signal
import sqlite3
import sys
import time
import urllib.error
import urllib.request
from datetime import datetime, timezone
from zoneinfo import ZoneInfo

from wcferry import Wcf


WCF_HOST = os.environ.get("WCF_HOST", "127.0.0.1")
WCF_COMMAND_PORT = int(os.environ.get("WCF_COMMAND_PORT", "10086"))
STATE_ROOT = os.environ.get("STATE_ROOT", "/var/lib/unidesk/personal-wechat")
ARCHIVE_CALLBACK_URL = os.environ.get("ARCHIVE_CALLBACK_URL", "")
ARCHIVE_CALLBACK_TOKEN = os.environ.get("ARCHIVE_CALLBACK_TOKEN", "")
ARCHIVE_REMOTE_ROOT = os.environ.get("ARCHIVE_REMOTE_ROOT", "/UniDesk/WeChatArchive")
ARCHIVE_TIMEZONE = os.environ.get("ARCHIVE_TIMEZONE", "Asia/Shanghai")
POLL_IDLE_SECONDS = float(os.environ.get("POLL_IDLE_SECONDS", "1"))
STATUS_INTERVAL_SECONDS = float(os.environ.get("STATUS_INTERVAL_SECONDS", "15"))

STOP = False


def now_iso():
    return datetime.now(timezone.utc).isoformat()


def log(event, **fields):
    payload = {"ts": now_iso(), "event": event, **fields}
    print(json.dumps(payload, ensure_ascii=False, sort_keys=True), flush=True)


def handle_signal(signum, _frame):
    global STOP
    STOP = True
    log("signal", signum=signum)


signal.signal(signal.SIGTERM, handle_signal)
signal.signal(signal.SIGINT, handle_signal)


def ensure_state():
    os.makedirs(STATE_ROOT, exist_ok=True)
    os.makedirs(os.path.join(STATE_ROOT, "messages"), exist_ok=True)
    db_path = os.path.join(STATE_ROOT, "collector.sqlite3")
    conn = sqlite3.connect(db_path, timeout=30)
    conn.execute("PRAGMA journal_mode=WAL")
    conn.execute(
        """
        CREATE TABLE IF NOT EXISTS messages (
          id TEXT PRIMARY KEY,
          observed_at TEXT NOT NULL,
          sender TEXT,
          roomid TEXT,
          type INTEGER,
          ts INTEGER,
          content_sha256 TEXT,
          payload_json TEXT NOT NULL,
          callback_status INTEGER,
          callback_error TEXT
        )
        """
    )
    conn.execute(
        """
        CREATE TABLE IF NOT EXISTS collector_state (
          key TEXT PRIMARY KEY,
          value TEXT NOT NULL,
          updated_at TEXT NOT NULL
        )
        """
    )
    conn.commit()
    return conn


def safe_message(msg):
    content = str(getattr(msg, "content", "") or "")
    xml = str(getattr(msg, "xml", "") or "")
    extra = str(getattr(msg, "extra", "") or "")
    thumb = str(getattr(msg, "thumb", "") or "")
    summary = {
        "id": str(getattr(msg, "id", "")),
        "ts": int(getattr(msg, "ts", 0) or 0),
        "type": int(getattr(msg, "type", 0) or 0),
        "sender": str(getattr(msg, "sender", "") or ""),
        "roomid": str(getattr(msg, "roomid", "") or ""),
        "fromSelf": bool(msg.from_self()) if hasattr(msg, "from_self") else False,
        "fromGroup": bool(msg.from_group()) if hasattr(msg, "from_group") else False,
        "contentSha256": hashlib.sha256(content.encode("utf-8", errors="replace")).hexdigest(),
        "contentPreview": content[:160],
        "xmlSha256": hashlib.sha256(xml.encode("utf-8", errors="replace")).hexdigest() if xml else "",
        "extraSha256": hashlib.sha256(extra.encode("utf-8", errors="replace")).hexdigest() if extra else "",
        "thumbSha256": hashlib.sha256(thumb.encode("utf-8", errors="replace")).hexdigest() if thumb else "",
    }
    return summary, content


def sanitize_path_segment(value, fallback="unknown"):
    text = str(value or fallback)
    cleaned = []
    for char in text:
        if char.isalnum() or char in "._-":
            cleaned.append(char)
        else:
            cleaned.append("-")
    normalized = "".join(cleaned).strip("-")[:120]
    return normalized or fallback


def archive_date(summary):
    try:
        zone = ZoneInfo(ARCHIVE_TIMEZONE)
    except Exception:
        zone = timezone.utc
    ts = int(summary.get("ts") or 0)
    if ts > 0:
        dt = datetime.fromtimestamp(ts, tz=zone)
    else:
        dt = datetime.now(zone)
    return dt.strftime("%Y-%m-%d")


def archive_payload(summary, content):
    message_id = summary["id"] or f"missing-{summary['ts']}-{summary['contentSha256'][:12]}"
    conversation_id = summary["roomid"] or summary["sender"] or "unknown"
    sender_id = summary["sender"] or "unknown"
    message_type = "text" if summary["type"] == 1 else "wcf-message"
    extension = "txt"
    filename = f"{sanitize_path_segment(message_id, 'message')}.{extension}"
    remote_path = "/".join([
        ARCHIVE_REMOTE_ROOT.rstrip("/"),
        archive_date(summary),
        sanitize_path_segment(conversation_id, "conversation"),
        filename,
    ])
    text = "\n".join([
        f"messageId={message_id}",
        f"observedAt={now_iso()}",
        f"sender={sender_id}",
        f"roomid={summary['roomid']}",
        f"type={summary['type']}",
        f"timestamp={summary['ts']}",
        f"contentSha256={summary['contentSha256']}",
        "",
        content,
    ])
    return {
        "messageId": message_id,
        "conversationId": conversation_id,
        "senderId": sender_id,
        "messageType": message_type,
        "text": text,
        "message": {
            "messageId": message_id,
            "conversationId": conversation_id,
            "senderId": sender_id,
            "messageType": message_type,
            "text": text,
            "receivedAt": now_iso(),
        },
        "archive": {
            "remotePath": remote_path,
            "filename": filename,
        },
        "metadata": {
            "source": "personal-wechat-collector",
            "wcfType": summary["type"],
            "contentSha256": summary["contentSha256"],
            "fromSelf": summary["fromSelf"],
            "fromGroup": summary["fromGroup"],
        },
    }


def post_callback(payload):
    if not ARCHIVE_CALLBACK_URL:
        return None, "callback-url-empty"
    body = json.dumps(payload, ensure_ascii=False).encode("utf-8")
    req = urllib.request.Request(ARCHIVE_CALLBACK_URL, data=body, method="POST")
    req.add_header("content-type", "application/json")
    if ARCHIVE_CALLBACK_TOKEN:
        req.add_header("authorization", f"Bearer {ARCHIVE_CALLBACK_TOKEN}")
        req.add_header("x-unidesk-wechat-archive-token", ARCHIVE_CALLBACK_TOKEN)
    try:
        with urllib.request.urlopen(req, timeout=30) as resp:
            resp.read(1024)
            return int(resp.status), ""
    except urllib.error.HTTPError as exc:
        return int(exc.code), str(exc)[:500]
    except Exception as exc:
        return None, f"{type(exc).__name__}: {exc}"[:500]


def save_message(conn, summary, callback_status, callback_error):
    message_id = summary["id"] or f"missing-{summary['ts']}-{summary['contentSha256'][:12]}"
    conn.execute(
        """
        INSERT INTO messages (
          id, observed_at, sender, roomid, type, ts, content_sha256, payload_json, callback_status, callback_error
        ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
        ON CONFLICT(id) DO UPDATE SET
          observed_at=excluded.observed_at,
          callback_status=excluded.callback_status,
          callback_error=excluded.callback_error
        """,
        (
            message_id,
            now_iso(),
            summary["sender"],
            summary["roomid"],
            summary["type"],
            summary["ts"],
            summary["contentSha256"],
            json.dumps(summary, ensure_ascii=False, sort_keys=True),
            callback_status,
            callback_error,
        ),
    )
    conn.commit()


def set_state(conn, key, value):
    conn.execute(
        "INSERT INTO collector_state(key, value, updated_at) VALUES (?, ?, ?) "
        "ON CONFLICT(key) DO UPDATE SET value=excluded.value, updated_at=excluded.updated_at",
        (key, str(value), now_iso()),
    )
    conn.commit()


def main():
    conn = ensure_state()
    log("collector-start", wcfHost=WCF_HOST, commandPort=WCF_COMMAND_PORT, stateRoot=STATE_ROOT, readOnly=True, sendCapability=False)
    wcf = Wcf(host=WCF_HOST, port=WCF_COMMAND_PORT, debug=False, block=False)
    receiving = False
    last_status = 0.0
    while not STOP:
      try:
        logged_in = bool(wcf.is_login())
        set_state(conn, "is_login", "true" if logged_in else "false")
        now = time.time()
        if now - last_status >= STATUS_INTERVAL_SECONDS:
            self_wxid = ""
            if logged_in:
                try:
                    self_wxid = wcf.get_self_wxid()
                    set_state(conn, "self_wxid_sha256", hashlib.sha256(self_wxid.encode()).hexdigest())
                except Exception as exc:
                    log("self-wxid-error", error=f"{type(exc).__name__}: {exc}"[:500])
            log("collector-status", isLogin=logged_in, receiving=receiving, selfWxidPresent=bool(self_wxid))
            last_status = now
        if logged_in and not receiving:
            receiving = bool(wcf.enable_receiving_msg())
            log("receiving-enabled", ok=receiving)
        if receiving:
            try:
                msg = wcf.get_msg(block=True)
            except queue.Empty:
                continue
            summary, content = safe_message(msg)
            payload = archive_payload(summary, content)
            status, error = post_callback(payload)
            save_message(conn, summary, status, error)
            log("message-archived", id=summary["id"], type=summary["type"], callbackStatus=status, callbackError=bool(error))
        else:
            time.sleep(POLL_IDLE_SECONDS)
      except Exception as exc:
        log("collector-loop-error", error=f"{type(exc).__name__}: {exc}"[:1000])
        time.sleep(3)
    log("collector-stop")


if __name__ == "__main__":
    main()
