import streamlit as st
from PIL import Image
import io
import time
from utils.aadhaar_ocr import extract_aadhaar_details
from utils.pan_ocr import extract_pan_details
from utils.face_match import match_faces

st.set_page_config(page_title="Smart KYC Verification", page_icon="🔍", layout="wide")

# Allowed file types and max file size (5MB)
ALLOWED_EXTENSIONS = ["jpg", "jpeg", "png"]
MAX_FILE_SIZE = 5 * 1024 * 1024  # 5MB

def show_temp_error(message, duration=5):
    """Show error message for a specified duration (in seconds) and then clear it."""
    placeholder = st.empty()
    placeholder.error(message)
    time.sleep(duration)
    placeholder.empty()

def validate_file(file):
    """
    Validate file extension and file size. Also attempts to open the image.
    Returns a tuple (image, error_message). If valid, error_message is None.
    """
    if not file:
        return None, None
    filename = file.name.lower()
    extension = filename.split('.')[-1]
    if extension not in ALLOWED_EXTENSIONS:
        return None, f"File extension .{extension} is not allowed. Allowed extensions: {ALLOWED_EXTENSIONS}"
    if file.size > MAX_FILE_SIZE:
        return None, f"File size {file.size} bytes exceeds the limit of {MAX_FILE_SIZE} bytes (5MB)."
    try:
        image = Image.open(io.BytesIO(file.getvalue()))
        image.verify()  # Check integrity
        image = Image.open(io.BytesIO(file.getvalue()))
        return image, None
    except Exception as e:
        return None, f"Error processing image: {e}"

# Display general instructions for the user.
st.markdown("## Welcome to Smart KYC Verification")
st.markdown("""
Please follow the steps below to complete your KYC process:
1. **Aadhaar Verification:** Upload your Aadhaar front and back images. Once uploaded, extract and verify the details.
2. **PAN Verification:** Upload your PAN card image and extract the details. Verify the information provided.
3. **Face Verification:** Upload your selfie, and we will match it with your Aadhaar or PAN details.

**Note:** Only valid image files (JPG/JPEG/PNG) up to 5MB are allowed. Follow the instructions and messages displayed on each step.
""")
st.markdown("---")

# Define the order of steps
steps = ["Aadhaar Verification", "PAN Verification", "Face Verification"]

# Initialize session state variables
if "current_step" not in st.session_state:
    st.session_state.current_step = steps[0]  # start at first step

if "aadhaar_data" not in st.session_state:
    st.session_state.aadhaar_data = None
if "pan_data" not in st.session_state:
    st.session_state.pan_data = None
if "face_match_result" not in st.session_state:
    st.session_state.face_match_result = None
if "aadhaar_verified" not in st.session_state:
    st.session_state.aadhaar_verified = False
if "pan_verified" not in st.session_state:
    st.session_state.pan_verified = False
if "selfie_uploaded" not in st.session_state:
    st.session_state.selfie_uploaded = False

# Initialize dynamic keys for file uploaders
if "aadhaar_front_key" not in st.session_state:
    st.session_state.aadhaar_front_key = 0
if "aadhaar_back_key" not in st.session_state:
    st.session_state.aadhaar_back_key = 0
if "pan_card_key" not in st.session_state:
    st.session_state.pan_card_key = 0
if "selfie_key" not in st.session_state:
    st.session_state.selfie_key = 0

# Also, store the uploaded files separately if needed
if "aadhaar_front_file" not in st.session_state:
    st.session_state.aadhaar_front_file = None
if "aadhaar_back_file" not in st.session_state:
    st.session_state.aadhaar_back_file = None
if "pan_card_file" not in st.session_state:
    st.session_state.pan_card_file = None

# Sidebar Navigation with current step pre-selected
st.sidebar.title("Smart KYC Verification")
page = st.sidebar.radio(
    "Select Step",
    steps,
    index=steps.index(st.session_state.current_step),
    disabled=True  # disable manual navigation to enforce flow
)

