import os
import shutil
import datetime
import random
from config import IMAGE_FOLDER, LABEL_FOLDER, TRAINED_MODEL_FOLDER, PRETRAINED_MODEL
from ultralytics import YOLO
import yaml
import requests
import glob
import sqlite3

# 🗂️ Configure paths
DATA_SPLIT_FOLDER = "model_datasets"
IMAGE_EXTENSION = ".png"  
LOG_FILE = "training_logs.log"

def get_label_names():
    conn = sqlite3.connect("my_database.db")
    conn.row_factory = sqlite3.Row  # Enables dictionary-like access to rows
    # Create a cursor object to execute SQL commands
    cursor = conn.cursor()

    # Fetch all names from the labels table
    cursor.execute("SELECT name FROM labels")
    labels = cursor.fetchall()
    conn.close()

    # Convert fetched rows into a simple list of strings
    label_names = [row["name"] for row in labels]

    return label_names

class_names = get_label_names()

# 📅 Get today's and yesterday's date
today = datetime.date.today()
yesterday = today - datetime.timedelta(days=1)

today_str = today.strftime("%Y-%m-%d")
yesterday_str = yesterday.strftime("%Y-%m-%d")

image_folder_today = os.path.join(IMAGE_FOLDER, today_str)
image_folder_yesterday = os.path.join(IMAGE_FOLDER, yesterday_str)

label_folder_today = os.path.join(LABEL_FOLDER, today_str)
label_folder_yesterday = os.path.join(LABEL_FOLDER, yesterday_str)

# 🔍 Validate folder existence
if os.path.exists(image_folder_today) and os.path.exists(label_folder_today):
    image_folder = image_folder_today
    label_folder = label_folder_today
    date_str = today_str
elif os.path.exists(image_folder_yesterday) and os.path.exists(label_folder_yesterday):
    image_folder = image_folder_yesterday
    label_folder = label_folder_yesterday
    date_str = yesterday_str
else:
    # print(image_folder)
    # print(label_folder)
    raise Exception("  No valid image & label folder found in the last two days!")

# 🏗️ Create dataset folders
dataset_folder = os.path.join(DATA_SPLIT_FOLDER, date_str)
model_folder_name = os.path.join(TRAINED_MODEL_FOLDER)
nolable_img_folder = os.path.join(dataset_folder, "nolable")
train_img_folder = os.path.join(dataset_folder, "train", "images")
train_lbl_folder = os.path.join(dataset_folder, "train", "labels")
val_img_folder = os.path.join(dataset_folder, "val", "images")
val_lbl_folder = os.path.join(dataset_folder, "val", "labels")

for folder in [train_img_folder, train_lbl_folder, val_img_folder, val_lbl_folder, nolable_img_folder]:
    os.makedirs(folder, exist_ok=True)

# 📂 Get image list
image_files = [f for f in os.listdir(image_folder) if f.endswith(IMAGE_EXTENSION)]

# 🌀 Shuffle images
random.shuffle(image_files)

# 📊 Split ratio
train_ratio = 0.8
split_idx = int(len(image_files) * train_ratio)

train_files = image_files[:split_idx]
val_files = image_files[split_idx:]



def log_message(message: str):
    """Ghi log vào file"""
    timestamp = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    log_entry = f"[{timestamp}] {message}\n"
    with open(LOG_FILE, "a") as log:
        log.write(log_entry)

def get_latest_model(trained_model_folder: str, default_model: str = "train5/weights/best.pt"):
    """Tìm model gần nhất (trước ngày hôm nay), nếu không có thì dùng model mặc định."""
    if not os.path.exists(trained_model_folder):
        log_message(f"  Folder {trained_model_folder} does not exist. Using default model: {default_model}")
        return default_model

    today_str = datetime.datetime.now().strftime("%Y-%m-%d")

    # Lấy danh sách các thư mục theo ngày, loại bỏ thư mục của ngày hôm nay
    subfolders = [
        f for f in os.listdir(trained_model_folder) 
        if os.path.isdir(os.path.join(trained_model_folder, f)) and f < today_str
    ]
    subfolders = sorted(subfolders, reverse=True)  # Sắp xếp giảm dần (mới nhất trước)

    for folder in subfolders:
        model_path = os.path.join(trained_model_folder, folder, "weights", "best.pt")
        if os.path.exists(model_path):
            log_message(f" Found latest model: {model_path}")
            return model_path

    log_message(f"  No valid models found in {trained_model_folder}. Using default model: {default_model}")
    return default_model

def create_dataset_yaml(dataset_folder: str, classes: list):
    dataset_path = os.path.abspath(dataset_folder)  
    data_yaml_path = os.path.join(dataset_folder, "data.yaml")

    data = {
        "train": os.path.join(dataset_path, "train", "images"),
        "val": os.path.join(dataset_path, "val", "images"),
        "test": os.path.join(dataset_path, "test", "images"),  # Nếu có test set
        "nc": len(classes),
        "names": list(classes)
    }

    with open(data_yaml_path, "w") as f:
        yaml.dump(data, f, default_flow_style=True)

    return data_yaml_path

