ui/parsing/parser.py

"""Parse ChatGPT and Claude exports (.json or .zip) into a normalized `ParsedExport`.

Mirrors the real export shapes:
- ChatGPT: list of conversations, each with a `mapping` node-tree (author.role, content.parts,
  create_time as unix epoch). We walk parent->first-child from the root to recover order.
- Claude: list of conversations, each with `chat_messages` (sender: human|assistant, text,
  created_at as ISO-8601), conversation `created_at`.

A .zip is the raw download from either provider; we locate the first `conversations.json`
(or any top-level .json that looks like an export).
"""
from __future__ import annotations

import datetime as _dt
import io
import json
import os
import re
import zipfile
from dataclasses import dataclass, field

from .language import is_english

_MONTHS = ["Jan", "Feb", "Mar", "Apr", "May", "Jun", "Jul", "Aug", "Sep", "Oct", "Nov", "Dec"]


class ParseError(Exception):
    """Raised when the input is not a recognizable ChatGPT/Claude export."""


@dataclass
class Turn:
    role: str            # "user" | "assistant"
    text: str
    epoch: float | None = None   # normalized unix seconds, when available


@dataclass
class Conversation:
    turns: list[Turn] = field(default_factory=list)


@dataclass
class ParsedExport:
    source: str                                  # "chatgpt" | "claude"
    conversations: list[Conversation] = field(default_factory=list)

    # ---- aggregate facts (used by the processing reveal) ----
    @property
    def conversation_count(self) -> int:
        return len(self.conversations)

    @property
    def all_turns(self) -> list[Turn]:
        return [t for c in self.conversations for t in c.turns]

    @property
    def turn_count(self) -> int:
        return len(self.all_turns)

    @property
    def user_turns(self) -> list[Turn]:
        return [t for t in self.all_turns if t.role == "user"]

    @property
    def english_turn_count(self) -> int:
        return sum(1 for t in self.all_turns if is_english(t.text))

    @property
    def other_turn_count(self) -> int:
        return self.turn_count - self.english_turn_count

    def date_range(self) -> str | None:
        epochs = [t.epoch for t in self.all_turns if t.epoch]
        if not epochs:
            return None
        lo, hi = min(epochs), max(epochs)
        return f"{_fmt_month(lo)} – {_fmt_month(hi)}"

    def busiest_slot(self) -> str | None:
        """Best-effort 'most active on <weekday> <period>' from timestamps, else None."""
        epochs = [t.epoch for t in self.all_turns if t.epoch]
        if len(epochs) < 3:
            return None
        from collections import Counter
        days, periods = Counter(), Counter()
        for e in epochs:
            d = _dt.datetime.fromtimestamp(e, _dt.timezone.utc)
            days[d.strftime("%A")] += 1
            periods[_period(d.hour)] += 1
        return f"Most active on {days.most_common(1)[0][0]} {periods.most_common(1)[0][0]}"


def _period(hour: int) -> str:
    if 5 <= hour < 12:
        return "mornings"
    if 12 <= hour < 17:
        return "afternoons"
    if 17 <= hour < 22:
        return "evenings"
    return "nights"


def _fmt_month(epoch: float) -> str:
    d = _dt.datetime.fromtimestamp(epoch, _dt.timezone.utc)
    return f"{_MONTHS[d.month - 1]} {d.year}"


def _to_epoch(value) -> float | None:
    """Accept unix epoch (int/float/str, seconds OR milliseconds) or ISO-8601 string. Returns a
    plausible unix-seconds float, or None for missing/garbage/out-of-range values."""
    if value is None:
        return None
    raw = None
    if isinstance(value, (int, float)):
        raw = float(value)
    else:
        s = str(value).strip()
        if not s:
            return None
        try:
            raw = float(s)
        except ValueError:
            try:
                return _sane(_dt.datetime.fromisoformat(s.replace("Z", "+00:00")).timestamp())
            except ValueError:
                return None
    if raw is not None and raw > 1e11:   # milliseconds (some new ChatGPT nodes) -> seconds
        raw /= 1000.0
    return _sane(raw)


def _sane(epoch: float | None) -> float | None:
    """Keep only plausible timestamps (2001-01-01 .. 2035-12-31); reject garbage so datetime can't throw."""
    if epoch is None or not (1_000_000_000 <= epoch <= 2_080_000_000):
        return None
    return epoch


def _norm_role(raw: str) -> str | None:
    r = (raw or "").lower()
    if r in ("user", "human"):
        return "user"
    if r in ("assistant", "model", "ai", "bot"):
        return "assistant"
    return None


# --------------------------------------------------------------------------------------
# Provider parsers
# --------------------------------------------------------------------------------------

def _parse_chatgpt(data: list) -> list[Conversation]:
    convs: list[Conversation] = []
    for conv in data:
        if not isinstance(conv, dict):
            continue
        mapping = conv.get("mapping")
        if not isinstance(mapping, dict):
            continue
        convs.append(Conversation(turns=_walk_chatgpt(mapping, conv.get("current_node"))))
    return convs


