#!/usr/bin/env python3
"""
Shared engine for importing coding-agent artifacts into the Réplica brain.

Each agent (Kimi, Hermes, Claude Code, Antigravity, Qwen Code, MiMo Code) has an
adapter in `agent_importers/` that yields normalized `NormalizedSession` objects.
This module writes them to the unified brain layout used by the Aurélio extension
and the Theia backend:

  <replica>/.aurelio/brain/session-<source>-<id>/
    session.jsonl                       # session_start / session_end events
    session_memory.md                   # human summary + provenance
    .system_generated/chat_history.json # full normalized message list
    .system_generated/logs/overview.txt # one-line header

The writer is idempotent: it only touches a file when its content would change,
and derives stable timestamps from the source artifact so re-imports are no-ops.
"""

from __future__ import annotations

import hashlib
import json
import re
from dataclasses import dataclass, field
from datetime import datetime, timedelta, timezone
from pathlib import Path
from typing import Any, Iterable, Optional

# engine.py lives at <repo>/scripts/agent_importers/engine.py. The repo root is
# the nearest ancestor that contains a `.aurelio` directory; fall back to the
# three-levels-up heuristic only if that search fails.
def _resolve_repo_root() -> Path:
    here = Path(__file__).resolve()
    for candidate in (here.parent, *here.parents):
        if (candidate / ".aurelio").is_dir():
            return candidate
    return here.parents[2]


REPLICA_ROOT = _resolve_repo_root()
BRAIN_DIR = REPLICA_ROOT / ".aurelio" / "brain"


# ---------------------------------------------------------------------------
# Normalized shapes
# ---------------------------------------------------------------------------

@dataclass
class NormalizedMessage:
    role: str                 # user | assistant | tool_call | tool_result | system
    content: str
    timestamp: Optional[str] = None
    model: Optional[str] = None
    id: Optional[str] = None


@dataclass
class NormalizedSession:
    source: str               # kimi | hermes | claude-code | antigravity | qwen-code | mimocode
    session_id: str           # stable id within the source
    title: str
    created: str              # ISO-8601 Z
    modified: str             # ISO-8601 Z
    workspace_path: Optional[str]
    messages: list[NormalizedMessage]
    source_locator: str       # human-readable pointer back to the original artifact
    model_type: Optional[str] = None
    extra: dict[str, Any] = field(default_factory=dict)

    @property
    def brain_id(self) -> str:
        return f"session-{self.source}-{self.session_id}"


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def iso_now() -> str:
    return datetime.now(timezone.utc).isoformat().replace("+00:00", "Z")


def to_iso_z(ts: Any) -> Optional[str]:
    """Best-effort normalization of many timestamp shapes to ISO-8601 Z."""
    if ts is None:
        return None
    if isinstance(ts, (int, float)):
        # Heuristic: >1e12 is milliseconds
        seconds = ts / 1000.0 if ts > 1e12 else float(ts)
        try:
            return datetime.fromtimestamp(seconds, tz=timezone.utc).isoformat().replace("+00:00", "Z")
        except (OSError, ValueError):
            return None
    if isinstance(ts, str):
        s = ts.strip()
        if not s:
            return None
        # Already ISO-ish
        try:
            s2 = s.replace("Z", "+00:00")
            dt = datetime.fromisoformat(s2)
            if dt.tzinfo is None:
                dt = dt.replace(tzinfo=timezone.utc)
            return dt.astimezone(timezone.utc).isoformat().replace("+00:00", "Z")
        except ValueError:
            return None
    return None


def flatten_content(content: Any) -> str:
    """Flatten OpenAI/Anthropic/Kimi style content (str | list[part]) to text."""
    if content is None:
        return ""
    if isinstance(content, str):
        return content
    if isinstance(content, list):
        parts: list[str] = []
        for part in content:
            if isinstance(part, dict):
                ptype = part.get("type")
                if ptype in ("text", "input_text", "output_text"):
                    parts.append(part.get("text", ""))
                elif ptype == "tool_use":
                    parts.append(f"[tool_use: {part.get('name','?')}({json.dumps(part.get('input', {}), ensure_ascii=False)[:300]})]")
                elif ptype == "tool_result":
                    tr = part.get("content", "")
                    if isinstance(tr, list):
                        tr = " ".join(x.get("text", "") for x in tr if isinstance(x, dict))
                    parts.append(f"[tool_result] {str(tr)[:500]}")
                elif ptype == "image":
                    parts.append("[image]")
                else:
                    txt = part.get("text") or part.get("content")
                    parts.append(str(txt) if txt is not None else f"[{ptype or 'media'}]")
            else:
                parts.append(str(part))
        return "\n".join(p for p in parts if p)
    return str(content)


