#!/usr/bin/env python3
"""Split a Telegram HTML export into hash-bound, compiler-sized source artifacts."""

from __future__ import annotations

import argparse
import hashlib
import json
import re
from dataclasses import dataclass
from datetime import datetime
from html.parser import HTMLParser
from pathlib import Path, PurePosixPath
from typing import Any
from urllib.parse import unquote, urlsplit

SCHEMA = "livingip.telegramExportCorpus.v1"
PARSER_VERSION = "1"
DEFAULT_MAX_SEGMENT_CHARS = 100_000
ATTACHMENT_PREFIXES = frozenset(
    {"animations", "files", "photos", "round_video_messages", "stickers", "video_files", "voice_messages"}
)
VOID_TAGS = frozenset({"area", "base", "br", "col", "embed", "hr", "img", "input", "link", "meta", "param", "source", "track", "wbr"})
SECRET_VALUE_RE = re.compile(
    r"(?i)(\b(?:[a-z][a-z0-9_]*(?:_?key|_?token|_?secret|_?password|_?passcode|_?credential)|"
    r"x-api-key|authorization)\b\s*[:=]\s*(?:bearer\s+)?)([\"']?)([a-z0-9_./+=:-]{8,})([\"']?)"
)


class ExportError(RuntimeError):
    """Raised when an export cannot be converted without losing provenance."""


@dataclass(frozen=True)
class TelegramMessage:
    message_id: str
    timestamp: str
    sender: str
    reply_to: str | None
    text: str
    links: tuple[str, ...]


def sha256_bytes(value: bytes) -> str:
    return hashlib.sha256(value).hexdigest()


def canonical_json_bytes(value: Any) -> bytes:
    return json.dumps(value, ensure_ascii=False, separators=(",", ":"), sort_keys=True).encode("utf-8")


def _clean_text(chunks: list[str]) -> str:
    text = "".join(chunks).replace("\r\n", "\n").replace("\r", "\n")
    lines = [line.rstrip() for line in text.splitlines()]
    return "\n".join(lines).strip()


def _timestamp(value: str) -> str:
    try:
        return datetime.strptime(value, "%d.%m.%Y %H:%M:%S UTC%z").isoformat()
    except ValueError as exc:
        raise ExportError(f"unsupported Telegram timestamp: {value}") from exc


class TelegramHtmlParser(HTMLParser):
    def __init__(self) -> None:
        super().__init__(convert_charrefs=True)
        self.stack: list[tuple[str, frozenset[str]]] = []
        self.current: dict[str, Any] | None = None
        self.message_depth: int | None = None
        self.last_sender: str | None = None
        self.messages: list[TelegramMessage] = []
        self.title_chunks: list[str] = []

    def _inside(self, class_name: str) -> bool:
        return any(class_name in classes for _tag, classes in self.stack)

    def handle_starttag(self, tag: str, attrs: list[tuple[str, str | None]]) -> None:
        values = {key: value or "" for key, value in attrs}
        classes = frozenset(values.get("class", "").split())
        if tag in VOID_TAGS:
            if self.current is not None and tag == "br" and self._inside("text"):
                self.current["text"].append("\n")
            return
        if tag == "div" and {"message", "default"}.issubset(classes):
            if self.current is not None:
                raise ExportError("nested Telegram message containers are unsupported")
            raw_id = values.get("id", "")
            if not re.fullmatch(r"message\d+", raw_id):
                raise ExportError(f"invalid Telegram message id: {raw_id or '<missing>'}")
            self.current = {
                "message_id": raw_id.removeprefix("message"),
                "timestamp": None,
                "sender": [],
                "reply_to": None,
                "text": [],
                "links": [],
            }
            self.message_depth = len(self.stack)
        self.stack.append((tag, classes))
        if self.current is None:
            return
        if tag == "div" and "date" in classes and values.get("title"):
            self.current["timestamp"] = _timestamp(values["title"])
        if tag == "a" and values.get("href"):
            href = values["href"].strip()
            if self._inside("reply_to"):
                match = re.search(r"(?:go_to_message|message)(\d+)", href)
                if match:
                    self.current["reply_to"] = match.group(1)
            elif href and href not in self.current["links"]:
                self.current["links"].append(href)
    def handle_startendtag(self, tag: str, attrs: list[tuple[str, str | None]]) -> None:
        self.handle_starttag(tag, attrs)
        if tag not in VOID_TAGS:
            self.handle_endtag(tag)

    def handle_data(self, data: str) -> None:
        if self.current is None:
            if self._inside("page_header") and self._inside("text"):
                self.title_chunks.append(data)
            return
        if self._inside("from_name") and not self._inside("forwarded"):
            self.current["sender"].append(data)
        elif self._inside("text"):
            self.current["text"].append(data)

    def handle_endtag(self, tag: str) -> None:
        if not self.stack:
            return
        closing_message = tag == "div" and self.message_depth == len(self.stack) - 1
        self.stack.pop()
        if closing_message:
            self._finish_message()

    def _finish_message(self) -> None:
        assert self.current is not None
        sender = _clean_text(self.current["sender"]) or self.last_sender
        timestamp = self.current["timestamp"]
        if not sender:
            raise ExportError(f"message {self.current['message_id']} has no sender and no joined-message predecessor")
        if not timestamp:
            raise ExportError(f"message {self.current['message_id']} has no timestamp")
        self.last_sender = sender
        self.messages.append(
            TelegramMessage(
                message_id=self.current["message_id"],
                timestamp=timestamp,
                sender=sender,
                reply_to=self.current["reply_to"],
                text=_clean_text(self.current["text"]),
                links=tuple(self.current["links"]),
            )
        )
        self.current = None
        self.message_depth = None

    @property
    def chat_title(self) -> str:
        return _clean_text(self.title_chunks)