# Helper: redirect to a given step and rerun
def redirect_to_step(step):
    st.session_state.current_step = step
    st.rerun()

# ===================================
# Step 1: Aadhaar Verification
# ===================================
if st.session_state.current_step == "Aadhaar Verification":
    st.title("Step 1: Aadhaar Verification")
    
    if st.session_state.aadhaar_verified:
        st.success("Aadhaar Verification Completed!")
        with st.expander("Aadhaar Details"):
            st.json(st.session_state.aadhaar_data)
        if st.button("Proceed to PAN Verification"):
            redirect_to_step("PAN Verification")
    else:
        col1, col2 = st.columns(2)
        with col1:
            aadhaar_front_file = st.file_uploader(
                "Upload Aadhaar Front",
                type=ALLOWED_EXTENSIONS,
                key=f"aadhaar_front_{st.session_state.aadhaar_front_key}"
            )
        with col2:
            aadhaar_back_file = st.file_uploader(
                "Upload Aadhaar Back",
                type=ALLOWED_EXTENSIONS,
                key=f"aadhaar_back_{st.session_state.aadhaar_back_key}"
            )
        
        with col1:
            if aadhaar_front_file:
                image_front, error_front = validate_file(aadhaar_front_file)
                if error_front:
                    show_temp_error(f"Aadhaar Front: {error_front}")
                else:
                    st.session_state.aadhaar_front_file = aadhaar_front_file
                    st.image(image_front, caption="Aadhaar Front", width=300)
                
        with col2:
            if aadhaar_back_file:
                image_back, error_back = validate_file(aadhaar_back_file)
                if error_back:
                    show_temp_error(f"Aadhaar Back: {error_back}")
                else:
                    st.session_state.aadhaar_back_file = aadhaar_back_file
                    st.image(image_back, caption="Aadhaar Back", width=300)

        if aadhaar_front_file and aadhaar_back_file and st.session_state.aadhaar_data is None:
            if st.button("Extract Aadhaar Details"):
                with st.spinner("Extracting Aadhaar details..."):
                    try:
                        data = extract_aadhaar_details(aadhaar_front_file, aadhaar_back_file)
                        if "error" in data:
                            show_temp_error(f"Aadhaar OCR Error: {data['error']}")
                        else:
                            st.session_state.aadhaar_data = data
                            st.success("Aadhaar details extracted successfully!")
                    except Exception as e:
                        show_temp_error(f"Failed to process Aadhaar: {e}")
                        
        if st.session_state.aadhaar_data:
            with st.expander("Aadhaar Details"):
                st.json(st.session_state.aadhaar_data)
            col_btn = st.columns([1,7,1.21])
            with col_btn[0]:
                if st.button("Correct"):
                    st.session_state.aadhaar_verified = True
                    st.success("Aadhaar details verified!")
                    redirect_to_step("PAN Verification")
            with col_btn[2]:
                if st.button("Not Correct"):
                    st.session_state.aadhaar_data = None
                    st.session_state.aadhaar_front_file = None
                    st.session_state.aadhaar_back_file = None
                    st.session_state.aadhaar_verified = False
                    st.session_state.aadhaar_front_key += 1
                    st.session_state.aadhaar_back_key += 1
                    st.warning("Aadhaar data cleared. Please re-upload Aadhaar images.")
                    redirect_to_step("Aadhaar Verification")