def _node_to_turn(node: dict) -> Turn | None:
    msg = node.get("message") or {}
    role = _norm_role((msg.get("author") or {}).get("role", ""))
    text = _text_from_parts(msg.get("content") or {})
    return Turn(role=role, text=text, epoch=_to_epoch(msg.get("create_time"))) if role and text else None


def _walk_chatgpt(mapping: dict, current_node: str | None = None) -> list[Turn]:
    """Order a ChatGPT conversation's mapping nodes. Old exports carry `children` (walk root→child[0]);
    new exports dropped `children` and keep only `parent` (walk the active leaf `current_node` up to root).
    Falls back to create_time order if neither path is available."""
    order: list[str] = []
    if any(n.get("children") for n in mapping.values()):  # old format: forward via children
        root = next((nid for nid, n in mapping.items() if not n.get("parent")), None) or next(iter(mapping), None)
        cur, seen = root, set()
        while cur and cur in mapping and cur not in seen:
            seen.add(cur); order.append(cur)
            kids = mapping[cur].get("children") or []
            cur = kids[0] if kids else None
    elif current_node and current_node in mapping:        # new format: backward via parent
        cur, seen = current_node, set()
        while cur and cur in mapping and cur not in seen:
            seen.add(cur); order.append(cur)
            cur = mapping[cur].get("parent")
        order.reverse()
    if not order:                                         # fallback: chronological by create_time
        order = sorted((nid for nid, n in mapping.items() if n.get("message")),
                       key=lambda nid: (mapping[nid]["message"].get("create_time") or 0))
    turns = []
    for nid in order:
        t = _node_to_turn(mapping[nid])
        if t:
            turns.append(t)
    return turns


def _text_from_parts(content: dict) -> str:
    if not isinstance(content, dict):
        return ""
    parts = content.get("parts") or []
    return "\n".join(p for p in parts if isinstance(p, str)).strip()


def _parse_claude(data: list) -> list[Conversation]:
    convs: list[Conversation] = []
    for conv in data:
        if not isinstance(conv, dict):
            continue
        msgs = conv.get("chat_messages")
        if not isinstance(msgs, list):
            continue
        turns = []
        for m in msgs:
            if not isinstance(m, dict):
                continue
            role = _norm_role(m.get("sender", ""))
            text = (m.get("text") or "").strip()
            if role and text:
                turns.append(Turn(role=role, text=text, epoch=_to_epoch(m.get("created_at"))))
        convs.append(Conversation(turns=turns))
    return convs


def _detect(data) -> str:
    sample = data[0] if isinstance(data, list) and data else data
    if isinstance(sample, dict):
        if "mapping" in sample:
            return "chatgpt"
        if "chat_messages" in sample:
            return "claude"
    raise ParseError("Unrecognized export — expected a ChatGPT or Claude conversations JSON.")


# --------------------------------------------------------------------------------------
# Public entry point
# --------------------------------------------------------------------------------------

def parse_export(path: str) -> ParsedExport:
    """Parse a ChatGPT/Claude export file (.json or .zip). Raises ParseError if unrecognized."""
    raw = _load_json(path)
    if not isinstance(raw, list):
        raise ParseError("Export must be a JSON list of conversations.")
    source = _detect(raw)
    convs = _parse_chatgpt(raw) if source == "chatgpt" else _parse_claude(raw)
    if not any(c.turns for c in convs):
        raise ParseError("No conversation turns found in this export.")
    return ParsedExport(source=source, conversations=convs)


def _load_json(path: str):
    if path is None or not os.path.exists(path):
        raise ParseError("No file provided.")
    if zipfile.is_zipfile(path):
        with zipfile.ZipFile(path) as zf:
            members = _conversation_members(zf.namelist())
            if not members:
                raise ParseError("Zip has no conversations.json (or conversations-NNN.json shards).")
            merged = []  # new ChatGPT exports SHARD conversations across conversations-000.json, -001, …
            for name in members:
                with zf.open(name) as fh:
                    part = json.load(io.TextIOWrapper(fh, encoding="utf-8"))
                merged.extend(part if isinstance(part, list) else [part])
            return merged
    try:
        with open(path, "r", encoding="utf-8") as fh:
            return json.load(fh)
    except (json.JSONDecodeError, UnicodeDecodeError) as e:
        raise ParseError(f"File is not valid JSON: {e}") from e


# export metadata files that are NOT conversation data (must not be mistaken for the export)
_META_JSON = {"user.json", "user_settings.json", "export_manifest.json", "shared_conversations.json",
              "conversation_asset_file_names.json", "library_files.json", "users.json", "memories.json"}


def _conversation_members(names: list[str]) -> list[str]:
    """The conversation JSON member(s) in an export zip: a single `conversations.json`, OR the
    `conversations-NNN.json` shards (new ChatGPT format), OR a best-effort non-metadata fallback."""
    base = lambda n: n.rsplit("/", 1)[-1]
    exact = [n for n in names if base(n) == "conversations.json"]
    if exact:
        return exact[:1]
    shards = sorted(n for n in names if re.match(r"conversations-\d+\.json$", base(n)))
    if shards:
        return shards
    other = [n for n in names if n.endswith(".json") and not n.startswith("__MACOSX")
             and "/projects/" not in n and base(n) not in _META_JSON]
    return other[:1]