"""Move legacy screenshot BLOBs to server files.

Run from the project root with::

    python backend/scripts/migrate_screenshots_to_files.py

Each row is committed only after the written file matches the source byte
length and SHA-256 digest. Re-running is safe: rows with a valid file_path and
an empty image_data value are skipped.
"""

import hashlib
import sys
from pathlib import Path


PROJECT_ROOT = Path(__file__).resolve().parents[2]
BACKEND_ROOT = PROJECT_ROOT / "backend"
if str(BACKEND_ROOT) not in sys.path:
    sys.path.insert(0, str(BACKEND_ROOT))

from sqlalchemy import inspect, select, text  # noqa: E402

from app.config import settings  # noqa: E402
from app.database import SessionLocal, engine  # noqa: E402
from app.models import AnswerScreenshot, AskAssignmentBatchScreenshot, AskAssignmentScreenshot  # noqa: E402
from app.screenshot_storage import file_digest, store_screenshot  # noqa: E402


TARGETS = (
    (AnswerScreenshot, "answer_screenshots", "answers"),
    (AskAssignmentBatchScreenshot, "ask_assignment_batch_screenshots", "ask-batches"),
    (AskAssignmentScreenshot, "ask_assignment_screenshots", "ask-assignments"),
)


def ensure_columns() -> None:
    for _model, table, _category in TARGETS:
        columns = {column["name"]: column for column in inspect(engine).get_columns(table)}
        if "file_path" not in columns:
            with engine.begin() as connection:
                connection.execute(text(f"ALTER TABLE {table} ADD COLUMN file_path VARCHAR(1000) NULL"))
        if engine.dialect.name == "mysql" and not columns["image_data"].get("nullable", False):
            with engine.begin() as connection:
                connection.execute(text(f"ALTER TABLE {table} MODIFY COLUMN image_data MEDIUMBLOB NULL"))


def migrate() -> None:
    ensure_columns()
    settings.screenshot_storage_dir.mkdir(parents=True, exist_ok=True)
    migrated_rows = 0
    migrated_bytes = 0
    for model, table, category in TARGETS:
        table_rows = 0
        table_bytes = 0
        while True:
            with SessionLocal() as db:
                image = db.scalar(
                    select(model)
                    .where(model.image_data.is_not(None))
                    .order_by(model.id)
                    .limit(1)
                )
                if image is None:
                    break
                source = bytes(image.image_data)
                source_digest = hashlib.sha256(source).hexdigest()
                relative_path = store_screenshot(source, image.filename, image.content_type, category)
                stored_size, stored_digest = file_digest(relative_path)
                if stored_size != len(source) or stored_digest != source_digest:
                    raise RuntimeError(f"{table}#{image.id} 文件校验失败")
                image.file_path = relative_path
                image.image_data = None
                db.commit()
                table_rows += 1
                table_bytes += len(source)
                if table_rows % 100 == 0:
                    print(f"{table}: migrated {table_rows} rows", flush=True)
        migrated_rows += table_rows
        migrated_bytes += table_bytes
        print(f"{table}: complete rows={table_rows} bytes={table_bytes}", flush=True)
    print(f"TOTAL rows={migrated_rows} bytes={migrated_bytes}", flush=True)
    print(f"FILES {settings.screenshot_storage_dir.resolve()}", flush=True)


if __name__ == "__main__":
    migrate()