from __future__ import annotations

import hashlib
from pathlib import Path
from typing import Any

import nbformat

from ocr_pipeline.utils.logging import get_logger

logger = get_logger(__name__)


def parse_notebook_for_images(notebook_path: Path) -> list[Path]:
    image_paths = []
    try:
        with open(notebook_path) as f:
            nb = nbformat.read(f, as_version=4)

        for cell in nb.cells:
            if cell.cell_type == "code":
                for output in cell.outputs:
                    if output.output_type == "stream" and output.name == "stdout":
                        text = output.text
                        for line in text.split("\n"):
                            line = line.strip()
                            if line.startswith("/") and any(
                                line.endswith(ext)
                                for ext in (".png", ".jpg", ".jpeg", ".tiff", ".bmp")
                            ):
                                path = Path(line)
                                if path.exists():
                                    image_paths.append(path)
                            elif "SCR-" in line and any(
                                line.endswith(ext) for ext in (".png", ".jpg", ".jpeg")
                            ):
                                for part in line.split():
                                    if part.startswith("SCR-") and any(
                                        part.endswith(ext) for ext in (".png", ".jpg", ".jpeg")
                                    ):
                                        path = Path("/Users/Aman/Pictures") / part
                                        if path.exists():
                                            image_paths.append(path)

    except Exception as e:
        logger.warning("notebook_parse_failed", path=str(notebook_path), error=str(e))

    return list(set(image_paths))


def compute_file_hash(path: Path) -> str:
    hasher = hashlib.sha256()
    with open(path, "rb") as f:
        for chunk in iter(lambda: f.read(8192), b""):
            hasher.update(chunk)
    return hasher.hexdigest()


def migrate_notebook(notebook_path: Path, pipeline: Any) -> dict[str, Any]:
    logger.info("starting_notebook_migration", notebook=str(notebook_path))

    image_paths = parse_notebook_for_images(notebook_path)
    logger.info("found_images_in_notebook", count=len(image_paths))

    results = {
        "total": len(image_paths),
        "processed": 0,
        "skipped": 0,
        "failed": 0,
        "errors": [],
    }

    for img_path in image_paths:
        try:
            img_hash = compute_file_hash(img_path)
            existing = pipeline.db.get_processed(img_hash)
            if existing:
                logger.info("skipping_already_processed", path=str(img_path), hash=img_hash[:16])
                results["skipped"] += 1
                continue

            logger.info("migrating_image", path=str(img_path))
            pipeline.process_single(img_path)
            results["processed"] += 1

        except Exception as e:
            logger.error("migration_failed", path=str(img_path), error=str(e))
            results["failed"] += 1
            results["errors"].append({"path": str(img_path), "error": str(e)})

    logger.info("migration_complete", **results)
    return results
