"""Export annotated frames as a plain YOLO dataset the user can keep.

This is deliberately independent of the master dataset. A batch can be exported
the moment it has annotations — before review, before merge — because the point
is that the user owns a copy of their own work and does not have to trust this
application to keep it.

Layout inside the zip:

    images/<batch>__<frame>.jpg
    labels/<batch>__<frame>.txt
    data.yaml
    classes.txt
    MANIFEST.json      what was exported, and what was skipped
"""

import json
import os
import shutil
import tempfile
import time
from typing import List, Optional

from backend import batches, config, db, projects, review


class ExportError(Exception):
    pass


def _label_line(class_id: int, geometry: dict, label_type: str) -> str:
    from backend.dataset import _label_line as build
    return build(class_id, geometry, label_type)


def annotated_frames(project_id: int, batch_ids: Optional[List[int]] = None,
                     approved_only: bool = False) -> list:
    where = "b.project_id = ?"
    args: list = [project_id]
    if batch_ids:
        where += f" AND b.id IN ({','.join('?' for _ in batch_ids)})"
        args += list(batch_ids)
    if approved_only:
        where += " AND f.review_status = 'approved'"
    with db.cursor() as cur:
        cur.execute(
            f"""SELECT f.id, f.filename, f.review_status, b.id AS batch_id,
                       b.batch_label, b.date_label
                FROM frames f JOIN batches b ON b.id = f.batch_id
                WHERE {where}
                ORDER BY b.id, f.idx""",
            args,
        )
        return [dict(row) for row in cur.fetchall()]


def build_zip(project_id: int, batch_ids: Optional[List[int]] = None,
              approved_only: bool = False, include_empty: bool = False) -> str:
    """Write the zip and return its path.

    `include_empty` decides whether frames with no annotation are exported as
    negative samples. Off by default: an empty label file is a claim that the
    image contains nothing, and for a frame that was simply never annotated that
    claim is false.
    """
    project = projects.get(project_id)
    if project is None:
        raise ExportError("No such project")

    frames = annotated_frames(project_id, batch_ids, approved_only)
    if not frames:
        raise ExportError("There is nothing to export — no frames match that selection")

    staging = tempfile.mkdtemp(prefix="export-")
    images_dir = os.path.join(staging, "images")
    labels_dir = os.path.join(staging, "labels")
    os.makedirs(images_dir)
    os.makedirs(labels_dir)

    written = 0
    skipped_empty = 0
    missing = 0
    shapes = 0
    per_batch: dict = {}

    for frame in frames:
        annotations = review.listing(frame["id"])
        if not annotations and not include_empty:
            skipped_empty += 1
            continue

        source = os.path.join(
            batches.frames_dir(project["slug"], frame["batch_id"]), frame["filename"])
        if not os.path.isfile(source):
            missing += 1
            continue

        stem = f"{frame['batch_id']}__{os.path.splitext(frame['filename'])[0]}"
        shutil.copyfile(source, os.path.join(images_dir, f"{stem}.jpg"))
        lines = [_label_line(item["class_id"], item["geometry"], project["label_type"])
                 for item in annotations]
        with open(os.path.join(labels_dir, f"{stem}.txt"), "w", encoding="utf-8") as handle:
            handle.write("\n".join(lines) + ("\n" if lines else ""))

        written += 1
        shapes += len(lines)
        label = f"{frame['date_label']}/{frame['batch_label']}"
        per_batch[label] = per_batch.get(label, 0) + 1

    if written == 0:
        shutil.rmtree(staging, ignore_errors=True)
        raise ExportError(
            "Every matching frame was empty or missing its image, so the export would "
            "contain no labels"
        )

    names = ", ".join(f"'{item['name']}'" for item in project["classes"])
    with open(os.path.join(staging, "data.yaml"), "w", encoding="utf-8") as handle:
        handle.write("path: .\ntrain: images\nval: images\n\n")
        handle.write(f"nc: {len(project['classes'])}\n")
        handle.write(f"names: [{names}]\n")

    with open(os.path.join(staging, "classes.txt"), "w", encoding="utf-8") as handle:
        handle.write("\n".join(item["name"] for item in project["classes"]) + "\n")

    manifest = {
        "project": project["name"],
        "label_type": project["label_type"],
        "exported_at": time.strftime("%Y-%m-%d %H:%M:%S"),
        "approved_only": approved_only,
        "include_empty": include_empty,
        "images": written,
        "shapes": shapes,
        "per_batch": per_batch,
        "skipped_without_annotations": skipped_empty,
        "skipped_image_missing": missing,
        "classes": [{"class_id": c["class_id"], "name": c["name"]} for c in project["classes"]],
    }
    with open(os.path.join(staging, "MANIFEST.json"), "w", encoding="utf-8") as handle:
        json.dump(manifest, handle, indent=2)

    project_root = config.project_dir(project["slug"])
    # Each export is ~200 MB. Without this, every click leaves another copy
    # behind and the project folder grows without bound.
    for stale in os.listdir(project_root):
        if stale.startswith("export-") and stale.endswith(".zip"):
            try:
                os.unlink(os.path.join(project_root, stale))
            except OSError:
                pass

    stamp = time.strftime("%Y%m%d-%H%M%S")
    path = shutil.make_archive(os.path.join(project_root, f"export-{stamp}"), "zip", staging)
    shutil.rmtree(staging, ignore_errors=True)
    return path


