import os
import traceback

# ==========================================
# 🚨 ANTI-DEADLOCK CPU LIMITERS 🚨
os.environ["OMP_NUM_THREADS"] = "1"
os.environ["OPENBLAS_NUM_THREADS"] = "1"
os.environ["MKL_NUM_THREADS"] = "1"
os.environ["VECLIB_MAXIMUM_THREADS"] = "1"
os.environ["NUMEXPR_NUM_THREADS"] = "1"
# ==========================================

import cv2
import base64
import torch
import torch.nn as nn
import numpy as np
import pandas as pd
import joblib
import smtplib
import ssl
from email.message import EmailMessage
import threading
import uuid
import time
import requests
import datetime
from urllib.parse import urlparse
from flask import Flask, request, render_template, jsonify, Response, send_from_directory
from werkzeug.utils import secure_filename
from werkzeug.middleware.proxy_fix import ProxyFix

# Force PyTorch to use 1 thread safely
torch.set_num_threads(1)
try:
    torch.set_num_interop_threads(1)
except:
    pass

# --- Auto-install missing libraries ---
try:
    from ultralytics import YOLO
    import easyocr
    import yt_dlp
except ModuleNotFoundError:
    import sys
    import subprocess
    print("Installing missing libraries... This might take a minute...", flush=True)
    subprocess.check_call([sys.executable, "-m", "pip", "install", "ultralytics", "easyocr", "yt-dlp"])
    from ultralytics import YOLO
    import easyocr
    import yt_dlp

from torchvision.models.video import r3d_18

print("--- BOOT SEQUENCE INITIATED ---", flush=True)

app = Flask(__name__)
# FIX: Tells Flask it is running behind a Hugging Face proxy so it generates the correct URL paths
app.wsgi_app = ProxyFix(app.wsgi_app, x_for=1, x_proto=1, x_host=1, x_prefix=1)

ALERT_EMAIL_SENDER = "gowreeshgowri50@gmail.com" 
ALERT_EMAIL_PASSWORD = "oynu ulet pynk xsza".replace(" ", "") 
ALERT_EMAIL_RECEIVER = "mcblackdevil12342@gmail.com"
ENABLE_EMAIL_ALERTS = True

UPLOAD_FOLDER = 'uploads'
MODEL_FOLDER = 'models'

app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER
os.makedirs(UPLOAD_FOLDER, exist_ok=True)

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
classes = ['major', 'minor', 'moderate']
active_streams = {}

# --- Model Loading ---
try:
    model_yolo = YOLO(os.path.join(MODEL_FOLDER, 'yolov8_accident_model.pt'))
    model_tracker = YOLO('yolov8n.pt') 

    model_3d = r3d_18()
    model_3d.fc = nn.Linear(model_3d.fc.in_features, 3)
    model_3d.load_state_dict(torch.load(os.path.join(MODEL_FOLDER, '3dcnn_accident_model.pth'), map_location=device))
    model_3d.to(device).eval()

    model_svm = joblib.load(os.path.join(MODEL_FOLDER, 'ensemble_svm_model.pkl'))
    ocr_reader = easyocr.Reader(['en'], gpu=torch.cuda.is_available())
    models_loaded = True
    print("✅ All AI Engines Loaded.", flush=True)
except Exception as e:
    print(f"⚠️ Model Load Error: {e}", flush=True)
    models_loaded = False

try:
    model_traffic = joblib.load(os.path.join(MODEL_FOLDER, 'traffic_predictor.pkl'))
    print("✅ Traffic Risk Predictor Loaded.", flush=True)
except:
    model_traffic = None

# --- Helper Functions ---

def get_geo_info(ip_or_url=None):
    try:
        target = ""
        if ip_or_url:
            parsed = urlparse(ip_or_url)
            target = parsed.netloc.split(':')[0] if parsed.netloc else ip_or_url
        
        res = requests.get(f"http://ip-api.com/json/{target}", timeout=5).json()
        if res.get("status") == "success":
            city = res.get("city", "Unknown")
            country = res.get("countryCode", "UN")
            lat, lon = res.get("lat"), res.get("lon")
            location = f"{city}, {country}"
            
            wx = requests.get(f"https://api.open-meteo.com/v1/forecast?latitude={lat}&longitude={lon}&current_weather=true", timeout=5).json()
            temp = wx["current_weather"]["temperature"]
            return location, f"{temp}°C, Active"
    except: pass
    return "Surveillance Zone", "Standard Conditions"

