from flask import Flask, request
from flask_restx import Api, Resource, fields
import os
from datetime import datetime
from yoloseg.inference_ import Inference, overlay_mask
import cv2
import time

app = Flask(__name__)
api = Api(app, version='1.0', title='CCTV Image Upload API',
          description='A simple API for receiving CCTV images')

# Namespace definition
ns = api.namespace('cctv', description='CCTV operations')

model_path = 'yoloseg/weight/best.onnx'
classes_txt_file = 'config/yolo_config.txt'
image_path = 'yoloseg/img3.jpg'

model_input_shape = (640, 640)
inference_engine = Inference(
    onnx_model_path=model_path,
    model_input_shape=model_input_shape,
    classes_txt_file=classes_txt_file,
    run_with_cuda=True
)

# Define the expected model for incoming data
image_upload_model = api.model('ImageUpload', {
    'image': fields.String(required=True, description='Image file', dt='File'),
    'x-cctv-info': fields.String(required=False, description='CCTV identifier'),
    'x-time-sent': fields.String(required=False, description='Time image was sent'),
    'x-cctv-latitude': fields.String(required=False, description='Latitude of CCTV'),
    'x-cctv-longitude': fields.String(required=False, description='Longitude of CCTV')
})

# Define the directory where images will be saved
IMAGE_DIR = "network_test"
if not os.path.exists(IMAGE_DIR):
    os.makedirs(IMAGE_DIR)

@ns.route('/infer', )
class ImageUpload(Resource):
    # @ns.expect(image_upload_model, validate=True)
    @ns.response(200, 'Success')
    @ns.response(400, 'Validation Error')
    def post(self):
        if 'file' not in request.files:
            ns.abort(400, 'No image part in the request')
        image = request.files['file']
        cctv_info = request.headers.get('x-cctv-info', '')
        time_sent = request.headers.get('x-time-sent', '')
        cctv_latitude = request.headers.get('x-cctv-latitude', 'Not provided')
        cctv_longitude = request.headers.get('x-cctv-longitude', 'Not provided')

        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        image = cv2.imdecode(image)
        filename = f"{timestamp}_{cctv_info}.png"

        t1 = time.time()
        detections, mask_maps = inference_engine.run_inference(image)
        t2 = time.time()

        if mask_maps is not None:
            seg_image = overlay_mask(image, mask_maps[0], color=(0, 255, 0), alpha=0.3)

        if image.filename == '':
            ns.abort(400, 'No selected image')

        # Use current timestamp to avoid filename conflicts
        # image_path = os.path.join(IMAGE_DIR, filename)
        # image.save(image_path)

        return {"message": f"Image {filename} uploaded successfully!"}

    def send_result(self):
        pass
if __name__ == '__main__':
    app.run(debug=True, port=12345)