"""Centralized attachment registry for Mugi.

Maps short, model-friendly attachment IDs to absolute filesystem paths so
the LLM never needs to copy raw paths.  The registry is process-scoped
(in-memory dict), which matches the lifecycle of temp_uploads (purged on
startup and exit).

Usage:
    from attachment_registry import register, resolve, resolve_if_id, purge, list_for_thread
"""
from __future__ import annotations

import logging
import mimetypes
import os
import threading
import time
from dataclasses import dataclass, field
from typing import Optional

_log = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Schema
# ---------------------------------------------------------------------------

@dataclass(frozen=True, slots=True)
class AttachmentRecord:
    attachment_id: str
    abs_path: str
    original_name: str
    mime: str
    size_bytes: int
    thread_id: str
    source_route: str
    created_at: float = field(default_factory=time.time)


# Attachment IDs use this prefix so the normalizer can recognise them cheaply
# in free-text tool args without false-positives on regular paths.
ATTACHMENT_ID_PREFIX = "attachment://"

# ---------------------------------------------------------------------------
# Registry (process-scoped, thread-safe)
# ---------------------------------------------------------------------------

_lock = threading.Lock()
_store: dict[str, AttachmentRecord] = {}

# Telemetry counters (atomic via GIL, but guarded for clarity)
_counters: dict[str, int] = {
    "register": 0,
    "resolve_ok": 0,
    "resolve_fail": 0,
    "legacy_path_fallback": 0,
    "malformed_path_repair": 0,
}


def _next_short_id() -> str:
    """Return a short, monotonic, collision-free ID like ``att_1``, ``att_2``."""
    with _lock:
        seq = _counters["register"] + 1
        _counters["register"] = seq
    return f"att_{seq}"


# ---------------------------------------------------------------------------
# Public API
# ---------------------------------------------------------------------------

def register(
    abs_path: str,
    original_name: str,
    thread_id: str = "",
    source_route: str = "",
) -> AttachmentRecord:
    """Register a new attachment and return its record (including ``attachment_id``)."""
    aid = _next_short_id()
    mime = mimetypes.guess_type(original_name)[0] or "application/octet-stream"
    try:
        size = os.path.getsize(abs_path) if os.path.isfile(abs_path) else 0
    except OSError:
        size = 0
    rec = AttachmentRecord(
        attachment_id=aid,
        abs_path=abs_path,
        original_name=original_name,
        mime=mime,
        size_bytes=size,
        thread_id=thread_id,
        source_route=source_route,
    )
    with _lock:
        _store[aid] = rec
    _log.debug("attachment registered: %s -> %s (%s)", aid, abs_path, original_name)
    return rec


def resolve(attachment_id: str) -> Optional[str]:
    """Return the absolute path for *attachment_id*, or ``None`` if unknown.

    Accepts both bare IDs (``att_3``) and prefixed (``attachment://att_3``).
    """
    aid = attachment_id
    if aid.startswith(ATTACHMENT_ID_PREFIX):
        aid = aid[len(ATTACHMENT_ID_PREFIX):]
    with _lock:
        rec = _store.get(aid)
    if rec is not None:
        _bump("resolve_ok")
        return rec.abs_path
    _bump("resolve_fail")
    return None


def resolve_if_id(value: str) -> str:
    """If *value* looks like an attachment reference, resolve it.

    Returns the resolved absolute path on success, or the stripped *value*
    if it is not an attachment reference or cannot be resolved.
    This is the primary entry-point used by the dispatcher normalizer.
    """
    if not value:
        return value
    v = value.strip()
    if v.startswith(ATTACHMENT_ID_PREFIX):
        resolved = resolve(v)
        if resolved is not None:
            return resolved
        return v
    if v.startswith("att_") and v[4:].isdigit():
        resolved = resolve(v)
        if resolved is not None:
            return resolved
    # Always return the stripped form so leading/trailing whitespace from the
    # tokenizer never leaks into tool arguments, even for non-attachment values.
    return v


def get_record(attachment_id: str) -> Optional[AttachmentRecord]:
    """Return the full record, or ``None``."""
    aid = attachment_id
    if aid.startswith(ATTACHMENT_ID_PREFIX):
        aid = aid[len(ATTACHMENT_ID_PREFIX):]
    with _lock:
        return _store.get(aid)


def get_record_by_path(abs_path: str) -> Optional[AttachmentRecord]:
    """Return the most recently registered record matching *abs_path*, or ``None``.

    Compares using realpath + abspath so temp-upload paths match after symlink
    normalization (e.g. ``/tmp`` vs ``/private/tmp`` on macOS).
    """
    try:
        target = os.path.realpath(os.path.abspath(abs_path))
    except OSError:
        target = os.path.abspath(abs_path)
    with _lock:
        for rec in reversed(list(_store.values())):
            try:
                cand = os.path.realpath(os.path.abspath(rec.abs_path))
            except OSError:
                cand = rec.abs_path
            if cand == target:
                return rec
    return None


def list_for_thread(thread_id: str) -> list[AttachmentRecord]:
    """Return all attachments registered for *thread_id*."""
    with _lock:
        return [r for r in _store.values() if r.thread_id == thread_id]


def find_attachment_id(*, filename: str = "", url: str = "") -> Optional[str]:
    """Best-effort lookup when replaying stored ``file_attachment`` UI parts to the LLM.

    Returns the most recently registered ``attachment_id`` matching the attachment
    URL basename first (most specific), then *filename*, or ``None`` when the registry
    no longer holds the file.
    """
    fn = (filename or "").strip()
    url_base = os.path.basename((url or "").split("?", 1)[0]) if (url or "").strip() else ""
    if not fn and not url_base:
        return None
    with _lock:
        records = list(reversed(list(_store.values())))

    if url_base:
        for rec in records:
            if rec.original_name == url_base:
                return rec.attachment_id
            try:
                if os.path.basename(rec.abs_path) == url_base:
                    return rec.attachment_id
            except OSError:
                continue

    if fn:
        for rec in records:
            if rec.original_name == fn:
                return rec.attachment_id
    return None


def purge() -> int:
    """Clear the entire registry. Returns the number of records removed."""
    with _lock:
        n = len(_store)
        _store.clear()
    _log.debug("attachment registry purged (%d records)", n)
    return n


def get_counters() -> dict[str, int]:
    """Return a snapshot of telemetry counters."""
    with _lock:
        return dict(_counters)


def bump_counter(name: str) -> None:
    """Increment an external counter (used by the dispatcher normalizer)."""
    _bump(name)


def _bump(name: str) -> None:
    with _lock:
        _counters[name] = _counters.get(name, 0) + 1
