import streamlit as st
import cv2
import os
import numpy as np
from datetime import datetime
import matplotlib.pyplot as plt
from services.detection_service import detect_faults_solar, detect_faults_windmill
from services.anomaly_service import track_anomalies, predict_anomaly
from models.solar_model import load_solar_model
from models.windmill_model import load_windmill_model
from config.settings import VIDEO_FOLDER
from PIL import Image
import io

# Custom CSS for styling to match the screenshot
st.markdown(
    """
    <style>
    .main-header {
        text-align: center;
        font-size: 24px;
        font-weight: bold;
        color: #333;
    }
    .status {
        text-align: center;
        font-size: 16px;
        color: #333;
        margin-bottom: 20px;
    }
    .section-title {
        font-size: 16px;
        font-weight: bold;
        color: #333;
        text-transform: uppercase;
        margin-bottom: 10px;
    }
    .section-box {
        border: 1px solid #4A90E2;
        padding: 10px;
        border-radius: 5px;
        margin-bottom: 20px;
    }
    .log-entry {
        font-size: 14px;
        color: #333;
        margin-bottom: 5px;
    }
    .metrics-text {
        font-size: 14px;
        color: #333;
        margin-bottom: 5px;
    }
    .snapshot-img {
        max-width: 100%;
        height: auto;
        margin-bottom: 10px;
    }
    </style>
    """,
    unsafe_allow_html=True
)

# Initialize session state
if 'logs' not in st.session_state:
    st.session_state.logs = []
if 'anomaly_counts' not in st.session_state:
    st.session_state.anomaly_counts = []
if 'frame_numbers' not in st.session_state:
    st.session_state.frame_numbers = []
if 'total_detected' not in st.session_state:
    st.session_state.total_detected = 0
if 'snapshots' not in st.session_state:
    st.session_state.snapshots = []

# Create snapshots directory if it doesn't exist
SNAPSHOT_FOLDER = "./snapshots"
if not os.path.exists(SNAPSHOT_FOLDER):
    os.makedirs(SNAPSHOT_FOLDER)

