import os
os.environ["HF_HUB_OFFLINE"] = "1"
os.environ["TRANSFORMERS_OFFLINE"] = "1"
import time
import json
import shutil
import requests
from PIL import Image
import logging
import hashlib
import sys
import fcntl
import re

try:
    from . import database
    from . import embeddings
except ImportError:
    import database
    import embeddings

# Load configuration
CONFIG_PATH = os.path.join(os.path.dirname(__file__), "config.json")

def load_config():
    with open(CONFIG_PATH, 'r') as f:
        conf = json.load(f)
    # Resolve relative paths relative to this script's directory
    base_dir = os.path.dirname(os.path.abspath(__file__))
    for key in ["stage_dir", "output_dir", "failed_dir", "processed_file", "log_file"]:
        if not os.path.isabs(conf[key]):
            conf[key] = os.path.normpath(os.path.join(base_dir, "..", conf[key]))
    return conf

CONF = load_config()

# Setup logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler(CONF["log_file"]),
        logging.StreamHandler()
    ]
)

def get_processed_files():
    if os.path.exists(CONF["processed_file"]):
        try:
            with open(CONF["processed_file"], 'r') as f:
                data = json.load(f)
                if isinstance(data, list):
                    # Migration: convert old list format to dict
                    return {name: None for name in data}
                return data
        except Exception as e:
            logging.error(f"Error reading processed file: {e}")
            return {}
    return {}

def save_processed_files(processed):
    try:
        p = CONF["processed_file"]
        tmp = p + ".tmp"
        with open(tmp, 'w') as f:
            json.dump(processed, f, indent=2)
        os.replace(tmp, p)
    except Exception as e:
        logging.error(f"Error saving processed file: {e}")

def calculate_hash(filepath):
    """Calculate MD5 hash of a file."""
    hasher = hashlib.md5()
    try:
        with open(filepath, 'rb') as f:
            for chunk in iter(lambda: f.read(4096), b""):
                hasher.update(chunk)
        return hasher.hexdigest()
    except Exception as e:
        logging.error(f"Error calculating hash for {filepath}: {e}")
        return None

def crop_to_bbox(image_path, margin, top_margin=None, headroom=0.0):
    try:
        img = Image.open(image_path)
        if img.mode != 'RGBA':
            logging.info(f"Image {image_path} is mode {img.mode}, not RGBA. Skipping crop.")
            return img
        
        alpha = img.split()[-1]
        bbox = alpha.getbbox()
        if not bbox:
            logging.info(f"No non-transparent bbox found for {image_path}. Returning original.")
            return img
        
        if top_margin is None:
            top_margin = margin

        # Add margin
        left, upper, right, lower = bbox
        left = max(0, left - margin)
        upper = max(0, upper - top_margin)
        right = min(img.width, right + margin)
        lower = min(img.height, lower + margin)
        
        logging.info(f"Cropping {image_path} to {left, upper, right, lower} (margin={margin}, top_margin={top_margin})")
        cropped = img.crop((left, upper, right, lower))

        if headroom > 0:
            h_px = int(cropped.height * headroom)
            if h_px > 0:
                logging.info(f"Adding {h_px}px headroom to {image_path}")
                new_img = Image.new("RGBA", (cropped.width, cropped.height + h_px), (0, 0, 0, 0))
                new_img.paste(cropped, (0, h_px))
                return new_img
        
        return cropped
    except Exception as e:
        logging.error(f"Failed to crop {image_path}: {e}")
        raise

def is_file_stable(filepath):
    """Check if file size is stable for at least 1 second."""
    try:
        size1 = os.path.getsize(filepath)
        time.sleep(1)
        size2 = os.path.getsize(filepath)
        return size1 == size2 and size1 > 0
    except OSError:
        return False

def flag_image(filename):
    input_path = os.path.join(CONF["stage_dir"], filename)
    timestamp = time.strftime("%Y%m%d_%H%M%S")
    failed_filename = f"{timestamp}_{filename}"
    failed_path = os.path.join(CONF["failed_dir"], failed_filename)
    try:
        os.makedirs(CONF["failed_dir"], exist_ok=True)
        logging.info(f"Flagging image {filename} (moving to failed directory as {failed_filename})")
        shutil.move(input_path, failed_path)
    except Exception as e:
        logging.error(f"Failed to move {filename} to failed directory: {e}")

