import json
import logging
import os
import xml.etree.ElementTree as ET
from datetime import datetime
from typing import Dict, Any
import boto3
from PIL import Image
import io
import urllib.parse
import hashlib

logger = logging.getLogger(__name__)
s3_client = boto3.client('s3')

def ensure_jsonl_format(manifest_content: str) -> str:
    """
    Ensure the manifest content is in proper JSONL format
    (each line is a valid JSON object with no trailing commas)
    """
    try:
        data = json.loads(manifest_content)
        if isinstance(data, list):
            # Convert array to JSONL
            return '\n'.join(json.dumps(item) for item in data)
        elif isinstance(data, dict):
            # Single object
            return json.dumps(data)
        else:
            logger.error(f"Unexpected manifest content type: {type(data)}")
            return manifest_content
    except json.JSONDecodeError:
        # Already might be JSONL, try parsing line by line
        lines = manifest_content.strip().split('\n')
        valid_lines = []
        for line in lines:
            if not line.strip():
                continue
            try:
                entry = json.loads(line)
                valid_lines.append(json.dumps(entry))
            except json.JSONDecodeError:
                logger.warning(f"Skipping invalid JSON line: {line[:100]}...")
        
        if not valid_lines:
            logger.error("No valid JSON entries found in manifest")
            raise ValueError("Invalid manifest format")
            
        return '\n'.join(valid_lines)

def ensure_valid_s3_uri(uri: str) -> str:
    """
    Ensure the S3 URI is in the correct format expected by Rekognition.
    Fix common issues with S3 URIs without URL encoding.
    """
    # Check if already processed to avoid double-encoding
    if '%25' in uri:
        logger.warning("URI appears to be encoded, attempting to decode")
        # Try to decode the URI
        try:
            # First decode all %25 sequences (encoded % signs)
            while '%25' in uri:
                uri = uri.replace('%25', '%')
            
            # Then decode the URI
            uri = urllib.parse.unquote(uri)
        except Exception as e:
            logger.error(f"Error trying to decode URI: {str(e)}")
            return uri
    
    # Handle non-encoded URIs
    if not uri.startswith('s3://'):
        if uri.startswith('s3:/'):
            uri = 's3://' + uri[4:]  # Fix missing slash
        else:
            return uri
    
    # Split into components
    parts = uri.split('/', 3)  # Split into ['s3:', '', 'bucket', 'path']
    if len(parts) < 4:
        return uri
        
    # Get bucket and key
    bucket = parts[2]
    key = parts[3]
    
    # Reconstruct the URI without encoding
    return f"s3://{bucket}/{key}"

def create_test_manifest(training_bucket: str, test_image_key: str) -> Dict[str, Any]:
    """
    Create a minimal test manifest entry following AWS's format exactly.
    """
    manifest_entry = {
        "source-ref": ensure_valid_s3_uri(f"s3://{training_bucket}/{test_image_key}"),
        "bounding-box": {
            "annotations": [{
                "class_id": 0,
                "left": 0.2,  # Larger bounding box for better detection
                "top": 0.2,
                "width": 0.6,
                "height": 0.6
            }],
            "image_size": [{
                "width": 500,
                "height": 500,
                "depth": 3
            }]
        },
        "bounding-box-metadata": {
            "class-map": {"0": "test-class"},
            "type": "groundtruth/object-detection",
            "human-annotated": "yes",
            "creation-date": datetime.utcnow().isoformat(),
            "job-name": "test-manifest",
            "objects": [{"confidence": 1}]
        }
    }
    return manifest_entry

def create_test_image(training_bucket: str, test_image_key: str) -> None:
    """
    Create and upload a simple test image to S3.
    """
    try:
        # Create a simple test image
        img = Image.new('RGB', (500, 500), color=(255, 0, 0))
        buffer = io.BytesIO()
        img.save(buffer, format='JPEG')
        buffer.seek(0)
        
        # Upload the image to S3
        s3_client.upload_fileobj(
            buffer,
            training_bucket,
            test_image_key,
            ExtraArgs={'ContentType': 'image/jpeg'}
        )
        logger.info(f"Created test image: s3://{training_bucket}/{test_image_key}")
        
    except Exception as e:
        logger.error(f"Error creating test image: {str(e)}")
        raise