def main():
    # Header
    st.markdown('<div class="main-header">THERMAL ANOMALY MONITORING DASHBOARD</div>', unsafe_allow_html=True)
    st.markdown('<div class="status">🟢 RUNNING</div>', unsafe_allow_html=True)

    # Sidebar for video selection and detection type
    st.sidebar.header("Settings")
    video_files = [f for f in os.listdir(VIDEO_FOLDER) if f.endswith('.mp4')]
    if not video_files:
        st.error("No videos found in the 'data' folder. Please add .mp4 files.")
        return
    video_file = st.sidebar.selectbox("Select Video", video_files)
    detection_type = st.sidebar.selectbox("Detection Type", ["Solar Panel", "Windmill"])

    # Load the appropriate model
    model = load_solar_model() if detection_type == "Solar Panel" else load_windmill_model()

    # Layout: Two columns for video feed and metrics
    col1, col2 = st.columns([3, 1])

    with col1:
        st.markdown('<div class="section-title">LIVE VIDEO FEED</div>', unsafe_allow_html=True)
        with st.container():
            st.markdown('<div class="section-box">', unsafe_allow_html=True)
            video_placeholder = st.empty()
            st.markdown('</div>', unsafe_allow_html=True)

    with col2:
        st.markdown('<div class="section-title">LIVE METRICS</div>', unsafe_allow_html=True)
        with st.container():
            st.markdown('<div class="section-box">', unsafe_allow_html=True)
            metrics_placeholder = st.empty()
            prediction_placeholder = st.empty()
            st.markdown('</div>', unsafe_allow_html=True)

    # Layout: Two columns for logs and trends
    col3, col4 = st.columns([1, 2])

    with col3:
        # Live Logs
        st.markdown('<div class="section-title">LIVE LOGS</div>', unsafe_allow_html=True)
        with st.container():
            st.markdown('<div class="section-box">', unsafe_allow_html=True)
            logs_placeholder = st.empty()
            st.markdown('</div>', unsafe_allow_html=True)

        # Last 5 Captured Events
        st.markdown('<div class="section-title">LAST 5 CAPTURED EVENTS</div>', unsafe_allow_html=True)
        with st.container():
            st.markdown('<div class="section-box">', unsafe_allow_html=True)
            events_placeholder = st.empty()
            st.markdown('</div>', unsafe_allow_html=True)

    with col4:
        st.markdown('<div class="section-title">DETECTION TRENDS</div>', unsafe_allow_html=True)
        with st.container():
            st.markdown('<div class="section-box">', unsafe_allow_html=True)
            st.markdown('<div style="font-size: 14px; font-weight: bold; margin-bottom: 10px;">Anomalies Over Time</div>', unsafe_allow_html=True)
            trends_placeholder = st.empty()
            st.markdown('</div>', unsafe_allow_html=True)

    # Process video
    if video_file:
        video_path = os.path.join(VIDEO_FOLDER, video_file)
        cap = cv2.VideoCapture(video_path)
        if not cap.isOpened():
            st.error("Error: Could not open video file.")
            return

        frame_count = 0
        while cap.isOpened():
            ret, frame = cap.read()
            if not ret:
                break

            frame_count += 1
            frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)

            # Detect faults
            faults = detect_faults_solar(model, frame_rgb) if detection_type == "Solar Panel" else detect_faults_windmill(model, frame_rgb)
            num_anomalies = len(faults)

            # Draw bounding boxes and labels
            annotated_frame = frame_rgb.copy()
            for fault in faults:
                x, y = int(fault['location'][0]), int(fault['location'][1])
                cv2.rectangle(annotated_frame, (x-30, y-30), (x+30, y+30), (255, 0, 0), 2)
                cv2.putText(annotated_frame, f"{fault['type']}", (x, y-40), 
                            cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 0, 0), 2)

            # Save snapshot if faults are detected
            if faults:
                timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
                fault_types = "_".join([fault['type'].replace(" ", "_") for fault in faults])
                snapshot_filename = f"snapshot_{timestamp}_frame_{frame_count}_{fault_types}.png"
                snapshot_path = os.path.join(SNAPSHOT_FOLDER, snapshot_filename)
                cv2.imwrite(snapshot_path, cv2.cvtColor(annotated_frame, cv2.COLOR_RGB2BGR))
                st.session_state.snapshots.append({
                    "path": snapshot_path,
                    "log": f"{timestamp} - Frame {frame_count} - Anomalies: {num_anomalies} ({fault_types})"
                })

            # Update video feed with timestamp
            timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
            video_placeholder.image(annotated_frame, channels="RGB", caption=f"{timestamp}")

            # Update logs and metrics
            log_entry = f"{timestamp} - Frame {frame_count} - Anomalies: {num_anomalies}"
            st.session_state.logs.append(log_entry)
            st.session_state.total_detected += num_anomalies
            st.session_state.anomaly_counts.append(num_anomalies)
            st.session_state.frame_numbers.append(frame_count)

            # Keep only the last 100 frames for trends and last 5 snapshots
            if len(st.session_state.frame_numbers) > 100:
                st.session_state.frame_numbers.pop(0)
                st.session_state.anomaly_counts.pop(0)
            if len(st.session_state.snapshots) > 5:
                st.session_state.snapshots.pop(0)

            # Update live logs
            with logs_placeholder.container():
                for log in st.session_state.logs[::-1]:
                    st.markdown(f'<div class="log-entry">{log}</div>', unsafe_allow_html=True)

            # Update last 5 captured events with snapshots
            with events_placeholder.container():
                for snapshot in st.session_state.snapshots[::-1]:
                    st.markdown(f'<div class="log-entry">{snapshot["log"]}</div>', unsafe_allow_html=True)
                    st.image(snapshot["path"], caption="Fault Snapshot", use_column_width=True)

            # Update live metrics
            metrics_placeholder.markdown(
                f'<div class="metrics-text"><b>anomalies:</b> {num_anomalies}</div>'
                f'<div class="metrics-text"><b>total_detected:</b> {st.session_state.total_detected}</div>',
                unsafe_allow_html=True
            )

            # Predictive anomaly detection
            prediction = predict_anomaly(st.session_state.anomaly_counts)
            if prediction:
                prediction_placeholder.warning("**Prediction:** Potential issue detected - anomaly spike detected!")
            else:
                prediction_placeholder.empty()

            # Update trends graph
            fig, ax = plt.subplots(figsize=(6, 3))
            ax.plot(st.session_state.frame_numbers, st.session_state.anomaly_counts, marker='o', color='blue')
            ax.set_xlabel("Frame", fontsize=10)
            ax.set_ylabel("Count", fontsize=10)
            ax.grid(True)
            ax.tick_params(axis='both', which='major', labelsize=8)
            trends_placeholder.pyplot(fig)
            plt.close(fig)

        cap.release()
        st.success("Video processing completed.")

if __name__ == "__main__":
    main()