def process_image(filename):
    # Reload config in case it changed
    global CONF
    try:
        CONF = load_config()
    except:
        pass

    input_path = os.path.join(CONF["stage_dir"], filename)
    timestamp = time.strftime("%Y%m%d_%H%M%S")
    output_filename = f"{timestamp}_{filename}"
    output_path = os.path.join(CONF["output_dir"], output_filename)
    
    temp_path = input_path + ".tmp.png"
    try:
        logging.info(f"Starting processing for {filename}...")
        cropped_img = crop_to_bbox(
            input_path, 
            CONF["margin"], 
            top_margin=CONF.get("top_margin"), 
            headroom=CONF.get("headroom", 0.0)
        )
        
        # Save temporary cropped image for upload
        cropped_img.save(temp_path, format="PNG")
        
        prompt = CONF.get("prompt")
        if not prompt:
            bp = CONF.get("base_prompts", [])
            if bp and isinstance(bp, list) and len(bp) > 0:
                prompt = bp[0]
            else:
                prompt = "high quality, masterpiece"

        with open(temp_path, 'rb') as f:
            files = {'image': (filename, f, 'image/png')}
            data = {
                'prompt': prompt, 
                'seed': CONF.get("seed", -1),
                'max_area': CONF.get("max_area", 0)
            }
            logging.info(f"Calling API for {filename} -> {output_filename} with prompt: {prompt}")
            response = requests.post(CONF["api_url"], files=files, data=data, timeout=600)
            
        if response.status_code == 200:
            with open(output_path, 'wb') as f:
                f.write(response.content)
            logging.info(f"Successfully processed {filename} -> {output_path}")
            
            # Register in DB
            try:
                embedding = embeddings.generate_embedding(output_path)
                gid = filename
                database.upsert_person(output_filename, filepath=output_path, embedding=embedding, group_id=gid, is_source=True)
                
                # Also trigger tagging to get auto-name and clip description
                tag_url = CONF["api_url"].replace("/edit", "/tag")
                try:
                    requests.post(tag_url, json={"filename": output_filename, "group_id": gid}, timeout=30)
                except Exception as tag_err:
                    logging.error(f"Error triggering tagging for {output_filename}: {tag_err}")
            except Exception as db_err:
                logging.error(f"Database error registering {output_filename}: {db_err}")
                
            if os.path.exists(temp_path):
                os.remove(temp_path)
            return True
        else:
            logging.error(f"API failed for {filename}: {response.status_code} - {response.text}")
            if os.path.exists(temp_path):
                os.remove(temp_path)
            return False
    except requests.exceptions.ConnectionError as e:
        logging.error(f"Connection error while processing {filename}: {e}")
        if os.path.exists(temp_path):
            os.remove(temp_path)
        return None
    except Exception as e:
        logging.error(f"Error processing {filename}: {str(e)}", exc_info=True)
        if os.path.exists(temp_path):
            os.remove(temp_path)
        return False

def update_car_html():
    output_dir = CONF["output_dir"]
    car_html_path = os.path.join(output_dir, "car.html")
    if not os.path.exists(car_html_path):
        logging.warning(f"car.html not found at {car_html_path}")
        return

    try:
        # Use database to list only non-archived images
        persons = database.list_persons(include_archived=False)
        db_images = {p[0] for p in persons}
        
        # List images in output_dir
        extensions = ('.jpg', '.jpeg', '.png', '.gif', '.webp', '.bmp', '.svg')
        images = [f for f in os.listdir(output_dir) 
                  if f.lower().endswith(extensions) and f != "car.html" and f in db_images]
        
        # Sort by mtime, newest first
        images.sort(key=lambda x: os.path.getmtime(os.path.join(output_dir, x)), reverse=True)

        with open(car_html_path, 'r') as f:
            content = f.read()

        images_json = json.dumps(images, indent=12).strip()
        # Ensure it looks nice in the JS
        images_json = images_json.replace('\n', '\n            ')

        pattern = r'// --- HYDRATION_START ---.*?// --- HYDRATION_END ---'
        replacement = f'// --- HYDRATION_START ---\n        const PRELOADED_IMAGES = {images_json};\n        // --- HYDRATION_END ---'
        
        new_content = re.sub(pattern, replacement, content, flags=re.DOTALL)

        tmp_path = car_html_path + ".tmp"
        with open(tmp_path, 'w') as f:
            f.write(new_content)
        os.replace(tmp_path, car_html_path)
        logging.info(f"Updated {car_html_path} atomically with {len(images)} images")
    except Exception as e:
        logging.error(f"Failed to update car.html: {e}")

def main():
    # Prevent multiple instances
    lock_file = os.path.join(os.path.dirname(os.path.abspath(__file__)), "watcher.lock")
    fp = open(lock_file, 'w')
    try:
        fcntl.lockf(fp, fcntl.LOCK_EX | fcntl.LOCK_NB)
    except IOError:
        print("Another instance of watcher.py is already running. Exiting.")
        sys.exit(1)

    processed = get_processed_files()
    
    # Ensure directories exist
    os.makedirs(CONF["stage_dir"], exist_ok=True)
    os.makedirs(CONF["output_dir"], exist_ok=True)
    os.makedirs(CONF["failed_dir"], exist_ok=True)

    logging.info(f"Watcher started. Monitoring {CONF['stage_dir']}...")
    logging.info(f"Output directory: {CONF['output_dir']}")
    logging.info(f"API URL: {CONF['api_url']}")
    
    while True:
        try:
            files = [f for f in os.listdir(CONF["stage_dir"]) 
                     if f.lower().endswith(('.png', '.jpg', '.jpeg')) 
                     and not f.endswith('.tmp.png')]
            
            for f in files:
                input_path = os.path.join(CONF["stage_dir"], f)
                
                # Check if file is stable (not still being copied)
                if not is_file_stable(input_path):
                    continue
                
                # Calculate current file hash
                current_hash = calculate_hash(input_path)
                if not current_hash:
                    continue

                # Check if already processed
                if f in processed:
                    stored_hash = processed[f]
                    if stored_hash == current_hash:
                        continue
                    if stored_hash is None:
                        # Migration case: filename exists but no hash.
                        # Skip to avoid mass re-processing, but update the hash.
                        logging.info(f"Updating hash for previously processed {f}")
                        processed[f] = current_hash
                        save_processed_files(processed)
                        continue
                
                res = process_image(f)
                if res is True:
                    processed[f] = current_hash
                    save_processed_files(processed)
                    update_car_html()
                elif res is False:
                    flag_image(f)
                    # We don't add to processed here so that if the user 
                    # moves the file back to stage, it will be retried.
        except Exception as e:
            logging.error(f"Main loop error: {e}")
            
        time.sleep(CONF["poll_interval"])

if __name__ == "__main__":
    main()