def write_if_changed(path: Path, content: str) -> bool:
    path.parent.mkdir(parents=True, exist_ok=True)
    new = content.encode("utf-8")
    if path.exists() and path.read_bytes() == new:
        return False
    path.write_bytes(new)
    return True


_TAG_RE = re.compile(r"<[^>]+>")
_BLOCK_TAG_RE = re.compile(r"<(system|current_focus|note)>.*?</\1>", re.S)
FILE_REF_RE = re.compile(
    r"`([^`]+?\.(?:py|js|ts|tsx|jsx|c|cpp|h|hpp|md|yaml|yml|json|toml|rs|go|java|kt|swift|dart))`",
    re.I,
)
CODE_FENCE_RE = re.compile(r"```(\w+)")


def derive_title(messages: list[NormalizedMessage], workspace_path: Optional[str], fallback: str) -> str:
    for m in messages:
        if m.role != "user":
            continue
        text = _BLOCK_TAG_RE.sub("", m.content).strip()
        text = _TAG_RE.sub("", text).strip()
        if text:
            title = text.split("\n", 1)[0][:80]
            if len(title) == 80:
                title += "..."
            return title
    if workspace_path:
        return f"{fallback} in {Path(workspace_path).name}"
    return fallback


def stable_per_message_timestamps(messages: list[NormalizedMessage], base_iso: str) -> list[str]:
    """Assign deterministic timestamps to messages that lack them by offsetting
    from the session base by index. Messages that already carry a timestamp keep it."""
    try:
        base_dt = datetime.fromisoformat(base_iso.replace("Z", "+00:00"))
    except ValueError:
        base_dt = datetime.now(timezone.utc)
    out: list[str] = []
    cursor = base_dt
    for m in messages:
        if m.timestamp:
            iso = to_iso_z(m.timestamp) or base_iso
            out.append(iso)
            try:
                cursor = datetime.fromisoformat(iso.replace("Z", "+00:00"))
            except ValueError:
                pass
        else:
            cursor = cursor + timedelta(seconds=1)
            out.append(cursor.isoformat().replace("+00:00", "Z"))
    return out


# ---------------------------------------------------------------------------
# Rendering
# ---------------------------------------------------------------------------

def render_session_memory(sess: NormalizedSession, import_ts: str) -> str:
    title = sess.title or derive_title(sess.messages, sess.workspace_path, f"{sess.source} session")
    user_msgs = [m for m in sess.messages if m.role == "user"]
    asst_msgs = [m for m in sess.messages if m.role == "assistant"]
    tool_msgs = [m for m in sess.messages if m.role in ("tool_call", "tool_result")]

    code_refs: set[str] = set()
    file_refs: set[str] = set()
    for m in sess.messages:
        code_refs.update(CODE_FENCE_RE.findall(m.content))
        file_refs.update(FILE_REF_RE.findall(m.content))

    last_user = (user_msgs[-1].content[:300].replace("\n", " ") if user_msgs else "_(none)_")
    last_asst = (asst_msgs[-1].content[:300].replace("\n", " ") if asst_msgs else "_(none)_")

    lines = [
        "# Session Memory",
        "",
        f"_Imported from {sess.source} agent artifacts._",
        "",
        "---",
        "",
        "## Session Goal",
        title,
        "",
        "## Workspace",
        f"- **Path**: `{sess.workspace_path or 'unknown'}`",
        f"- **{sess.source} session ID**: `{sess.session_id}`",
        "",
        "## Key Facts",
        f"- **Total messages**: {len(sess.messages)}",
        f"- **User messages**: {len(user_msgs)}",
        f"- **Assistant messages**: {len(asst_msgs)}",
        f"- **Tool messages**: {len(tool_msgs)}",
    ]
    if sess.model_type:
        lines.append(f"- **Model**: `{sess.model_type}`")

    if code_refs:
        lines += ["", "## Languages / Technologies Mentioned", ""]
        lines += [f"- `{l}`" for l in sorted(code_refs)]
    if file_refs:
        lines += ["", "## Files Referenced", ""]
        lines += [f"- `{f}`" for f in sorted(file_refs)[:30]]

    lines += [
        "",
        "## Last Exchange",
        "",
        f"**User**: {last_user}",
        "",
        f"**Assistant**: {last_asst}",
        "",
        "## Source",
        f"- Original artifact: `{sess.source_locator}`",
        f"- Imported at: {import_ts}",
        "",
    ]
    return "\n".join(lines)