def validate_bounding_box(ann: Dict[str, Any], img_width: int, img_height: int) -> bool:
    """
    Validate that a bounding box is reasonable in size.
    Returns True if valid, False if too small.
    """
    # Get pixel measurements directly (since we're now using pixel values)
    width_px = ann["width"]
    height_px = ann["height"]
    area_px = width_px * height_px
    
    # Minimum size requirements - very lenient
    MIN_WIDTH = 1  # minimum 1 pixel wide
    MIN_HEIGHT = 1  # minimum 1 pixel high
    MIN_AREA = 1   # minimum 1 pixel area
    
    # Maximum size requirements - prevent unreasonable boxes
    MAX_WIDTH = img_width - 1  # max width should be within image
    MAX_HEIGHT = img_height - 1  # max height should be within image
    
    # Log the actual values for debugging
    logger.info(f"Validating bounding box - Width: {width_px}px ({(width_px/img_width)*100:.1f}% of image), "
               f"Height: {height_px}px ({(height_px/img_height)*100:.1f}% of image), "
               f"Area: {area_px}px² ({(area_px/(img_width*img_height))*100:.1f}% of image)")
    
    if width_px < MIN_WIDTH:
        logger.warning(f"Bounding box too narrow: {width_px}px (minimum {MIN_WIDTH}px)")
        return False
        
    if height_px < MIN_HEIGHT:
        logger.warning(f"Bounding box too short: {height_px}px (minimum {MIN_HEIGHT}px)")
        return False
        
    if area_px < MIN_AREA:
        logger.warning(f"Bounding box too small: {area_px}px² (minimum {MIN_AREA}px²)")
        return False
        
    if width_px > MAX_WIDTH:
        logger.warning(f"Bounding box too wide: {width_px}px (maximum {MAX_WIDTH}px)")
        return False
        
    if height_px > MAX_HEIGHT:
        logger.warning(f"Bounding box too tall: {height_px}px (maximum {MAX_HEIGHT}px)")
        return False
        
    # Validate coordinates are within bounds
    left_px = ann["left"]
    top_px = ann["top"]
    right_px = left_px + width_px
    bottom_px = top_px + height_px
    
    if left_px < 0 or top_px < 0 or right_px > img_width or bottom_px > img_height:
        logger.warning(f"Bounding box coordinates out of bounds: "
                      f"left={left_px}, top={top_px}, right={right_px}, bottom={bottom_px}")
        return False
        
    return True