def process_accident_async(stream_id, frame, location, confidence, severity):
    detected_plates = []
    try:
        ocr_results = ocr_reader.readtext(frame)
        for (_, text, _) in ocr_results:
            if len(text) > 4: detected_plates.append({"text": text.upper()})
    except: pass

    if stream_id in active_streams:
        active_streams[stream_id]["plates"] = detected_plates

    if ENABLE_EMAIL_ALERTS:
        try:
            _, buffer = cv2.imencode('.jpg', frame)
            msg = EmailMessage()
            msg['Subject'] = f"🚨 ALERT: {severity.upper()} ACCIDENT DETECTED"
            msg['From'] = ALERT_EMAIL_SENDER
            msg['To'] = ALERT_EMAIL_RECEIVER
            msg.set_content(f"Incident Report\nLocation: {location}\nSeverity: {severity}\nConfidence: {confidence}%\nTime: {datetime.datetime.now()}\n\nSystem has locked this stream for investigation.")
            
            msg.add_attachment(buffer.tobytes(), maintype='image', subtype='jpeg', filename='incident.jpg')
            
            # --- DUAL-PORT FIREWALL BYPASS ---
            try:
                # 1. Try Port 465 (Strict SSL - Works best for local VS Code)
                context = ssl.create_default_context()
                with smtplib.SMTP_SSL('smtp.gmail.com', 465, context=context, timeout=10) as server:
                    server.login(ALERT_EMAIL_SENDER, ALERT_EMAIL_PASSWORD)
                    server.send_message(msg)
                print(f"📧 Dispatch Email Sent (Port 465) for Stream {stream_id}", flush=True)
            except Exception as e_ssl:
                print(f"⚠️ Port 465 blocked by firewall. Falling back to Port 587... ({e_ssl})", flush=True)
                
                # 2. Try Port 587 (STARTTLS - Works best for Cloud environments like Hugging Face)
                server = smtplib.SMTP('smtp.gmail.com', 587, timeout=10)
                server.ehlo()
                server.starttls()
                server.login(ALERT_EMAIL_SENDER, ALERT_EMAIL_PASSWORD)
                server.send_message(msg)
                server.quit()
                print(f"📧 Dispatch Email Sent (Port 587) for Stream {stream_id}", flush=True)
                
        except Exception as e:
            print(f"📧 CRITICAL: Email Failed. Hugging Face might be strictly blocking all outbound SMTP connections: {str(e)}", flush=True)

def run_temporal_analysis(stream_id, frames_3d_copy, prob_max, prob_mean, prob_min, yolo_max_conf, raw_frame, location):
    state = active_streams.get(stream_id)
    if not state or state.get("final_decision"): return

    try:
        tensor_3d = torch.tensor(np.array(frames_3d_copy), dtype=torch.float32).permute(3, 0, 1, 2).unsqueeze(0).to(device) / 255.0
        with torch.no_grad():
            out_3d = model_3d(tensor_3d)
            prob_3d = torch.nn.functional.softmax(out_3d, dim=1).cpu().numpy()[0]
            
        combined = np.concatenate((prob_mean, prob_max, prob_min, prob_3d)).reshape(1, -1)
        
        try:
            ensemble_probs = model_svm.predict_proba(combined)[0]
            final_idx = model_svm.predict(combined)[0]
            confidence = float(ensemble_probs[final_idx] * 100)
            severity = classes[final_idx]
        except:
            final_idx = min(int(np.argmax(prob_max)), 2)
            confidence = float(np.max(prob_max)) * 100
            severity = classes[final_idx]

        state["severity"] = severity.capitalize()
        state["confidence"] = round(confidence, 1)
        state["cnn"] = round(yolo_max_conf, 1)
        state["rcnn"] = round(float(np.max(prob_3d)) * 100, 1)

        is_live_stream = state.get("is_live", False)
        threshold = 85 if is_live_stream else 75

        if confidence > threshold:
            state["final_decision"] = True
            state["label"] = f"Incident Logged: {severity.capitalize()}"
            threading.Thread(target=process_accident_async, args=(stream_id, raw_frame, location, confidence, severity)).start()
        else:
            state["label"] = f"Analyzing... ({severity.capitalize()})"

    except Exception as e:
        print(f"AI Error: {e}", flush=True)
    finally:
        if state: state["is_analyzing"] = False