def parse_export(path: Path) -> tuple[str, list[TelegramMessage]]:
    try:
        raw = path.read_text(encoding="utf-8", errors="strict")
    except (OSError, UnicodeDecodeError) as exc:
        raise ExportError(f"could not read strict UTF-8 Telegram export: {exc}") from exc
    parser = TelegramHtmlParser()
    try:
        parser.feed(raw)
        parser.close()
    except (ValueError, AssertionError) as exc:
        raise ExportError(f"could not parse Telegram export: {exc}") from exc
    if parser.current is not None:
        raise ExportError("Telegram export ended inside a message")
    if not parser.messages:
        raise ExportError("Telegram export contains no default messages")
    return parser.chat_title, parser.messages


def redact_secret_values(value: str) -> tuple[str, int]:
    count = 0

    def replacement(match: re.Match[str]) -> str:
        nonlocal count
        count += 1
        opening = match.group(2)
        closing = match.group(4) if match.group(4) == opening else opening
        return f"{match.group(1)}{opening}[REDACTED_SECRET]{closing}"

    return SECRET_VALUE_RE.sub(replacement, value), count


def _redact_messages(messages: list[TelegramMessage]) -> tuple[list[TelegramMessage], list[str], int]:
    redacted = []
    affected_ids = []
    value_count = 0
    for message in messages:
        text, text_count = redact_secret_values(message.text)
        links = []
        link_count = 0
        for link in message.links:
            safe_link, count = redact_secret_values(link)
            links.append(safe_link)
            link_count += count
        count = text_count + link_count
        if count:
            affected_ids.append(message.message_id)
            value_count += count
        redacted.append(
            TelegramMessage(
                message_id=message.message_id,
                timestamp=message.timestamp,
                sender=message.sender,
                reply_to=message.reply_to,
                text=text,
                links=tuple(links),
            )
        )
    return redacted, affected_ids, value_count


def _format_message(message: TelegramMessage, source_sha256: str) -> str:
    metadata = {
        "id": message.message_id,
        "reply_to": message.reply_to,
        "sender": message.sender,
        "source_ref": f"telegram-export:sha256:{source_sha256}#message={message.message_id}",
        "timestamp": message.timestamp,
    }
    parts = [json.dumps(metadata, ensure_ascii=False, separators=(",", ":"), sort_keys=True)]
    parts.append(message.text or "[attachment-only message]")
    if message.links:
        parts.append("Links:\n" + "\n".join(f"- {link}" for link in message.links))
    return "\n".join(parts).strip() + "\n"


def _split_messages(messages: list[TelegramMessage], source_sha256: str, max_chars: int) -> list[tuple[str, ...]]:
    if max_chars < 1_000 or max_chars > 110_000:
        raise ExportError("max segment characters must be between 1000 and 110000")
    segments: list[list[str]] = []
    current: list[str] = []
    current_chars = 0
    for message in messages:
        rendered = _format_message(message, source_sha256)
        if len(rendered) > max_chars:
            raise ExportError(f"message {message.message_id} exceeds the segment character limit")
        separator_chars = 1 if current else 0
        if current and current_chars + separator_chars + len(rendered) > max_chars:
            segments.append(current)
            current = []
            current_chars = 0
            separator_chars = 0
        current.append(rendered)
        current_chars += separator_chars + len(rendered)
    if current:
        segments.append(current)
    return [tuple(segment) for segment in segments]


def _attachment_relative_path(href: str) -> PurePosixPath | None:
    parsed = urlsplit(href)
    if parsed.scheme or parsed.netloc or not parsed.path:
        return None
    relative = PurePosixPath(unquote(parsed.path))
    if not relative.parts or relative.parts[0] not in ATTACHMENT_PREFIXES:
        return None
    if relative.is_absolute() or any(part in {"", ".", ".."} for part in relative.parts):
        raise ExportError(f"unsafe Telegram attachment path: {href}")
    return relative