def convert_xml_to_rekognition_format(xml_content: bytes, source_bucket: str, image_key: str) -> Dict[str, Any]:
    """
    Convert XML annotation to Rekognition format.
    """
    try:
        # Parse the XML content
        root = ET.fromstring(xml_content)

        # Get image size    
        size = root.find('size')    
        img_width = int(size.find('width').text)    
        img_height = int(size.find('height').text)    
        
        logger.info(f"Processing image {image_key} with size {img_width}x{img_height}")
        
        # Prepare annotations list
        annotations = []
        class_map = {}
        class_id = 0
        
        # Track skipped boxes for logging
        skipped_boxes = []
        
        # Loop through each object in the XML    
        for obj in root.findall('object'):        
            label = obj.find('name').text
            
            # Clean and validate label name
            label = label.strip().lower().replace(' ', '_')
            if len(label) > 256:
                logger.warning(f"Label name too long, truncating: {label}")
                label = label[:256]
            
            logger.info(f"Processing object with label: {label}")
            
            # Add to class map if not already present
            if label not in class_map.values():
                class_map[str(class_id)] = label
                class_id += 1
                
            bndbox = obj.find('bndbox')        
            xmin = float(bndbox.find('xmin').text)        
            ymin = float(bndbox.find('ymin').text)        
            xmax = float(bndbox.find('xmax').text)        
            ymax = float(bndbox.find('ymax').text)        

            # Calculate width and height
            width = xmax - xmin
            height = ymax - ymin
            
            # Log original coordinates
            logger.info(f"Original coordinates: xmin={xmin}, ymin={ymin}, width={width}, height={height}")
            
            # Ensure coordinates are within image bounds
            xmin = max(0, min(xmin, img_width))
            ymin = max(0, min(ymin, img_height))
            xmax = max(0, min(xmax, img_width))
            ymax = max(0, min(ymax, img_height))
            
            # Recalculate width and height after bounding
            width = xmax - xmin
            height = ymax - ymin
            
            # Get class_id for this label
            class_id_for_label = next(k for k, v in class_map.items() if v == label)
            
            # Use pixel values directly instead of normalizing coordinates
            # Convert to integers as required by AWS Rekognition
            left = int(xmin)
            top = int(ymin)
            width_px = int(width)
            height_px = int(height)
            
            # Add annotation with pixel coordinates (not normalized)
            ann = {
                "class_id": int(class_id_for_label),
                "left": left,
                "top": top,
                "width": width_px,
                "height": height_px
            }
            
            # Log pixel coordinates
            logger.info(f"Pixel coordinates: left={ann['left']}, top={ann['top']}, width={ann['width']}, height={ann['height']}")
            
            # Validate the bounding box - checking if it's within image boundaries
            if (0 <= left < img_width and 
                0 <= top < img_height and 
                width_px > 0 and 
                height_px > 0 and
                left + width_px <= img_width and
                top + height_px <= img_height):
                annotations.append(ann)
                logger.info(f"Added valid annotation for {label}")
            else:
                skipped_boxes.append({
                    'label': label,
                    'width_px': width_px,
                    'height_px': height_px,
                    'area_px': width_px * height_px,
                    'coordinates': {
                        'left': left,
                        'top': top,
                        'width': width_px,
                        'height': height_px
                    }
                })
                logger.warning(f"Skipped invalid bounding box for {label}")
        
        if skipped_boxes:
            logger.warning(f"Skipped {len(skipped_boxes)} boxes that were invalid:")
            for box in skipped_boxes:
                logger.warning(
                    f"- {box['label']}: {box['width_px']}x{box['height_px']}px ({box['area_px']}px²) "
                    f"Coordinates: left={box['coordinates']['left']}, top={box['coordinates']['top']}, "
                    f"width={box['coordinates']['width']}, height={box['coordinates']['height']}"
                )
        
        if not annotations:
            logger.error(f"No valid annotations found for {image_key} - all bounding boxes were invalid")
            raise ValueError("No valid annotations")
            
        # Limit annotations to 50 as per AWS Rekognition requirements
        MAX_ANNOTATIONS = 50
        if len(annotations) > MAX_ANNOTATIONS:
            logger.warning(f"Limiting annotations from {len(annotations)} to {MAX_ANNOTATIONS} for image {image_key} to prevent ERROR_TOO_MANY_BOUNDING_BOXES")
            annotations = annotations[:MAX_ANNOTATIONS]
        
        # Create manifest entry with properly encoded source-ref
        source_ref = f"s3://{source_bucket}/{image_key}"
        
        # Clean up the source-ref to ensure proper format with double slashes after s3:
        # First make sure it starts with s3:// (double slash)
        if source_ref.startswith('s3:/') and not source_ref.startswith('s3://'):
            source_ref = 's3://' + source_ref[4:]
        
        # Fix any duplicate slashes in the path part (but not in the s3:// prefix)
        if '//' in source_ref[5:]:
            parts = source_ref.split('/', 3)
            if len(parts) >= 4:
                source_ref = f"s3://{parts[2]}/{parts[3].replace('//', '/')}"
        
        # Replace any URL-encoded spaces with actual spaces (S3 allows spaces in object keys)
        source_ref = source_ref.replace('%20', ' ')
        
        # Replace any other common URL encodings that shouldn't be in the path
        source_ref = source_ref.replace('%25', '%').replace('%2F', '/').replace('%3A', ':')
        
        logger.info(f"Using source-ref: {source_ref}")
        
        # Create manifest entry in the exact format required by Rekognition
        manifest_entry = {
            "source-ref": source_ref,
            "bounding-box": {
                "image_size": [{
                    "width": img_width,
                    "height": img_height,
                    "depth": 3
                }],
                "annotations": annotations
            },
            "bounding-box-metadata": {
                "objects": [{"confidence": 1} for _ in annotations],
                "class-map": class_map,
                "type": "groundtruth/object-detection",
                "human-annotated": "yes",
                "creation-date": datetime.utcnow().isoformat(),
                "job-name": "xml-import"
            }
        }
        
        # Validate the manifest entry format
        if not validate_manifest_format(manifest_entry):
            raise ValueError("Invalid manifest format")
        
        return manifest_entry
        
    except Exception as e:
        logger.error(f"Error converting XML to Rekognition format: {str(e)}")
        raise

