# transport.py — minio artifact client, shared by all workers.
# workers are stateless: pull inputs by uri to local disk, push outputs back, return uris.
# no sqlite here (see plan phase 2b option A: homesrv orchestrator owns state).
# uri form: "s3://<bucket>/<key...>" or bare "<bucket>/<key...>". first path segment = bucket.
import os
import json
import time
import logging
from minio import Minio
from minio.error import S3Error
from starlette.responses import Response

_client = None

# --- artifact layout ---------------------------------------------------------------------------
# one bucket per artifact class (`decisions/storage-layout.md#bucket-per-artifact`). Every worker
# formats its output uri from these, so moving a class between buckets is one edit here rather than
# a grep across five workers. `name` is the panel id, or 'p' when a worker has none.
PANEL_URI = "s3://panels/{manga_id}/{chapter_id}/panels/p{idx:03d}.png"
PAGE_PANEL_URI = "s3://panels/{manga_id}/{chapter_id}/panels/pg{page_index:03d}_p{idx:02d}.png"
AUDIO_URI = "s3://audio/{manga_id}/{chapter_id}/audio/{name}.wav"
AUDIO_FLAT_URI = "s3://audio/_audio/{name}.wav"
LAYER_URI = "s3://layers/{manga_id}/{chapter_id}/layers/{name}/{idx}.png"
CLIP_URI = "s3://video/{manga_id}/{chapter_id}/clips/{name}.mp4"
CHAPTER_URI = "s3://video/{manga_id}/{chapter_id}/chapter.mp4"
CHAR_PNG_URI = "s3://manga/{key}.png"
CHAR_NPY_URI = "s3://manga/{key}.npy"


def ids_from_uri(uri: str):
    """(manga_id, chapter_id) from any artifact uri: <bucket>/<manga_id>/<chapter_id>/...
    the orchestrator passes no ids to tts, layers or render, but every input uri encodes them."""
    parts = (uri.removeprefix("s3://")).split("/")
    if len(parts) < 3:
        raise ValueError(f"uri carries no manga/chapter: {uri!r}")
    return parts[1], parts[2]


def _summarize(body: bytes, limit=6) -> str:
    """compact one-line view of a json body for observability: uri inputs/outputs (basename, or
    key×N for lists) + short scalars; big lists/dicts (embeddings, panel arrays) shown as key[N].
    this is what makes the log answer 'what did the stage get / return / where did it go'."""
    try:
        obj = json.loads(body)
    except Exception:
        return "-"
    if not isinstance(obj, dict):
        return f"[{len(obj)}]" if isinstance(obj, list) else "-"
    parts = []
    for k, v in obj.items():
        if k == "panel_id":
            continue
        if isinstance(v, str) and (k.endswith(("uri", "url"))):
            parts.append(f"{k}={v.rsplit('/', 1)[-1]}")
        elif isinstance(v, list) and v and isinstance(v[0], str) and (k.endswith(("uris", "urls"))):
            parts.append(f"{k}×{len(v)}")
        elif isinstance(v, (int, float, bool)):
            parts.append(f"{k}={v}")
        elif isinstance(v, str):
            parts.append(f"{k}={v[:24]!r}" if len(v) > 24 else f"{k}={v!r}")
        elif isinstance(v, (list, dict)):
            parts.append(f"{k}[{len(v)}]")
        if len(parts) >= limit:
            parts.append("…")
            break
    return " ".join(parts) or "{}"


def install_logging(app, name: str):
    """one log line per request for any worker:
    `name METHOD /path panel=<id> in:{...} -> out:{...} 123ms 200` (or `... ERR: <exc>` on failure).
    call once right after app = FastAPI(). in/out summarize the json bodies (uris + short scalars,
    big arrays as key[N]); unparseable/empty bodies show `-`."""
    logging.basicConfig(
        level=logging.INFO,
        format="%(asctime)s %(levelname)s %(message)s",
        force=False,
    )
    log = logging.getLogger(name)

    @app.middleware("http")
    async def _log(request, call_next):
        body = await request.body()  # buffered; downstream re-reads from the cached bytes
        panel = "-"
        if body:
            try:
                panel = str(json.loads(body).get("panel_id") or "-")
            except Exception:
                pass
        req_in = _summarize(body) if body else "-"
        t0 = time.perf_counter()
        try:
            resp = await call_next(request)
        except Exception as e:
            ms = int((time.perf_counter() - t0) * 1000)
            log.error("%s %s panel=%s in:{%s} %dms ERR: %r",
                      request.method, request.url.path, panel, req_in, ms, e)
            raise
        # drain the response stream so we can summarize outputs, then hand back an identical response.
        out = "-"
        chunks = [c async for c in resp.body_iterator]
        raw = b"".join(chunks)
        if resp.headers.get("content-type", "").startswith("application/json"):
            out = _summarize(raw)
        resp = Response(content=raw, status_code=resp.status_code,
                        headers=dict(resp.headers), media_type=resp.media_type)
        ms = int((time.perf_counter() - t0) * 1000)
        lvl = log.error if resp.status_code >= 500 else log.info
        lvl("%s %s panel=%s in:{%s} -> out:{%s} %dms %d",
            request.method, request.url.path, panel, req_in, out, ms, resp.status_code)
        return resp


def _mc():
    global _client
    if _client is None:
        _client = Minio(
            os.environ.get("MINIO_ENDPOINT", "192.168.1.104:9000"),  # homesrv
            access_key=os.environ.get("MINIO_ACCESS_KEY", "admin"),
            secret_key=os.environ.get("MINIO_SECRET_KEY", "godforgiveus"),
            secure=False,
        )
    return _client