def _attachment_inventory(export_root: Path, messages: list[TelegramMessage]) -> list[dict[str, Any]]:
    references: dict[str, set[str]] = {}
    for message in messages:
        for href in message.links:
            relative = _attachment_relative_path(href)
            if relative is not None:
                references.setdefault(relative.as_posix(), set()).add(message.message_id)
    inventory = []
    root = export_root.resolve()
    for relative_text, message_ids in sorted(references.items()):
        path = export_root / relative_text
        try:
            resolved = path.resolve(strict=True)
        except OSError as exc:
            raise ExportError(f"referenced Telegram attachment is missing: {relative_text}") from exc
        if root not in resolved.parents or not resolved.is_file() or resolved.is_symlink():
            raise ExportError(f"unsafe Telegram attachment target: {relative_text}")
        data = resolved.read_bytes()
        inventory.append(
            {
                "bytes": len(data),
                "message_ids": sorted(message_ids, key=int),
                "path": relative_text,
                "sha256": sha256_bytes(data),
            }
        )
    return inventory


def prepare(export_path: Path, output_dir: Path, max_segment_chars: int) -> dict[str, Any]:
    if not export_path.is_file() or export_path.is_symlink():
        raise ExportError("export path must be a regular, non-symlink file")
    source = export_path.read_bytes()
    source_sha256 = sha256_bytes(source)
    title, messages = parse_export(export_path)
    messages, secret_message_ids, secret_value_count = _redact_messages(messages)
    segments = _split_messages(messages, source_sha256, max_segment_chars)
    attachments = _attachment_inventory(export_path.parent, messages)
    if output_dir.exists():
        raise ExportError("output directory already exists; choose a new private path")
    output_dir.mkdir(mode=0o700, parents=True)
    output_dir.chmod(0o700)
    segment_rows = []
    offset = 0
    for index, rendered_messages in enumerate(segments, start=1):
        filename = f"segment-{index:04d}.txt"
        text = "\n".join(rendered_messages)
        data = text.encode("utf-8")
        path = output_dir / filename
        path.write_bytes(data)
        path.chmod(0o600)
        count = len(rendered_messages)
        selected = messages[offset : offset + count]
        offset += count
        segment_rows.append(
            {
                "bytes": len(data),
                "characters": len(text),
                "first_message_id": selected[0].message_id,
                "last_message_id": selected[-1].message_id,
                "message_count": count,
                "path": filename,
                "sha256": sha256_bytes(data),
                "source_identity": f"telegram-export:sha256:{source_sha256}#segment={index:04d}",
            }
        )
    corpus_binding = {
        "attachments": attachments,
        "segments": segment_rows,
        "source_html_sha256": source_sha256,
    }
    manifest = {
        "schema": SCHEMA,
        "parser_version": PARSER_VERSION,
        "source": {
            "bytes": len(source),
            "chat_title": title,
            "filename": export_path.name,
            "sha256": source_sha256,
        },
        "bounds": {"max_segment_characters": max_segment_chars},
        "counts": {
            "attachments": len(attachments),
            "messages": len(messages),
            "participants": len({message.sender for message in messages}),
            "segments": len(segment_rows),
        },
        "participants": sorted({message.sender for message in messages}),
        "time_range": {"first": messages[0].timestamp, "last": messages[-1].timestamp},
        "secret_redaction": {
            "message_count": len(secret_message_ids),
            "message_ids": secret_message_ids,
            "value_count": secret_value_count,
        },
        "segments": segment_rows,
        "attachments": attachments,
        "corpus_sha256": sha256_bytes(canonical_json_bytes(corpus_binding)),
        "safety": {
            "database_write_performed": False,
            "model_call_performed": False,
            "message_text_in_manifest": False,
            "output_files_private": True,
            "secret_values_removed_before_segments": True,
        },
    }
    manifest_path = output_dir / "corpus-manifest.json"
    manifest_path.write_text(json.dumps(manifest, ensure_ascii=False, indent=2, sort_keys=True) + "\n", encoding="utf-8")
    manifest_path.chmod(0o600)
    return manifest


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--export", required=True, type=Path)
    parser.add_argument("--output-dir", required=True, type=Path)
    parser.add_argument("--max-segment-chars", default=DEFAULT_MAX_SEGMENT_CHARS, type=int)
    return parser.parse_args()


def main() -> int:
    args = parse_args()
    try:
        manifest = prepare(args.export, args.output_dir, args.max_segment_chars)
    except ExportError as exc:
        print(json.dumps({"schema": SCHEMA, "status": "error", "error": str(exc)}, sort_keys=True))
        return 2
    print(
        json.dumps(
            {
                "schema": SCHEMA,
                "status": "prepared",
                "messages": manifest["counts"]["messages"],
                "segments": manifest["counts"]["segments"],
                "corpus_sha256": manifest["corpus_sha256"],
            },
            sort_keys=True,
        )
    )
    return 0


if __name__ == "__main__":
    raise SystemExit(main())