def call_reload_model_api(base_url="http://localhost:5000"):
    """Gọi API để reload model"""
    url = f"{base_url}/reload_model"
    try:
        response = requests.post(url, timeout=10)
        response_data = response.json()
        log_message(f" API Response: {response_data} \n")
        return response_data
    except requests.RequestException as e:
        log_message(f" Error calling reload model API: {e} \n")
        return {"error": str(e)}

def clear_images_source():
    paths = [image_folder, label_folder]
    # paths = [image_folder, label_folder, train_img_folder, train_lbl_folder, val_img_folder, val_lbl_folder]

    for path in paths:
        for file in glob.glob(os.path.join(path, "*")):
            if os.path.isfile(file):  
                os.remove(file)

    
    log_message(f"Delete source image success \n")
    log_message("END " + ("=" * 20) + "\n")


# 🚀 Function to move files and log process
def move_files(file_list, dest_img_folder, dest_lbl_folder):
    copied_images = 0
    copied_labels = 0
    
    with open(LOG_FILE, "a") as log:
        for img_file in file_list:
            img_path = os.path.join(image_folder, img_file)
            label_file = os.path.splitext(img_file)[0] + ".txt"
            label_path = os.path.join(label_folder, label_file)
            
            if os.path.exists(img_path) and os.path.exists(label_path):
                shutil.copy(img_path, os.path.join(dest_img_folder, img_file))
                shutil.copy(label_path, os.path.join(dest_lbl_folder, label_file))
                copied_images += 1
                copied_labels += 1
            else:
                log.write(f"[{datetime.datetime.now()}]   Missing label for image: {img_file}\n")
                shutil.copy(img_path, os.path.join(nolable_img_folder, img_file))
                
                
    
    return copied_images, copied_labels


def train_yolo_model(pretrained_model: str, dataset_folder: str,project_name: str, name: str, epochs: int = 50, batch_size: int = 16, img_size: int = 640, lr: float = 0.001):
    dataset_yaml = os.path.join(dataset_folder, "data.yaml")

    if not os.path.exists(dataset_yaml):
        raise FileNotFoundError(f"  Not found file {dataset_yaml}. Plases check datasets!")

    if not os.path.exists(pretrained_model):
        log_message(f"  Model not found {pretrained_model}. Start train with 'yolov8n.pt'.")
        pretrained_model = "yolov8n.pt"

    # 🚀 Tạo model YOLOv8 và load model đã train trước đó
    model = YOLO(pretrained_model)

    # 🔥 Bắt đầu training
    model.train(
        data=dataset_yaml,
        epochs=epochs,
        batch=batch_size,
        imgsz=img_size,
        optimizer="AdamW",
        lr0=lr,
        weight_decay=0.0005,
        patience=10, 
        verbose=True,
        project=project_name,  
        name=name  
    )

    with open(LOG_FILE, "a") as log:
        log.write(f"\n[{datetime.datetime.now()}]  Train completed\n")

    call_reload_model_api()
    clear_images_source()
 



# 🚀 Copy files
train_copied_imgs, train_copied_lbls = move_files(train_files, train_img_folder, train_lbl_folder)
val_copied_imgs, val_copied_lbls = move_files(val_files, val_img_folder, val_lbl_folder)

# Create yml file
create_dataset_yaml(dataset_folder, class_names)


# 🏁 Log summary
with open(LOG_FILE, "a") as log:
    log.write(f"\n[{datetime.datetime.now()}]  Dataset split completed\n")
    log.write(f"Source folder: {image_folder}\n")
    log.write(f"Dataset saved in: {dataset_folder}\n")
    log.write(f"Planned Train: {len(train_files)} images, Val: {len(val_files)} images\n")
    log.write(f"Actual Train: {train_copied_imgs} images, {train_copied_lbls} labels\n")
    log.write(f"Actual Val: {val_copied_imgs} images, {val_copied_lbls} labels\n")


if(train_copied_imgs <=0 or train_copied_lbls <= 0 or val_copied_imgs <=0 or val_copied_lbls <=0):
    with open(LOG_FILE, "a") as log:
        log.write(f"\n[{datetime.datetime.now()}]  Data not qualified\n")
        log.write("=" * 50 + "\n")
else:
    train_yolo_model(pretrained_model = get_latest_model(
        trained_model_folder=TRAINED_MODEL_FOLDER, 
        default_model=PRETRAINED_MODEL), 
        dataset_folder = dataset_folder, epochs = 2, name=today_str, project_name=model_folder_name
        )

