# scripts/visualize_contours.py
import os
import sys
import json
import base64

# Configure matplotlib first
os.environ['MPLCONFIGDIR'] = os.path.join(os.path.dirname(__file__), 'mplconfig')
os.makedirs(os.environ['MPLCONFIGDIR'], exist_ok=True)

try:
    import cv2
    import numpy as np
    import matplotlib
    matplotlib.use('Agg')
    import matplotlib.pyplot as plt
except ImportError as e:
    print(json.dumps({
        "error": f"Missing Python packages: {str(e)}",
        "solution": "Run: pip install matplotlib numpy opencv-python"
    }))
    sys.exit(1)

def verify_files(image_path, json_path):
    if not os.path.exists(image_path):
        return f"Image file not found: {image_path}"
    if not os.path.exists(json_path):
        return f"JSON file not found: {json_path}"
    return None

def load_contours(json_path):
    try:
        with open(json_path, 'r') as f:
            contours = json.load(f)
            return contours, None  # Return data and no error
    except Exception as e:
        return None, f"Failed to load JSON: {str(e)}"  # Return no data and error

def visualize(image_path, contours_data):
    try:
        img = cv2.imread(image_path)
        if img is None:
            return None, "Could not read image file"

        img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        plt.figure(figsize=(14, 10))
        plt.imshow(img_rgb)

        for idx, contour in enumerate(contours_data.get('contours', [])):
            if not contour:
                continue

            pts = np.array([[p['x'], p['y']] for p in contour], dtype=np.int32)
            pts = pts.reshape((-1, 1, 2))

            cv2.polylines(img, [pts], isClosed=True, color=(0, 255, 0), thickness=2)

            centroid_x = int(np.mean(pts[:, 0, 0]))
            centroid_y = int(np.mean(pts[:, 0, 1]))

            plt.text(centroid_x, centroid_y, f"#{idx}",
                    color='red', fontsize=12, weight='bold',
                    bbox=dict(facecolor='white', alpha=0.7))

        plt.title(f"Contours: {len(contours_data.get('contours', []))} found")
        plt.axis('off')

        # Save to temporary file
        temp_img = os.path.join(os.path.dirname(image_path), 'visualization.png')
        plt.savefig(temp_img, bbox_inches='tight', pad_inches=0)
        plt.close()

        with open(temp_img, 'rb') as f:
            img_base64 = base64.b64encode(f.read()).decode('utf-8')

        os.remove(temp_img)
        return img_base64, None

    except Exception as e:
        return None, f"Visualization error: {str(e)}"

if __name__ == "__main__":
    if len(sys.argv) < 3:
        print(json.dumps({
            "error": "Usage: python visualize_contours.py <image_path> <json_path>",
            "received_args": sys.argv[1:]
        }))
        sys.exit(1)

    image_path = sys.argv[1]
    json_path = sys.argv[2]

    # Verify files exist
    file_error = verify_files(image_path, json_path)
    if file_error:
        print(json.dumps({"error": file_error}))
        sys.exit(1)

    # Load contours data
    contours_data, load_error = load_contours(json_path)
    if load_error:
        print(json.dumps({"error": load_error}))
        sys.exit(1)

    # Generate visualization
    visualization, viz_error = visualize(image_path, contours_data)
    if viz_error:
        print(json.dumps({"error": viz_error}))
        sys.exit(1)

    # Return results
    result = {
        "success": True,
        "contour_count": len(contours_data.get('contours', [])),
        "image_size": contours_data.get('image_size', {}),
        "visualization": visualization,
        "message": "Visualization completed successfully"
    }

    print(json.dumps(result))
    sys.exit(0)