def validate_manifest_format(entry: Dict[str, Any]) -> bool:
    """
    Validate that a manifest entry follows the exact format required by Rekognition.
    """
    try:
        # Check required top-level keys
        required_keys = ["source-ref", "bounding-box", "bounding-box-metadata"]
        for key in required_keys:
            if key not in entry:
                logger.error(f"Missing required key: {key}")
                return False
        
        # Validate source-ref
        if not isinstance(entry["source-ref"], str) or not entry["source-ref"].startswith("s3://"):
            logger.error("Invalid source-ref format")
            return False
        
        # Validate bounding-box
        bb = entry["bounding-box"]
        if not isinstance(bb, dict):
            logger.error("bounding-box must be a dictionary")
            return False
        
        # Validate image_size
        if "image_size" not in bb or not isinstance(bb["image_size"], list) or len(bb["image_size"]) != 1:
            logger.error("Invalid image_size format")
            return False
        
        size = bb["image_size"][0]
        for key in ["width", "height", "depth"]:
            if key not in size or not isinstance(size[key], int) or size[key] <= 0:
                logger.error(f"Invalid {key} in image_size")
                return False
        
        # Validate annotations
        if "annotations" not in bb or not isinstance(bb["annotations"], list) or not bb["annotations"]:
            logger.error("Invalid or empty annotations")
            return False
        
        # Validate bounding-box-metadata
        meta = entry["bounding-box-metadata"]
        required_meta = ["objects", "class-map", "type", "human-annotated", "creation-date", "job-name"]
        for key in required_meta:
            if key not in meta:
                logger.error(f"Missing {key} in metadata")
                return False
        
        if meta["type"] != "groundtruth/object-detection":
            logger.error("Invalid type in metadata")
            return False
        
        if len(meta["objects"]) != len(bb["annotations"]):
            logger.error("Number of objects doesn't match number of annotations")
            return False
        
        return True
        
    except Exception as e:
        logger.error(f"Error validating manifest format: {str(e)}")
        return False