# --- Route Handlers ---

@app.route('/init_upload', methods=['POST'])
def init_upload():
    file = request.files['file']
    filename = secure_filename(file.filename)
    path = os.path.join(app.config['UPLOAD_FOLDER'], filename)
    file.save(path)
    loc, wx, file_time = "Nil", "--", "--"
    stream_id = f"up_{uuid.uuid4().hex}"
    app.config[f'SRC_{stream_id}'] = path
    active_streams[stream_id] = {
        "label": "Initializing...", "final_decision": False, "confidence": 0, "severity": "--", 
        "plates": [], "show_tracking": True, "paused": False, "seek_to": None, "skip_val": 0, "progress": 0,
        "cnn": 0, "rcnn": 0, "force_update": False, "is_live": False
    }
    return jsonify({"stream_id": stream_id, "location": loc, "weather": wx, "time": file_time, "is_live": False})

@app.route('/init_stream', methods=['POST'])
def init_stream():
    url = request.json.get('url', '').strip()
    
    if 'youtube.com' in url or 'youtu.be' in url:
        try:
            ydl_opts = {'format': 'best', 'quiet': True, 'noplaylist': True}
            with yt_dlp.YoutubeDL(ydl_opts) as ydl:
                info = ydl.extract_info(url, download=False)
                if 'url' in info: 
                    url = info['url']
                elif 'formats' in info and len(info['formats']) > 0: 
                    for f in reversed(info['formats']):
                        if f.get('vcodec') != 'none':
                            url = f['url']
                            break
        except Exception as e:
            print(f"yt-dlp extraction failed: {e}", flush=True)

    loc, wx = get_geo_info(url)
    v_time = datetime.datetime.now().strftime("%H:%M:%S")
    stream_id = f"live_{uuid.uuid4().hex}"
    app.config[f'SRC_{stream_id}'] = url
    active_streams[stream_id] = {
        "label": "Connecting...", "final_decision": False, "confidence": 0, "severity": "--", 
        "plates": [], "show_tracking": True, "paused": False, "seek_to": None, "skip_val": 0, "progress": 0,
        "cnn": 0, "rcnn": 0, "force_update": False, "is_live": True
    }
    return jsonify({"stream_id": stream_id, "location": loc, "weather": wx, "time": v_time, "is_live": True})

@app.route('/video_control/<stream_id>', methods=['POST'])
def video_control(stream_id):
    state = active_streams.get(stream_id)
    if not state: return jsonify({"error": "not found"}), 404
    
    action = request.json.get('action')
    if action == 'toggle_tracking': 
        state["show_tracking"] = bool(request.json.get("track"))
        state["force_update"] = True 
    elif action == 'pause': state["paused"] = True
    elif action == 'play': state["paused"] = False
    elif action == 'seek': state["seek_to"] = request.json.get("value", 0.0)
    elif action == 'skip': state["skip_val"] = request.json.get("value", 0)
    
    return jsonify({"status": "ok"})

