from fastapi import FastAPI, HTTPException, File, UploadFile
from fastapi.responses import StreamingResponse, JSONResponse
from fastapi.middleware.cors import CORSMiddleware
import cv2
import numpy as np
import json
import uvicorn
from kafka import KafkaProducer, KafkaConsumer
from redis import Redis
import uuid
import os
from datetime import timedelta
import threading
import mediapipe as mp
from mediapipe.tasks import python
from mediapipe.tasks.python import vision

app = FastAPI()
mediapipe_app = FastAPI()
app.mount("/mediapipe", mediapipe_app)

# CORS configuration
ALLOWED_ORIGINS = 'https://beta.obscura.work'

app.add_middleware(
    CORSMiddleware,
    allow_origins=ALLOWED_ORIGINS,
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Configuration
MODEL_PATH = "/home/zydi/models/face_landmarker.task"  # Replace with your model path
KAFKA_BROKER = "222.186.10.253:9092"
KAFKA_TOPIC = "mediapipe"
KAFKA_GROUP_ID = "mediapipe_group"

REDIS_HOST = "222.186.10.253"
REDIS_PORT = 6379
REDIS_PASSWORD = "Obscura@2024" 
REDIS_DB = 10 

UPLOAD_DIR = "/www/wwwroot/beta.obscura.work/upload_files/upload"
RESULT_DIR = "/www/wwwroot/beta.obscura.work/upload_files/result"
MAX_FILE_AGE = timedelta(hours=1)

# Ensure directories exist
os.makedirs(UPLOAD_DIR, exist_ok=True)
os.makedirs(RESULT_DIR, exist_ok=True)

# Initialize Kafka
producer = KafkaProducer(bootstrap_servers=[KAFKA_BROKER])
consumer = KafkaConsumer(
    KAFKA_TOPIC,
    bootstrap_servers=[KAFKA_BROKER],
    group_id=KAFKA_GROUP_ID,
    auto_offset_reset='earliest',
    enable_auto_commit=True,
    value_deserializer=lambda x: json.loads(x.decode('utf-8'))
)

# Initialize Redis
redis_client = Redis(
    host=REDIS_HOST,
    port=REDIS_PORT,
    password=REDIS_PASSWORD,
    db=REDIS_DB
)

class mediapipeEmbedder:
    def __init__(self, model_path):
        base_options = python.BaseOptions(model_asset_path=model_path)
        options = vision.FaceLandmarkerOptions(
            base_options=base_options,
            output_face_blendshapes=True,
            output_facial_transformation_matrixes=True,
            num_faces=1
        )
        self.detector = vision.FaceLandmarker.create_from_options(options)

    def get_mediapipe_landmarks(self, image):
        mp_image = mp.Image(image_format=mp.ImageFormat.SRGB, data=image)
        detection_result = self.detector.detect(mp_image)
        if detection_result.face_landmarks:
            return np.array([(lm.x, lm.y, lm.z) for lm in detection_result.face_landmarks[0]])
        return None

    def process_image(self, image_data):
        img = cv2.imdecode(np.frombuffer(image_data, np.uint8), cv2.IMREAD_COLOR)
        landmarks = self.get_mediapipe_landmarks(img)
        
        if landmarks is not None:
            # Calculate a more detailed mediapipe embedding
            embedding = self.calculate_detailed_embedding(landmarks)
            
            # Draw landmarks on the image
            for lm in landmarks:
                cv2.circle(img, (int(lm[0]*img.shape[1]), int(lm[1]*img.shape[0])), 2, (0,255,0), -1)
            
            return {
                "embedding": embedding,
                "landmarks": landmarks.tolist()
            }, img
        else:
            return None, img
        
    def calculate_detailed_embedding(self, landmarks):
        # Calculate various statistical features
        mean = np.mean(landmarks, axis=0)
        std = np.std(landmarks, axis=0)
        median = np.median(landmarks, axis=0)
        min_vals = np.min(landmarks, axis=0)
        max_vals = np.max(landmarks, axis=0)
        
        # Calculate pairwise distances between key facial landmarks
        nose_tip = landmarks[4]
        left_eye = landmarks[159]
        right_eye = landmarks[386]
        left_mouth = landmarks[61]
        right_mouth = landmarks[291]
        
        eye_distance = np.linalg.norm(left_eye - right_eye)
        mouth_width = np.linalg.norm(left_mouth - right_mouth)
        nose_to_mouth = np.linalg.norm(nose_tip - (left_mouth + right_mouth) / 2)
        
        # Calculate face shape features
        face_width = np.max(landmarks[:, 0]) - np.min(landmarks[:, 0])
        face_height = np.max(landmarks[:, 1]) - np.min(landmarks[:, 1])
        face_depth = np.max(landmarks[:, 2]) - np.min(landmarks[:, 2])
        
        # Combine all features into a single embedding
        embedding = np.concatenate([
            mean, std, median, min_vals, max_vals,
            [eye_distance, mouth_width, nose_to_mouth, face_width, face_height, face_depth]
        ])
        
        return embedding.tolist()

embedder = mediapipeEmbedder(MODEL_PATH)

def process_image(image_data, filename):
    try:
        results, annotated_img = embedder.process_image(image_data)
        
        if results:
            # Save annotated image
            annotated_filename = f"mediapipe_{filename}"
            annotated_path = os.path.join(RESULT_DIR, annotated_filename)
            cv2.imwrite(annotated_path, annotated_img)

            return results, annotated_filename
        else:
            print(f"No face landmarks detected in image: {filename}")
            return None, None
    except Exception as e:
        print(f"Error processing image {filename}: {str(e)}")
        import traceback
        traceback.print_exc()
        return None, None

def process_video(video_data, filename):
    try:
        temp_video_path = os.path.join(UPLOAD_DIR, f"mediapipe_{filename}")
        with open(temp_video_path, 'wb') as temp_video:
            temp_video.write(video_data)

        cap = cv2.VideoCapture(temp_video_path)
        frame_count = 0
        results = []
        
        fps = int(cap.get(cv2.CAP_PROP_FPS))
        width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
        height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
        
        annotated_filename = f"mediapipe_{filename}"
        output_path = os.path.join(RESULT_DIR, annotated_filename)
        out = cv2.VideoWriter(output_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (width, height))

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

            if frame_count % fps == 0:
                frame_results, annotated_frame = embedder.process_image(cv2.imencode('.jpg', frame)[1].tobytes())
                if frame_results:
                    results.append({"frame": frame_count, "results": frame_results})
            else:
                annotated_frame = frame

            out.write(annotated_frame)
            frame_count += 1

        cap.release()
        out.release()

        os.remove(temp_video_path)

        return results, annotated_filename
    except Exception as e:
        print(f"Error processing video: {str(e)}")
        return None, None

@mediapipe_app.post("/upload")
async def upload_file(file: UploadFile = File(...)):
    content = await file.read()
    file_extension = os.path.splitext(file.filename)[1].lower()
    new_filename = f"{uuid.uuid4()}{file_extension}"
    
    original_file_path = os.path.join(UPLOAD_DIR, new_filename)
    with open(original_file_path, "wb") as f:
        f.write(content)

    producer.send(KAFKA_TOPIC, json.dumps({
        "filename": new_filename,
        "file_type": "image" if file_extension in ['.jpg', '.jpeg', '.png'] else "video"
    }).encode('utf-8'))

    redis_key = f"mediapipe_result:{new_filename}"
    redis_client.set(redis_key, json.dumps({"status": "queued"}))

    return JSONResponse(content={"message": "File uploaded and queued for processing", "filename": new_filename})

@mediapipe_app.get("/result/{filename}")
async def get_mediapipe_result(filename: str):
    redis_key = f"mediapipe_result:{filename}"
    result = redis_client.get(redis_key)
    if result:
        result_data = json.loads(result)
        return JSONResponse(content=result_data)
    else:
        raise HTTPException(status_code=404, detail="Result not found")

@mediapipe_app.get("/annotated/{filename}")
async def get_annotated_file(filename: str):
    redis_key = f"mediapipe_result:{filename}"
    result = redis_client.get(redis_key)
    if result:
        result_data = json.loads(result)
        if result_data["status"] == "completed":
            annotated_filename = result_data["annotated_filename"]
            file_path = os.path.join(RESULT_DIR, annotated_filename)
            if os.path.exists(file_path):
                def iterfile():
                    with open(file_path, mode="rb") as file_like:
                        yield from file_like
                file_extension = os.path.splitext(annotated_filename)[1].lower()
                return StreamingResponse(iterfile(), media_type=f"image/{file_extension[1:]}" if file_extension in ['.jpg', '.jpeg', '.png'] else "video/mp4")
    
    raise HTTPException(status_code=404, detail="Annotated file not found")

def process_task():
    for message in consumer:
        task = message.value
        filename = task['filename']
        file_type = task['file_type']
        
        file_path = os.path.join(UPLOAD_DIR, filename)
        
        redis_key = f"mediapipe_result:{filename}"
        redis_client.set(redis_key, json.dumps({"status": "processing"}))
        
        try:
            if file_type == "image":
                with open(file_path, 'rb') as f:
                    content = f.read()
                results, annotated_filename = process_image(content, filename)
            else:
                with open(file_path, 'rb') as f:
                    content = f.read()
                results, annotated_filename = process_video(content, filename)
            
            if results and annotated_filename:
                redis_client.set(redis_key, json.dumps({
                    "results": results,
                    "status": "completed",
                    "annotated_filename": annotated_filename
                }))
            else:
                redis_client.set(redis_key, json.dumps({"status": "failed"}))
        except Exception as e:
            print(f"Error processing task: {str(e)}")
            redis_client.set(redis_key, json.dumps({"status": "failed", "error": str(e)}))

def listen_redis_changes():
    pubsub = redis_client.pubsub()
    pubsub.psubscribe('__keyspace@3__:mediapipe_result:*')
    
    for message in pubsub.listen():
        if message['type'] == 'pmessage':
            key = message['channel'].decode('utf-8').split(':')[-1]
            operation = message['data'].decode('utf-8')
            
            if operation == 'set':
                value = redis_client.get(f"mediapipe_result:{key}")
                if value:
                    result = json.loads(value)
                    print(f"Status update for {key}: {result['status']}")

if __name__ == "__main__":
    threading.Thread(target=process_task, daemon=True).start()
    threading.Thread(target=listen_redis_changes, daemon=True).start()
    uvicorn.run(app, host="0.0.0.0", port=7006)