def validate_manifest_entry(entry: Dict[str, Any]) -> bool:
    """
    Validate that a manifest entry has the correct format.
    Returns True if valid, False if not.
    """
    try:
        # Required top-level keys
        required_keys = ["source-ref", "bounding-box", "bounding-box-metadata"]
        for key in required_keys:
            if key not in entry:
                logger.error(f"Missing required key: {key}")
                return False
                
        # Validate source-ref
        if not isinstance(entry["source-ref"], str):
            logger.error("source-ref must be a string")
            return False
            
        if not entry["source-ref"].startswith("s3://"):
            logger.error("source-ref must start with s3://")
            return False
            
        # Clean up source-ref - ensure no double slashes (except after s3://) and proper format
        entry["source-ref"] = entry["source-ref"].replace("//", "/")
        # Make sure it starts with s3:// (with double slash)
        if entry["source-ref"].startswith("s3:/") and not entry["source-ref"].startswith("s3://"):
            entry["source-ref"] = "s3://" + entry["source-ref"][4:]
        # Fix any duplicate slashes in the path part (but not in the s3:// prefix)
        if '//' in entry["source-ref"][5:]:
            parts = entry["source-ref"].split('/', 3)
            if len(parts) >= 4:
                entry["source-ref"] = f"s3://{parts[2]}/{parts[3].replace('//', '/')}"
        
        # Validate bounding-box structure
        bb = entry["bounding-box"]
        if not isinstance(bb, dict):
            logger.error("bounding-box must be a dictionary")
            return False
            
        # Validate image_size
        if "image_size" not in bb or not isinstance(bb["image_size"], list) or len(bb["image_size"]) != 1:
            logger.error("bounding-box must contain exactly one image_size entry")
            return False
            
        size = bb["image_size"][0]
        for key in ["width", "height", "depth"]:
            if key not in size:
                logger.error(f"image_size missing {key}")
                return False
            try:
                size[key] = int(size[key])
                if size[key] <= 0:
                    logger.error(f"image_size {key} must be positive")
                    return False
            except (ValueError, TypeError):
                logger.error(f"image_size {key} must be an integer")
                return False
                
        # Validate annotations
        if "annotations" not in bb or not isinstance(bb["annotations"], list) or not bb["annotations"]:
            logger.error("bounding-box must contain non-empty annotations list")
            return False
            
        # Check for maximum number of annotations (AWS Rekognition limit)
        MAX_ANNOTATIONS = 50
        if len(bb["annotations"]) > MAX_ANNOTATIONS:
            logger.error(f"Too many annotations: {len(bb['annotations'])}. Maximum allowed is {MAX_ANNOTATIONS} to prevent ERROR_TOO_MANY_BOUNDING_BOXES")
            return False
            
        # Validate each annotation
        img_width = size["width"]
        img_height = size["height"]
        
        for i, ann in enumerate(bb["annotations"]):
            if not isinstance(ann, dict):
                logger.error(f"Annotation {i} must be a dictionary")
                return False
                
            # Required fields
            for key in ["class_id", "left", "top", "width", "height"]:
                if key not in ann:
                    logger.error(f"Annotation {i} missing {key}")
                    return False
                    
            # Validate class_id
            try:
                ann["class_id"] = int(ann["class_id"])
            except (ValueError, TypeError):
                logger.error(f"Annotation {i} class_id must be an integer")
                return False
                
            # Validate coordinates - now checking pixel values instead of normalized values
            for key in ["left", "top", "width", "height"]:
                try:
                    ann[key] = int(ann[key])
                    # For left and top, must be within image dimensions
                    if key in ["left", "top"]:
                        max_val = img_width if key == "left" else img_height
                        if not 0 <= ann[key] < max_val:
                            logger.error(f"Annotation {i} {key} must be between 0 and {max_val-1}")
                            return False
                    # For width and height, must be positive and fit within image bounds
                    else:
                        if ann[key] <= 0:
                            logger.error(f"Annotation {i} {key} must be positive")
                            return False
                except (ValueError, TypeError):
                    logger.error(f"Annotation {i} {key} must be an integer")
                    return False
                    
            # Validate box bounds
            if ann["left"] + ann["width"] > img_width:
                logger.error(f"Annotation {i} box extends beyond right edge of image (width {img_width}px)")
                return False
            if ann["top"] + ann["height"] > img_height:
                logger.error(f"Annotation {i} box extends beyond bottom edge of image (height {img_height}px)")
                return False
                
        # Validate metadata
        meta = entry["bounding-box-metadata"]
        if not isinstance(meta, dict):
            logger.error("bounding-box-metadata must be a dictionary")
            return False
            
        # Required metadata fields
        required_meta = ["objects", "class-map", "type", "human-annotated", "creation-date", "job-name"]
        for key in required_meta:
            if key not in meta:
                logger.error(f"Metadata missing {key}")
                return False
                
        # Validate specific fields
        if meta["type"] != "groundtruth/object-detection":
            logger.error("Metadata type must be groundtruth/object-detection")
            return False
            
        if not isinstance(meta["objects"], list) or len(meta["objects"]) != len(bb["annotations"]):
            logger.error("Metadata objects list must match number of annotations")
            return False
            
        # Validate class-map
        class_map = meta["class-map"]
        if not isinstance(class_map, dict):
            logger.error("class-map must be a dictionary")
            return False
            
        # Ensure all class_ids are in class-map
        used_class_ids = {str(ann["class_id"]) for ann in bb["annotations"]}
        missing_classes = used_class_ids - set(class_map.keys())
        if missing_classes:
            logger.error(f"Class IDs {missing_classes} not found in class-map")
            return False
            
        # Validate objects list
        for obj in meta["objects"]:
            if not isinstance(obj, dict) or "confidence" not in obj:
                logger.error("Each object must have a confidence value")
                return False
            try:
                obj["confidence"] = float(obj["confidence"])
                if not 0 <= obj["confidence"] <= 1:
                    logger.error("Confidence must be between 0 and 1")
                    return False
            except (ValueError, TypeError):
                logger.error("Confidence must be a number")
                return False
                
        return True
        
    except Exception as e:
        logger.error(f"Error validating manifest entry: {str(e)}")
        return False