def _split(uri: str):
    """(bucket, key) from an s3-style or bare uri."""
    u = uri.removeprefix("s3://")
    bucket, _, key = u.partition("/")
    if not bucket or not key:
        raise ValueError(f"bad uri: {uri!r}")
    return bucket, key


_known_buckets = set()


def _ensure_bucket(c, bucket):
    # ponytail: process-local cache, buckets are never deleted at runtime
    if bucket in _known_buckets:
        return
    if not c.bucket_exists(bucket):
        c.make_bucket(bucket)
    _known_buckets.add(bucket)


def put(local_path: str, uri: str, client=None) -> str:
    """upload local file to minio at uri, return the s3:// uri. raises on failure."""
    c = client or _mc()
    bucket, key = _split(uri)
    _ensure_bucket(c, bucket)
    c.fput_object(bucket, key, local_path)
    return f"s3://{bucket}/{key}"


def get(uri: str, local_path: str, client=None) -> str:
    """download uri to local_path, return local_path. workers pull inputs to /dev/shm."""
    c = client or _mc()
    bucket, key = _split(uri)
    os.makedirs(os.path.dirname(local_path) or ".", exist_ok=True)
    c.fget_object(bucket, key, local_path)
    return local_path


def exists(uri: str, client=None) -> bool:
    """HEAD check for worker-level skip-if-exists (orchestrator is the real skip authority)."""
    c = client or _mc()
    bucket, key = _split(uri)
    try:
        c.stat_object(bucket, key)
        return True
    except S3Error:
        return False


def put_bytes(data: bytes, uri: str, client=None) -> str:
    """for generated artifacts that never touch disk (embeddings, json)."""
    import io
    c = client or _mc()
    bucket, key = _split(uri)
    _ensure_bucket(c, bucket)
    c.put_object(bucket, key, io.BytesIO(data), length=len(data))
    return f"s3://{bucket}/{key}"


if __name__ == "__main__":
    # self-check: fake in-memory minio, round-trip a file + bytes, exists true/false.
    import tempfile

    class _Fake:
        def __init__(self):
            self.store = {}
        def bucket_exists(self, b):
            return any(k[0] == b for k in self.store)
        def make_bucket(self, b):
            pass
        def fput_object(self, b, k, path):
            with open(path, "rb") as f:
                self.store[(b, k)] = f.read()
        def fget_object(self, b, k, path):
            with open(path, "wb") as f:
                f.write(self.store[(b, k)])
        def put_object(self, b, k, stream, length):
            self.store[(b, k)] = stream.read()
        def stat_object(self, b, k):
            if (b, k) not in self.store:
                raise S3Error("NoSuchKey", "missing", "", "", "", None)
            return True

    fake = _Fake()
    assert _split("s3://manga/a/b.png") == ("manga", "a/b.png")
    assert _split("manga/a/b.png") == ("manga", "a/b.png")

    src = tempfile.NamedTemporaryFile(delete=False)
    src.write(b"hello panel"); src.close()
    uri = put(src.name, "s3://manga/x/p001.png", client=fake)
    assert uri == "s3://manga/x/p001.png", uri
    assert exists(uri, client=fake)
    assert not exists("s3://manga/x/nope.png", client=fake)
    dst = src.name + ".out"
    get(uri, dst, client=fake)
    assert open(dst, "rb").read() == b"hello panel"

    put_bytes(b"\x01\x02", "s3://manga/x/e.npy", client=fake)
    assert exists("s3://manga/x/e.npy", client=fake)
    os.remove(src.name); os.remove(dst)

    # _summarize: uris -> basename, uri lists -> key×N, scalars kept, big arrays -> key[N], panel_id dropped
    s = _summarize(json.dumps({
        "panel_id": "p1", "audio_uri": "s3://manga/x/cl.wav",
        "panel_uris": ["s3://a/1.png", "s3://a/2.png"], "active": 2, "rtl": True,
        "weights": [0.1, 0.2, 0.7], "narration_text": "a very long line of narration here indeed",
    }).encode())
    assert "audio_uri=cl.wav" in s and "panel_uris×2" in s and "active=2" in s, s
    assert "weights[3]" in s and "panel_id" not in s, s
    assert _summarize(b"clip_uri", ) == "-"          # non-json
    assert _summarize(json.dumps({"clip_uri": "s3://m/c/clips/p1.mp4", "duration": 4.1}).encode()) \
        == "clip_uri=p1.mp4 duration=4.1"

    # artifact layout: templates format to the keys the workers wrote by hand before, and
    # ids_from_uri recovers the ids the orchestrator never sends.
    panel = PANEL_URI.format(manga_id="m1", chapter_id="c1", idx=7)
    assert panel == "s3://panels/m1/c1/panels/p007.png", panel
    assert PAGE_PANEL_URI.format(manga_id="m1", chapter_id="c1", page_index=2, idx=3) \
        == "s3://panels/m1/c1/panels/pg002_p03.png"
    assert CLIP_URI.format(manga_id="m1", chapter_id="c1", name="p003") \
        == "s3://video/m1/c1/clips/p003.mp4"
    assert LAYER_URI.format(manga_id="m1", chapter_id="c1", name="p003", idx=0) \
        == "s3://layers/m1/c1/layers/p003/0.png"
    assert ids_from_uri(panel) == ("m1", "c1")
    assert ids_from_uri("panels/m1/c1/panels/p007.png") == ("m1", "c1")
    try:
        ids_from_uri("s3://panels/p007.png")
        raise AssertionError("a uri with no chapter segment must raise")
    except ValueError:
        pass
    print("transport self-check ok")