import hashlib
import cv2
import tempfile
import os
import subprocess

# ----------------------------------------------------------------------
# Configuration
HASH_POS = (50, 50)
HASH_FONT = cv2.FONT_HERSHEY_SIMPLEX
HASH_SCALE = 0.6
HASH_COLOR = (0, 0, 255)      # red
HASH_THICKNESS = 2
RECT_PADDING = 5

def frame_to_bytes(frame):
    """Convert frame (numpy array) to bytes."""
    return frame.tobytes()

def compute_hash(frame):
    """Return SHA256 hex digest of the frame."""
    return hashlib.sha256(frame_to_bytes(frame)).hexdigest()

def draw_text_with_background(img, text):
    """Draw text with a solid white rectangle for OCR readability."""
    x, y = HASH_POS
    (tw, th), baseline = cv2.getTextSize(text, HASH_FONT, HASH_SCALE, HASH_THICKNESS)
    cv2.rectangle(img,
                  (x - RECT_PADDING, y - th - RECT_PADDING),
                  (x + tw + RECT_PADDING, y + baseline + RECT_PADDING),
                  (255, 255, 255), -1)
    cv2.putText(img, text, (x, y), HASH_FONT, HASH_SCALE, HASH_COLOR, HASH_THICKNESS)

# ----------------------------------------------------------------------
def encode_video(input_path, output_path):
    """Streaming encode: visible hash chain, audio copied later."""
    # 1. Create a temporary file for video‑only output
    with tempfile.NamedTemporaryFile(suffix='.mp4', delete=False) as tmp:
        temp_video = tmp.name

    # 2. Open input video with OpenCV.
    cap = cv2.VideoCapture(input_path)
    if not cap.isOpened():
        raise RuntimeError(f"Cannot open input video: {input_path}")

    frame_rate = cap.get(cv2.CAP_PROP_FPS)
    if not frame_rate or frame_rate <= 0:
        frame_rate = 30.0
    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))

    # 3. Open video writer for the temporary video-only file.
    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
    writer = cv2.VideoWriter(temp_video, fourcc, frame_rate, (width, height))
    if not writer.isOpened():
        cap.release()
        raise RuntimeError(f"Cannot open output video writer: {temp_video}")

    # 5. Process frames sequentially
    prev_frame = None
    frame_index = 0

    print("Processing frames (streaming)...")
    while True:
        ret, frame = cap.read()
        if not ret:
            break

        # Working copy
        out_frame = frame.copy()

        if prev_frame is not None:
            # Hash of previous frame (including its overlay)
            h = compute_hash(prev_frame)
            draw_text_with_background(out_frame, h)
        else:
            # Anchor for first frame
            anchor = hashlib.sha256(b"FIRST_FRAME_ANCHOR").hexdigest()
            draw_text_with_background(out_frame, anchor)

        # Write this frame to output video
        writer.write(out_frame)

        # Remember this frame (with its overlay) as "previous" for next frame
        prev_frame = out_frame
        frame_index += 1

        if frame_index % 100 == 0:
            print(f"  Processed {frame_index} frames...")

    cap.release()
    writer.release()

    print(f"Video processing complete. Total frames: {frame_index}")

    # 6. Copy audio from original to the temporary video (stream copy, no re‑encode)
    print("Copying audio stream...")
    subprocess.run([
        'ffmpeg',
        '-i', temp_video,
        '-i', input_path,
        '-c:v', 'copy',           # copy video from temp file (no re-encode)
        '-c:a', 'aac',            # re-encode audio
        '-map', '0:v:0',          # video from first input
        '-map', '1:a:0',          # audio from second input
        '-shortest',              # stop when the shortest stream ends
        '-y',                     # overwrite output
        output_path
    ], check=True)

    # Clean up temporary file
    os.unlink(temp_video)
    print(f"Final video with audio saved to {output_path}")

# ----------------------------------------------------------------------
def verify_video(input_path):
    """Verify the visible hash chain (streaming, chunk‑based)."""
    # For verification we still need OCR, which works best on a full frame.
    # We can stream frames but need to compare successive hashes.
    import pytesseract
    from PIL import Image

    reader = cv2.VideoCapture(input_path)
    if not reader.isOpened():
        raise RuntimeError(f"Cannot open input video: {input_path}")
    prev_frame = None
    frame_index = 0
    verified = True

    print("Verifying frames (streaming)...")
    while True:
        ret, current_frame = reader.read()
        if not ret:
            break

        if prev_frame is not None:
            # Extract visible hash from current_frame
            x, y = HASH_POS
            roi = current_frame[max(0, y-50):y+50, max(0, x-10):x+400]
            if roi.size == 0:
                print(f"Frame {frame_index}: cannot read hash region")
                verified = False
                continue

            gray = cv2.cvtColor(roi, cv2.COLOR_BGR2GRAY)
            _, thresh = cv2.threshold(gray, 150, 255, cv2.THRESH_BINARY_INV)
            pil_img = Image.fromarray(thresh)
            text = pytesseract.image_to_string(
                pil_img,
                config='--psm 7 -c tessedit_char_whitelist=0123456789abcdef'
            )
            displayed = ''.join(c for c in text if c in '0123456789abcdef')
            if len(displayed) >= 64:
                displayed = displayed[:64]
            else:
                print(f"Frame {frame_index}: extracted short hash: '{displayed}'")
                verified = False
                continue

            actual = compute_hash(prev_frame)
            if displayed != actual:
                print(f"Frame {frame_index}: MISMATCH!")
                print(f"  Displayed: {displayed}")
                print(f"  Actual:    {actual}")
                verified = False
        # else first frame – we could verify anchor, but skip for brevity
        prev_frame = current_frame
        frame_index += 1
        if frame_index % 100 == 0:
            print(f"  Verified {frame_index} frames...")

    reader.release()

    print(f"\nVerification {'PASSED' if verified else 'FAILED'} after {frame_index} frames")
    return verified

# ----------------------------------------------------------------------
if __name__ == "__main__":
    import sys
    if len(sys.argv) < 3:
        print("Usage:")
        print("  Encode: python tiger_video2.py encode input.mp4 output.mp4")
        print("  Verify: python tiger_video2.py verify encoded.mp4")
        sys.exit(1)

    mode = sys.argv[1]
    if mode == "encode":
        encode_video(sys.argv[2], sys.argv[3])
    elif mode == "verify":
        verify_video(sys.argv[2])
    else:
        print("Unknown mode.")