def write_session(sess: NormalizedSession, dry_run: bool = False) -> dict[str, Any]:
    result: dict[str, Any] = {
        "source": sess.source,
        "session_id": sess.session_id,
        "brain_id": sess.brain_id,
        "written": [],
        "skipped_empty": False,
    }
    if not sess.messages:
        result["skipped_empty"] = True
        return result

    import_ts = sess.modified or sess.created or iso_now()
    timestamps = stable_per_message_timestamps(sess.messages, sess.created or import_ts)

    chat_messages = []
    for idx, (m, ts) in enumerate(zip(sess.messages, timestamps)):
        entry: dict[str, Any] = {
            "id": m.id or f"m{idx:06d}",
            "role": m.role,
            "content": m.content,
            "timestamp": ts,
        }
        if m.model:
            entry["model"] = m.model
        if sess.model_type:
            entry["modelType"] = sess.model_type
        chat_messages.append(entry)

    title = sess.title or derive_title(sess.messages, sess.workspace_path, f"{sess.source} session")
    target_dir = BRAIN_DIR / sess.brain_id

    chat_history = {
        "id": sess.brain_id,
        "title": title,
        "created": sess.created,
        "modified": sess.modified,
        "source": sess.source,
        "workspacePath": sess.workspace_path,
        "messages": chat_messages,
        "modelType": sess.model_type or sess.source,
        "sourceLocator": sess.source_locator,
        **({"extra": sess.extra} if sess.extra else {}),
    }

    session_jsonl_lines = [
        json.dumps({
            "type": "session_start",
            "timestamp": sess.created,
            "data": {
                "sessionId": sess.brain_id,
                "workspacePath": sess.workspace_path,
                "source": sess.source,
                "sourceSessionId": sess.session_id,
            },
        }, ensure_ascii=False),
    ]
    # Full trajectory: one event per message, preserving tool_call/tool_result
    # sequence for RL training. This is the primary training data artifact.
    for idx, (m, ts) in enumerate(zip(sess.messages, timestamps)):
        event: dict[str, Any] = {
            "type": "message",
            "timestamp": ts,
            "data": {
                "id": m.id or f"m{idx:06d}",
                "role": m.role,
                "content": m.content,
            },
        }
        if m.model:
            event["data"]["model"] = m.model
        session_jsonl_lines.append(json.dumps(event, ensure_ascii=False))
    session_jsonl_lines.append(
        json.dumps({
            "type": "session_end",
            "timestamp": sess.modified,
            "data": {"sessionId": sess.brain_id, "totalEvents": len(sess.messages)},
        }, ensure_ascii=False),
    )

    # RL trajectory: compact (state, action, reward) triples derived from the
    # message sequence. Each tool_call is an action; the preceding user/assistant
    # message is the state; the tool_result is the immediate observation.
    trajectory_lines: list[str] = []
    for idx, m in enumerate(sess.messages):
        if m.role != "tool_call":
            continue
        # Find the nearest preceding user/assistant message as state
        state = ""
        for j in range(idx - 1, -1, -1):
            if sess.messages[j].role in ("user", "assistant"):
                state = sess.messages[j].content[:2000]
                break
        # Find the tool_result that follows
        observation = ""
        if idx + 1 < len(sess.messages) and sess.messages[idx + 1].role == "tool_result":
            observation = sess.messages[idx + 1].content[:2000]
        traj_event = {
            "type": "trajectory_step",
            "session_id": sess.brain_id,
            "step": idx,
            "timestamp": timestamps[idx] if idx < len(timestamps) else None,
            "state": state,
            "action": m.content[:2000],
            "observation": observation,
            "model": m.model or sess.model_type,
        }
        trajectory_lines.append(json.dumps(traj_event, ensure_ascii=False))

    overview_txt = (
        f"# Conversation: {title}\n"
        f"# Created: {sess.created}\n"
        f"# ID: {sess.brain_id}\n"
        f"# Source: {sess.source}\n"
        f"# Workspace: {sess.workspace_path or 'unknown'}\n"
        f"# Original: {sess.source_locator}\n\n"
    )

    files = {
        target_dir / "session.jsonl": "\n".join(session_jsonl_lines) + "\n",
        target_dir / "session_memory.md": render_session_memory(sess, import_ts),
        target_dir / ".system_generated" / "chat_history.json": json.dumps(chat_history, indent=2, ensure_ascii=False) + "\n",
        target_dir / ".system_generated" / "logs" / "overview.txt": overview_txt,
    }
    if trajectory_lines:
        files[target_dir / ".system_generated" / "trajectory.jsonl"] = "\n".join(trajectory_lines) + "\n"

    if dry_run:
        result["would_write"] = [str(p.relative_to(REPLICA_ROOT)) for p in files]
        return result

    for path, content in files.items():
        if write_if_changed(path, content):
            result["written"].append(str(path.relative_to(REPLICA_ROOT)))
    return result


def source_id_for(label: str, raw: str) -> str:
    """Deterministic, filesystem-safe id suffix derived from an arbitrary raw id."""
    safe = re.sub(r"[^A-Za-z0-9._-]", "-", raw).strip("-")
    if len(safe) <= 80 and safe:
        return safe
    return hashlib.sha1(raw.encode("utf-8")).hexdigest()[:16]