"""Recover the existing project snapshot and its original server outputs.

Never overwrites an existing target database or regenerates media.
"""
from concurrent.futures import ThreadPoolExecutor
import hashlib
import json
import os
from pathlib import Path
import sqlite3

import requests

ROOT = Path(__file__).resolve().parent.parent
TARGET = ROOT / "huobao-data"
SOURCE = ROOT / "integration-work/backups/huobao-before-multivideo-20261106.sqlite3"


def main():
    TARGET.mkdir(parents=True, exist_ok=True)
    dest = TARGET / "huobao.sqlite3"
    if dest.exists():
        raise RuntimeError("Recovery target already exists; inspect it instead of overwriting")
    empty = Path(os.environ["APPDATA"]) / "HuobaoDrama/data/huobao.sqlite3"
    snapshot = ROOT / "integration-work/backups/huobao-empty-profile-20261106.sqlite3"
    if empty.exists() and not snapshot.exists():
        with sqlite3.connect(empty.as_uri() + "?mode=ro", uri=True) as a, sqlite3.connect(snapshot) as b:
            a.backup(b)
    with sqlite3.connect(SOURCE.as_uri() + "?mode=ro", uri=True) as a, sqlite3.connect(dest) as b:
        a.backup(b)
    with sqlite3.connect(dest) as conn:
        assert conn.execute("pragma integrity_check").fetchone()[0] == "ok"
        jobs = conn.execute("select id,task_id,local_path from sys_task where type='image' and status='completed' and local_path is not null").fetchall()
    client = json.loads((ROOT / "model-bridge/client.json").read_text("utf-8"))

    def download(job):
        task_id, remote_id, relative = job
        target = (TARGET / relative.lstrip("/")).resolve()
        target.relative_to(TARGET.resolve())
        session = requests.Session()
        session.trust_env = False
        response = session.get(client["url"] + "/v1/images/task/" + remote_id,
                               headers={"Authorization": "Bearer " + client["token"]}, timeout=30)
        if response.status_code != 200:
            return {"task_id": task_id, "error": "Status query failed", "status": response.status_code}
        payload = response.json()
        if payload.get("status") != "completed":
            return {"task_id": task_id, "error": "Original image is unavailable"}
        url = payload["image_url"]
        if not url.startswith(client["url"] + "/media/"):
            raise RuntimeError("Unexpected media destination")
        with session.get(url, stream=True, timeout=120) as media:
            if media.status_code != 200:
                return {"task_id": task_id, "error": "Download failed", "status": media.status_code}
            target.parent.mkdir(parents=True, exist_ok=True)
            partial = target.with_name(target.name + ".partial")
            digest = hashlib.sha256()
            count = 0
            with partial.open("wb") as out:
                for data in media.iter_content(1024 * 1024):
                    out.write(data)
                    digest.update(data)
                    count += len(data)
            partial.replace(target)
        return {"task_id": task_id, "path": relative, "bytes": count, "sha256": digest.hexdigest()}

    with ThreadPoolExecutor(max_workers=3) as pool:
        outcomes = list(pool.map(download, jobs))
    report = {"source": str(SOURCE), "target": str(TARGET), "images": outcomes}
    (ROOT / "model-bridge/data-recovery-20261106.json").write_text(json.dumps(report, indent=2, ensure_ascii=False), encoding="utf-8")
    failed = [r for r in outcomes if "error" in r]
    print("Recovered image files:", len(outcomes) - len(failed), "/", len(outcomes))
    if failed:
        print("Unrecovered task IDs:", [r["task_id"] for r in failed])
        raise RuntimeError("Some images need additional recovery")
    print("Project database integrity: OK")


if __name__ == "__main__":
    main()