# ---- restore ------------------------------------------------------------

def _safe_members(archive) -> list:
    """Only plain files under images/ and labels/.

    A zip can name `../../etc/passwd`; extracting one blindly writes wherever it
    says. Every member is checked rather than trusted.
    """
    picked = []
    for member in archive.infolist():
        if member.is_dir():
            continue
        name = member.filename.replace("\\", "/")
        if name.startswith("/") or ".." in name.split("/"):
            continue
        if name.startswith(("images/", "labels/")) or name == "MANIFEST.json":
            picked.append((name, member))
    return picked


def _points_from_label(parts: List[str], label_type: str) -> Optional[dict]:
    values = [float(v) for v in parts]
    if label_type == "bbox":
        if len(values) != 4:
            return None
        cx, cy, w, h = values
        return {"type": "bbox",
                "points": [cx - w / 2, cy - h / 2, cx + w / 2, cy + h / 2]}
    if len(values) < 6 or len(values) % 2:
        return None
    return {"type": "polygon",
            "points": [[values[i], values[i + 1]] for i in range(0, len(values), 2)]}


def restore_zip(project_id: int, zip_path: str, batch_label: str = "") -> dict:
    """Load an exported zip back in as a fresh batch, ready to keep reviewing.

    The frames land in a new batch rather than being merged back into the ones
    they came from: the originals may still exist, and silently overwriting a
    batch the user is working in would destroy the very work this feature is
    meant to protect.
    """
    import zipfile

    from PIL import Image

    project = projects.get(project_id)
    if project is None:
        raise ExportError("No such project")

    by_name = {item["name"]: item["class_id"] for item in project["classes"]}
    stamp = time.strftime("%Y%m%d-%H%M%S")
    label = batch_label or f"restored-{stamp}"

    with zipfile.ZipFile(zip_path) as archive:
        members = _safe_members(archive)
        names = {name for name, _ in members}
        if not any(name.startswith("images/") for name in names):
            raise ExportError("This zip has no images/ folder — is it an export from this app?")

        manifest = {}
        if "MANIFEST.json" in names:
            manifest = json.loads(archive.read("MANIFEST.json"))
        source_type = manifest.get("label_type", project["label_type"])
        if source_type != project["label_type"]:
            raise ExportError(
                f"This export holds {source_type} labels but the project is "
                f"{project['label_type']} — importing it would produce wrong shapes"
            )

        # Classes come back by name, so an id that shifted since the export does
        # not silently relabel every shape.
        remap = {}
        for item in manifest.get("classes", []):
            if item["name"] in by_name:
                remap[item["class_id"]] = by_name[item["name"]]
            else:
                raise ExportError(
                    f"The export uses class '{item['name']}', which this project does not "
                    "have. Add the class first, then import."
                )

        with db.cursor() as cur:
            cur.execute(
                """INSERT INTO batches (project_id, video_path, date_label, batch_label,
                                        start_sec, end_sec, fps, status, created_at)
                   VALUES (?, '', 'restored', ?, 0, 0, 0, 'extracted', ?)""",
                (project_id, label, time.time()),
            )
            batch_id = cur.lastrowid

        target_dir = batches.frames_dir(project["slug"], batch_id)
        os.makedirs(target_dir, exist_ok=True)

        restored, shapes, skipped = 0, 0, 0
        image_members = sorted(n for n in names if n.startswith("images/"))
        for index, name in enumerate(image_members):
            stem = os.path.splitext(os.path.basename(name))[0]
            if not stem:
                continue
            filename = f"{stem}.jpg"
            destination = os.path.join(target_dir, filename)
            with archive.open(name) as source, open(destination, "wb") as handle:
                shutil.copyfileobj(source, handle)

            try:
                with Image.open(destination) as image:
                    width, height = image.size
            except Exception:
                os.unlink(destination)
                skipped += 1
                continue

            with db.cursor() as cur:
                cur.execute(
                    """INSERT INTO frames (batch_id, idx, filename, width, height, review_status)
                       VALUES (?, ?, ?, ?, ?, 'pending')""",
                    (batch_id, index, filename, width, height),
                )
                frame_id = cur.lastrowid

                label_name = f"labels/{stem}.txt"
                if label_name in names:
                    for line in archive.read(label_name).decode("utf-8").splitlines():
                        parts = line.split()
                        if len(parts) < 5:
                            continue
                        geometry = _points_from_label(parts[1:], project["label_type"])
                        if geometry is None:
                            continue
                        class_id = remap.get(int(parts[0]), int(parts[0]))
                        cur.execute(
                            """INSERT INTO annotations (frame_id, class_id, geometry, score,
                                                        source, created_at)
                               VALUES (?, ?, ?, 1.0, 'manual', ?)""",
                            (frame_id, class_id, json.dumps(geometry), time.time()),
                        )
                        shapes += 1
            restored += 1

    with db.cursor() as cur:
        cur.execute("UPDATE batches SET frame_count = ?, status = 'reviewing' WHERE id = ?",
                    (restored, batch_id))

    return {"batch_id": batch_id, "batch_label": label, "frames": restored,
            "shapes": shapes, "skipped": skipped}