import argparse
import sys
from pathlib import Path

ROOT_DIR = Path(__file__).resolve().parents[1]
if str(ROOT_DIR) not in sys.path:
    sys.path.insert(0, str(ROOT_DIR))

from core.stego import StegoManager


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Embed/extract 36-byte context payload in PNG."
    )
    parser.add_argument("image", type=Path, help="Path to PNG image")
    parser.add_argument(
        "--data",
        type=str,
        help="ASCII payload to embed (max 36 bytes)",
    )
    parser.add_argument(
        "--output",
        type=Path,
        default=Path("output_context.png"),
        help="Output PNG path for embedded image",
    )
    parser.add_argument(
        "--extract",
        action="store_true",
        help="Extract payload from image",
    )

    args = parser.parse_args()
    manager = StegoManager()

    if args.extract:
        result = manager.extract_data_from_png(args.image)
        if not result.get("success"):
            raise SystemExit(f"Extract failed: {result.get('error')}")
        data = result.get("data") or ""
        print(data)
        return

    if args.data is None:
        raise SystemExit("--data is required when not using --extract")

    payload = args.data.encode("ascii", errors="strict")
    if len(payload) > 36:
        raise SystemExit("Payload exceeds 36 bytes")

    output = manager.embed_data(args.image, payload, "context")
    if output != args.output:
        output.replace(args.output)
        output = args.output
    print(f"Wrote {output}")


if __name__ == "__main__":
    main()