def add_to_training_manifest(manifest_entry: Dict[str, Any], training_bucket: str) -> None:
    """
    Add entry to training manifest without first deleting the existing file.
    This prevents data loss and handles duplicate entries properly.
    """
    try:
        # Validate the entry
        if not validate_manifest_entry(manifest_entry):
            logger.error("Invalid manifest entry, skipping")
            logger.error(f"Invalid entry content: {json.dumps(manifest_entry, indent=2)}")
            raise ValueError("Invalid manifest entry")
        
        # Use correct manifest filename for AWS Rekognition
        manifest_key = "manifest.jsonl"
        
        # Get the source-ref to check for duplicates
        source_ref = manifest_entry.get("source-ref", "")
        if not source_ref:
            logger.error("Entry has no source-ref, cannot add to manifest")
            raise ValueError("Entry missing source-ref")
            
        logger.info(f"Adding entry to manifest for: {source_ref}")
        
        # Try to get existing manifest
        existing_entries = []
        existing_source_refs = set()  # Track source refs we've seen
        has_duplicate = False
        
        try:
            # Read existing manifest
            response = s3_client.get_object(Bucket=training_bucket, Key=manifest_key)
            content = response['Body'].read().decode('utf-8')
            
            # Log the size and line count of the existing manifest
            line_count = len(content.strip().split('\n')) if content.strip() else 0
            logger.info(f"Reading existing manifest with size {len(content)} bytes, {line_count} lines")
            
            # Check for duplicates and parse valid entries
            for i, line in enumerate(content.strip().split('\n'), 1):
                if not line.strip():
                    logger.warning(f"Line {i}: Empty line in manifest, skipping")
                    continue
                    
                try:
                    entry = json.loads(line)
                    entry_source_ref = entry.get('source-ref', '')
                    
                    if not entry_source_ref:
                        logger.warning(f"Line {i}: Entry has no source-ref, skipping")
                        continue
                        
                    # Log source_ref for debugging
                    logger.debug(f"Line {i}: Processing entry with source-ref: {entry_source_ref}")
                    
                    if entry_source_ref == source_ref:
                        # Skip duplicate entries - we'll add the new one at the end
                        has_duplicate = True
                        logger.info(f"Line {i}: Found duplicate entry for {source_ref}, will replace it")
                    elif entry_source_ref in existing_source_refs:
                        # Skip duplicate entries that we've already seen
                        logger.warning(f"Line {i}: Duplicate source-ref in manifest: {entry_source_ref}, keeping first occurrence")
                    else:
                        # Keep non-duplicate entries
                        existing_entries.append(entry)
                        existing_source_refs.add(entry_source_ref)
                        logger.debug(f"Line {i}: Kept entry for {entry_source_ref}")
                except json.JSONDecodeError as e:
                    logger.warning(f"Line {i}: Invalid JSON: {str(e)}, line content: {line[:100]}...")
                except Exception as e:
                    logger.warning(f"Line {i}: Error processing entry: {str(e)}")
                    
            logger.info(f"Read {len(existing_entries)} valid entries from manifest")
            
        except s3_client.exceptions.NoSuchKey:
            logger.info("No existing manifest found, will create new one")
        except Exception as e:
            logger.error(f"Error reading existing manifest: {str(e)}")
            # Continue with empty existing entries rather than failing
            logger.warning("Will proceed with empty existing entries")
        
        # Add the new entry to the list
        existing_entries.append(manifest_entry)
        existing_source_refs.add(source_ref)
        
        # Log all source refs for debugging
        logger.info(f"Final source refs in manifest: {sorted(existing_source_refs)}")
        
        # Create the manifest content in proper JSONL format
        manifest_content = ''
        for entry in existing_entries:
            manifest_content += json.dumps(entry, separators=(',', ':')) + '\n'
        
        # Log manifest statistics
        logger.info(f"Final manifest will contain {len(existing_entries)} entries, {len(manifest_content)} bytes")
        
        # Upload manifest with correct content type (WITHOUT deleting first)
        try:
            s3_client.put_object(
                Bucket=training_bucket,
                Key=manifest_key,
                Body=manifest_content,
                ContentType='application/x-amazon-s3-object-manifest-jsonl'
            )
            
            action = "updated" if has_duplicate else "added new"
            logger.info(f"Successfully {action} entry in manifest, now contains {len(existing_entries)} entries")
            
            # Verify by reading back the manifest
            verify_response = s3_client.get_object(Bucket=training_bucket, Key=manifest_key)
            verify_content = verify_response['Body'].read().decode('utf-8')
            verify_line_count = len(verify_content.strip().split('\n')) if verify_content.strip() else 0
            
            if verify_line_count != len(existing_entries):
                logger.warning(f"Verification warning: Expected {len(existing_entries)} entries, but read back {verify_line_count}")
            else:
                logger.info(f"Verified manifest contains {verify_line_count} entries as expected")
            
            # Verify the content type is correct
            content_type = verify_response.get('ContentType', 'unknown')
            if content_type != 'application/x-amazon-s3-object-manifest-jsonl':
                logger.warning(f"Incorrect content type detected: {content_type}, fixing...")
                s3_client.copy_object(
                    CopySource={'Bucket': training_bucket, 'Key': manifest_key},
                    Bucket=training_bucket,
                    Key=manifest_key,
                    ContentType='application/x-amazon-s3-object-manifest-jsonl',
                    MetadataDirective='REPLACE'
                )
                logger.info("Content type fixed")
        except Exception as e:
            logger.error(f"Error uploading manifest: {str(e)}")
            raise
            
    except Exception as e:
        logger.error(f"Error updating manifest: {str(e)}")
        raise 