from __future__ import annotations

import sys
from pathlib import Path


ROOT = Path(__file__).resolve().parent
sys.path.insert(0, str(ROOT / "src"))

from egg_damage.config import load_config
from egg_damage.gradio_app import build_app


EXPECTED_PARTS = {
    "mobilenet_v3.pt": 2,
    "xception.pt": 80,
}


def restore_chunked_model(filename: str) -> None:
    target = ROOT / "models" / filename
    parts_dir = ROOT / "models" / f"{filename}.parts"
    if target.exists() or not parts_dir.exists():
        return
    part_paths = sorted(parts_dir.glob("*.part"))
    if not part_paths:
        return
    expected = EXPECTED_PARTS.get(filename)
    if expected is not None and len(part_paths) < expected:
        return
    with target.open("wb") as output:
        for part_path in part_paths:
            output.write(part_path.read_bytes())


restore_chunked_model("mobilenet_v3.pt")
restore_chunked_model("xception.pt")
config = load_config(ROOT / "configs" / "space.yaml")
demo = build_app(config)


if __name__ == "__main__":
    demo.launch()