def video_stream_gen(stream_id, source, location):
    cap = cv2.VideoCapture(source)
    fps = cap.get(cv2.CAP_PROP_FPS)
    if not fps or fps == 0: fps = 25.0
    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    is_live = str(source).startswith('http')
    
    frames_3d = []
    yolo_probs = []
    frame_count = 0
    last_buffer = None
    retry_count = 0

    start_sync_time = time.time()
    current_frame_idx = 0

    if stream_id in active_streams:
        active_streams[stream_id]["label"] = "Scanning Stream..."

    while True:
        state = active_streams.get(stream_id)
        if not state: break

        force_read = False

        if state.get('skip_val', 0) != 0 and not is_live:
            target = max(0, min(current_frame_idx + (state['skip_val'] * fps), total_frames - 1))
            cap.set(cv2.CAP_PROP_POS_FRAMES, target)
            current_frame_idx = int(target)
            start_sync_time = time.time() - (current_frame_idx / fps)
            state['skip_val'] = 0
            frames_3d.clear(); yolo_probs.clear()
            force_read = True

        if state.get('seek_to') is not None and not is_live:
            target = int(state['seek_to'] * total_frames)
            cap.set(cv2.CAP_PROP_POS_FRAMES, target)
            current_frame_idx = target
            start_sync_time = time.time() - (current_frame_idx / fps)
            state['seek_to'] = None
            frames_3d.clear(); yolo_probs.clear()
            force_read = True
            
        if state.get('force_update'):
            force_read = True
            state['force_update'] = False

        if state.get('paused') and not force_read:
            start_sync_time = time.time() - (current_frame_idx / fps)
            time.sleep(0.1)
            if last_buffer: yield (b'--frame\r\nContent-Type: image/jpeg\r\n\r\n' + last_buffer + b'\r\n')
            continue

        if not is_live and not force_read:
            elapsed = time.time() - start_sync_time
            target_frame = int(elapsed * fps)
            frames_to_skip = target_frame - current_frame_idx
            if frames_to_skip > 0:
                for _ in range(min(frames_to_skip, 3)):
                    cap.grab()
                    current_frame_idx += 1

        ret, frame = cap.read()
        current_frame_idx += 1
        
        if not ret: 
            if is_live: 
                retry_count += 1
                if retry_count > 30: break
                cap.release(); time.sleep(0.5)
                cap = cv2.VideoCapture(source); continue
            elif total_frames > 0: 
                cap.set(cv2.CAP_PROP_POS_FRAMES, 0); current_frame_idx = 0
                start_sync_time = time.time(); frames_3d.clear(); yolo_probs.clear()
                continue
            else: break
            
        retry_count = 0
                
        h, w = frame.shape[:2]
        if w > 1280: frame = cv2.resize(frame, (1280, int(h * 1280 / w)))

        if is_live:
            state['progress'] = 100 
        elif total_frames > 0:
            state['progress'] = (current_frame_idx / total_frames) * 100

        try:
            if state.get("show_tracking", True):
                track_res = model_tracker(frame, classes=[2, 3, 5, 7], verbose=False, conf=0.3)[0]
                display_frame = track_res.plot()
            else:
                display_frame = frame.copy()
        except:
            display_frame = frame.copy()

        if not state.get("final_decision"):
            try:
                res = model_yolo(frame, verbose=False)[0]
                probs = np.zeros(4)
                if res.probs is not None:
                    data = res.probs.data.cpu().numpy()
                    probs[:min(len(data), 4)] = data[:min(len(data), 4)]
                elif res.boxes is not None and len(res.boxes) > 0:
                    for box in res.boxes:
                        cls_id = int(box.cls[0].item())
                        if cls_id < 4: probs[cls_id] = max(probs[cls_id], float(box.conf[0].item()))
                
                yolo_probs.append(probs)
                if len(yolo_probs) > 10: yolo_probs.pop(0)

                if len(yolo_probs) > 0:
                    curr_max = np.max(yolo_probs, axis=0)
                    acc_conf = float(np.max(curr_max[:3])) * 100 if len(curr_max) >= 3 else float(np.max(curr_max)) * 100
                    state["cnn"] = round(acc_conf, 1)
                    if state["cnn"] > 0 and state["confidence"] == 0:
                        state["label"] = "Scanning Spatial Features..."

                f_3d = cv2.resize(frame, (112, 112))
                frames_3d.append(cv2.cvtColor(f_3d, cv2.COLOR_BGR2RGB))
                
                if len(frames_3d) > 16: 
                    frames_3d.pop(0)

                analyze_interval = 10 if is_live else 5
                
                if len(frames_3d) == 16 and frame_count % analyze_interval == 0 and not state.get("is_analyzing"):
                    state["is_analyzing"] = True
                    
                    analysis_frames = list(frames_3d)
                    p_max = np.max(yolo_probs, axis=0) if len(yolo_probs) > 0 else np.zeros(4)
                    p_mean = np.mean(yolo_probs, axis=0) if len(yolo_probs) > 0 else np.zeros(4)
                    p_min = np.min(yolo_probs, axis=0) if len(yolo_probs) > 0 else np.zeros(4)
                    
                    acc_max_conf = float(np.max(p_max[:3])) * 100 if len(p_max) >= 3 else float(np.max(p_max)) * 100
                    
                    threading.Thread(target=run_temporal_analysis, args=(
                        stream_id, analysis_frames, p_max, p_mean, p_min, acc_max_conf, frame.copy(), location
                    )).start()
            except Exception as e: 
                state["is_analyzing"] = False

        _, buffer = cv2.imencode('.jpg', display_frame)
        last_buffer = buffer.tobytes()
        yield (b'--frame\r\nContent-Type: image/jpeg\r\n\r\n' + last_buffer + b'\r\n')
        frame_count += 1
        
        if not is_live:
            wait = (current_frame_idx / fps) - (time.time() - start_sync_time)
            if wait > 0: time.sleep(wait)
    
    cap.release()

