from fastapi import FastAPI, UploadFile, File, Form
from fastapi.responses import JSONResponse
from ultralytics import YOLO
import numpy as np
from PIL import Image
from datetime import datetime, timedelta
import psycopg2
import io

app = FastAPI()

model = YOLO('./best v3(100).pt')

'''def init_connection():
    return psycopg2.connect(
        host="localhost",
        database="water_meter",
        user="postgres",
        password="16ecr193@kec"
    )

def insert_detection(conn, detection_date, water_meter_column, detected_values_str):
    with conn.cursor() as cur:
        created_time = datetime.now().strftime("%H:%M:%S")

        cur.execute(f"""
            INSERT INTO water_meter_readings (detection_date, {water_meter_column}, created_time)
            VALUES (%s, %s, %s)
            ON CONFLICT (detection_date) DO UPDATE
            SET {water_meter_column} = EXCLUDED.{water_meter_column},
                created_time = EXCLUDED.created_time
        """, (detection_date, int(detected_values_str), created_time))

        meter_number = water_meter_column.split('_')[-1]
        difference_column = f'difference_{meter_number}'

        yesterday = detection_date - timedelta(days=1)

        cur.execute(f"""
            SELECT {water_meter_column}
            FROM water_meter_readings
            WHERE detection_date = %s
        """, (yesterday,))
        yesterday_result = cur.fetchone()

        if yesterday_result and yesterday_result[0] is not None:
            difference = int(detected_values_str) - yesterday_result[0]
        else:
            difference = None

        if difference is not None:
            cur.execute(f"""
                UPDATE water_meter_readings
                SET {difference_column} = %s
                WHERE detection_date = %s
            """, (difference, detection_date))
        else:
            cur.execute(f"""
                UPDATE water_meter_readings
                SET {difference_column} = NULL
                WHERE detection_date = %s
            """, (detection_date,))

        conn.commit()'''

@app.get("/")
def read_root():
    return {"message": "Welcome to Water Meter Detection API"}

@app.post("/detect/")
async def detect_water_meter(
    #detection_date: str = Form(...),
    #water_meter: int = Form(...),
    file: UploadFile = File(...)
):
    image_data = await file.read()
    image = Image.open(io.BytesIO(image_data))
    img_array = np.array(image)

    results = model.predict(source=img_array, conf=0.25, imgsz=640)

    detected_numbers = []
    for result in results:
        for box in result.boxes:
            label = result.names[box.cls[0].item()]
            x1, _, _, _ = box.xyxy[0].tolist()

            if label.isdigit():
                detected_numbers.append((x1, label))

    detected_numbers.sort(key=lambda x: x[0])
    sorted_numbers = [num for _, num in detected_numbers]
    detected_values_str = ''.join(sorted_numbers)

    '''if detected_values_str:
        water_meter_column = f'water_meter_{water_meter}'
        detection_date_obj = datetime.strptime(detection_date, "%Y-%m-%d").date()

        conn = init_connection()
        insert_detection(conn, detection_date_obj, water_meter_column, detected_values_str)
        conn.close()'''

    return JSONResponse(
        content={
            "status": "success",
            "detected_values": detected_values_str if sorted_numbers else "No numeric values detected."
        }
    )

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="127.0.0.1", port=8000)