import torch
import numpy as np
import cv2
import tempfile, base64


def readb64(uri):
    encoded_data = uri.split(',')[-1]
    nparr = np.frombuffer(base64.b64decode(encoded_data), np.uint8)
    img = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
    return img

def img2base64(img, extension="jpg"):
    _, img_encoded = cv2.imencode(f".{extension}", img)
    img_base64 = base64.b64encode(img_encoded)
    img_base64 = img_base64.decode('utf-8')
    return img_base64

def binary2video(video_binary):
    # byte_arr = BytesIO()
    # byte_arr.write(video_binary)

    temp_ = tempfile.NamedTemporaryFile(suffix='.mp4')
    # decoded_string = base64.b64decode(video_binary)

    temp_.write(video_binary)
    video_capture = cv2.VideoCapture(temp_.name)
    ret, frame = video_capture.read()
    return video_capture

def extract_frames(data_path, interval=30, max_frames=50):
    """Method to extract frames"""
    cap = cv2.VideoCapture(data_path)
    frame_num = 0
    frames = list()

    while cap.isOpened():
        success, image = cap.read()
        if not success:
            break
        # image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        # image = torch.tensor(image) - torch.tensor([104, 117, 123])
        if frame_num % interval == 0:
            frames.append(image)
        frame_num += 1
        if len(frames) > max_frames:
            break
    cap.release()
    # if len(frames) > max_frames:
    #     samples = np.random.choice(
    #         np.arange(0, len(frames)), size=max_frames, replace=False)
    #     return [frames[_] for _ in samples]
    return frames

"""FilePicker for streamlit. 
Still doesn't seem to be a good solution for a way to select files to process from the server Streamlit is running on.
Here's a pretty functional solution. 
Usage:
```
import streamlit as st
from filepicker import st_file_selector
tif_file = st_file_selector(st, key = 'tif', label = 'Choose tif file')
```
"""

import os
import streamlit as st

def update_dir(key):
    choice = st.session_state[key]
    if os.path.isdir(os.path.join(st.session_state[key+'curr_dir'], choice)):
        st.session_state[key+'curr_dir'] = os.path.normpath(os.path.join(st.session_state[key+'curr_dir'], choice))
        files = sorted(os.listdir(st.session_state[key+'curr_dir']))
        if "images" in files:
          files.remove("images")
        st.session_state[key+'files'] = files

def st_file_selector(st_placeholder, path='.', label='Select a file/folder', key = 'selected'):
    if key+'curr_dir' not in st.session_state:
        base_path = '.' if path is None or path == '' else path
        base_path = base_path if os.path.isdir(base_path) else os.path.dirname(base_path)
        base_path = '.' if base_path is None or base_path == '' else base_path

        files = sorted(os.listdir(base_path))
        files.insert(0, 'Choose a file...')
        if "images" in files:
          files.remove("images")
        st.session_state[key+'files'] = files
        st.session_state[key+'curr_dir'] = base_path
    else:
        base_path = st.session_state[key+'curr_dir']

    selected_file = st_placeholder.selectbox(label=label, 
                                        options=st.session_state[key+'files'], 
                                        key=key, 
                                        on_change = lambda: update_dir(key))
    
    if selected_file == "Choose a file...":
        return None

    return selected_file