@app.route('/video_feed/<stream_id>')
def video_feed(stream_id):
    source = app.config.get(f'SRC_{stream_id}')
    location = request.args.get('loc', 'Unknown')
    return Response(video_stream_gen(stream_id, source, location), mimetype='multipart/x-mixed-replace; boundary=frame')

@app.route('/stream_status/<stream_id>')
def stream_status(stream_id):
    return jsonify(active_streams.get(stream_id, {}))

@app.route('/predict_traffic_risk', methods=['POST'])
def predict_traffic_risk():
    data = request.json
    try:
        if model_traffic is None: return jsonify({"error": "Tabular model missing"}), 500
        
        csv_cols = [
            'Weather', 'Road_Type', 'Time_of_Day', 'Traffic_Density', 'Speed_Limit',
            'Number_of_Vehicles', 'Driver_Alcohol', 'Road_Condition', 'Vehicle_Type',
            'Driver_Age', 'Driver_Experience', 'Road_Light_Condition'
        ]
        
        row_data = {}
        for col in csv_cols:
            val = data.get(col, "")
            if col in ['Traffic_Density', 'Speed_Limit', 'Number_of_Vehicles', 'Driver_Alcohol', 'Driver_Age', 'Driver_Experience']:
                try:
                    row_data[col] = float(val)
                except (ValueError, TypeError):
                    row_data[col] = 0.0 
            else:
                row_data[col] = val if val != "" else "Unknown"
                
        df = pd.DataFrame([row_data], columns=csv_cols)
        
        probs = model_traffic.predict_proba(df)[0]
        model_classes = list(model_traffic.classes_)
        
        if 'High' in model_classes:
            prob = float(probs[model_classes.index('High')] * 100)
        elif 'Major' in model_classes:
            prob = float(probs[model_classes.index('Major')] * 100)
        elif 1 in model_classes:
            prob = float(probs[model_classes.index(1)] * 100)
        else:
            prob = float(probs[-1] * 100)
            
        status = "Major" if prob >= 70 else "Moderate" if prob >= 35 else "Minor"
        
        return jsonify({"status": status, "risk_probability_percentage": prob})
    except Exception as e: 
        print(f"Risk Predictor Error: {e}", flush=True)
        return jsonify({"error": str(e)}), 500

@app.route('/favicon.ico')
@app.route('/app_logo')
def serve_logo():
    # FIX: The "Ultimate Logo Finder". Checks both the static folder AND the root folder.
    root_dir = os.path.dirname(os.path.abspath(__file__))
    static_dir = os.path.join(root_dir, 'static')
    
    possible_paths = [
        os.path.join(static_dir, 'logo.png'),
        os.path.join(static_dir, 'logo.PNG'),
        os.path.join(root_dir, 'logo.png'),
        os.path.join(root_dir, 'logo.PNG')
    ]
    
    for path in possible_paths:
        if os.path.exists(path):
            return send_from_directory(os.path.dirname(path), os.path.basename(path), mimetype='image/png')
            
    return "Logo not found", 404

@app.route('/')
def index(): return render_template('index.html')

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=7860, threaded=True)