# ===================================
# Step 2: PAN Verification (Accessible only after Aadhaar Verification)
# ===================================
elif st.session_state.current_step == "PAN Verification":
    if not st.session_state.aadhaar_verified:
        st.warning("Please complete Aadhaar verification first!")
        redirect_to_step("Aadhaar Verification")
        st.stop()

    st.title("Step 2: PAN Verification")
    
    if st.session_state.pan_verified:
        st.success("PAN Verification Completed!")
        with st.expander("PAN Details"):
            st.json(st.session_state.pan_data)
        if st.button("Proceed to Face Verification"):
            redirect_to_step("Face Verification")
    else:
        pan_card_file = st.file_uploader(
            "Upload PAN Card",
            type=ALLOWED_EXTENSIONS,
            key=f"pan_card_{st.session_state.pan_card_key}"
        )
        if pan_card_file:
            image_pan, error_pan = validate_file(pan_card_file)
            if error_pan:
                show_temp_error(f"PAN Card: {error_pan}")
            else:
                st.session_state.pan_card_file = pan_card_file
                st.image(image_pan, caption="PAN Card", width=300)
        if pan_card_file and st.session_state.pan_data is None:
            if st.button("Extract PAN Details"):
                with st.spinner("Extracting PAN details..."):
                    try:
                        data = extract_pan_details(pan_card_file)
                        if "error" in data:
                            show_temp_error(f"PAN OCR Error: {data['error']}")
                        else:
                            st.session_state.pan_data = data
                            st.success("PAN details extracted successfully!")
                    except Exception as e:
                        show_temp_error(f"Failed to process PAN: {e}")
        if st.session_state.pan_data:
            with st.expander("PAN Details"):
                st.json(st.session_state.pan_data)
            col_btn = st.columns([1,7,1.21])
            with col_btn[0]:
                if st.button("Correct"):
                    st.session_state.pan_verified = True
                    st.success("PAN details verified!")
                    redirect_to_step("Face Verification")
            with col_btn[2]:
                if st.button("Not Correct"):
                    st.session_state.pan_data = None
                    st.session_state.pan_card_file = None
                    st.session_state.pan_verified = False
                    st.session_state.pan_card_key += 1
                    st.warning("PAN data cleared. Please re-upload PAN image.")
                    redirect_to_step("PAN Verification")

# ===================================
# Step 3: Face Verification (Accessible only after PAN Verification)
# ===================================
elif st.session_state.current_step == "Face Verification":
    if not st.session_state.pan_verified:
        st.warning("Please complete PAN verification first!")
        redirect_to_step("PAN Verification")
        st.stop()

    st.title("Step 3: Face Verification")
    if st.session_state.face_match_result:
        st.info(st.session_state.face_match_result)
    selfie_file = st.file_uploader(
        "Upload Your Selfie",
        type=ALLOWED_EXTENSIONS,
        key=f"selfie_{st.session_state.selfie_key}"
    )
    if selfie_file:
        image_selfie, error_selfie = validate_file(selfie_file)
        if error_selfie:
            show_temp_error(f"Selfie: {error_selfie}")
        else:
            st.image(image_selfie, caption="Uploaded Selfie", width=300)
            st.session_state.selfie_uploaded = True

    if st.session_state.selfie_uploaded:
        if st.button("Match Faces"):
            if selfie_file and (st.session_state.aadhaar_data or st.session_state.pan_data):
                try:
                    aadhaar_match = match_faces(st.session_state.aadhaar_front_file, selfie_file) if st.session_state.aadhaar_data else False
                    pan_match = match_faces(st.session_state.pan_card_file, selfie_file) if st.session_state.pan_data else False

                    if aadhaar_match or pan_match:
                        st.session_state.face_match_result = "KYC Verified: Face matched successfully!"
                        st.success(st.session_state.face_match_result)
                        st.balloons()
                        st.info("Redirecting to main page in 5 seconds...")
                        # Auto-redirect to main page after 5 seconds using meta refresh
                        st.markdown("<meta http-equiv='refresh' content='5;url=/'/>", unsafe_allow_html=True)
                    else:
                        st.session_state.face_match_result = "KYC Failed: Face does not match. Please re-upload your selfie."
                        show_temp_error(st.session_state.face_match_result)
                        st.session_state.selfie_uploaded = False
                        st.session_state.selfie_key += 1
                        redirect_to_step("Face Verification")
                except Exception as e:
                    show_temp_error(f"Face matching failed: {e}")

st.sidebar.info("The current step is locked. Complete each step to proceed. Previous steps are frozen.")