src/main.py

import magic
import exifread
import xml.etree.ElementTree as ET
from PIL import Image
from PIL.ExifTags import TAGS, GPSTAGS
from PIL import ExifTags
import re
from typing import Dict, Any, Optional, List, Union, Tuple

from src.metadata_parser import MetadataParser
from src.helpers.makernote_parser import MakerNoteHelper

from src.helpers.image_utility_helper import load_image_resource
from src.helpers.logger import logger

# -------------------------
# Standalone functions
# -------------------------
def load_image_data(image_path: str, input_type: str = None) -> tuple:
    """
    Separate function to load image raw bytes and PIL object.
    Supports URL, Base64 and Local Path.
    """
    return load_image_resource(image_path, input_type)


class ImageMetadataDetector:
    def __init__(self):
        self.image_path = None
        self.image = None
        self.raw_bytes = None

    # -------------------------
    # Public API
    # -------------------------
    def predict(self, image_path: str, input_type: str = None) -> dict:
        """
        Main entry point for the FastAPI app.
        """
        result = self.detect(image_path, input_type)
        structured = MetadataParser.parse(result)

        return structured.model_dump(mode="json")

    def detect(self, image_path: str, input_type: str = None) -> dict:
        """
        Main entry point.
        Detects file type, metadata types, and parses all metadata.
        """
        logger.info(f"Detecting metadata for {image_path} (type: {input_type})")
        self.image_path = image_path
        self.load_image(image_path, input_type)

        result = {
            "file_path": self.image_path,
            "mime_type": self._detect_mime_type(),
            "image_format": self.image.format if self.image else None,
            "size": self.image.size if self.image else None,
            "size_bytes": len(self.raw_bytes),
            "color_space": self._detect_color_space(),
            "metadata_types": {},
            "metadata": {}
        }
        # result["metadata"]["additional_metadata"] = {}

        meta_types = self._detect_metadata_types()
        result["metadata_types"] = meta_types

        if meta_types["exif"]:
            exif_res = self._parse_exif()
            result["metadata"]["exif"] = exif_res
            
            # Check for MakerNote specific parsing (Apple, etc)
            if "MakerNote" in exif_res or "MakerNote" in exif_res.get("ExifIFD", {}):
                result["metadata"]["makernote"] = self._parse_makernote(exif_res)

        if meta_types["iptc"]:
            result["metadata"]["iptc"] = self._parse_iptc()

        if meta_types["xmp"]:
            result["metadata"]["xmp"] = self._parse_xmp()

        # if meta_types["icc"]:
        #     result["metadata"]["icc"] = self._parse_icc()

        if meta_types["png_text"]:
            raw, parsed_xmp = self._parse_png_text()
            result["metadata"]["png_text"] = raw
            if parsed_xmp:
                result["metadata"]["png_text_parsed"] = {
                    "xmp": parsed_xmp
                }



        # return result
        return self._make_json_safe(result)



    # -------------------------
    # Internal helpers
    # -------------------------
    def load_image(self, image_path: str, input_type: str = None):
        """
        Loads the image and stores it in the instance.
        """
        self.image_path = image_path
        self.image, self.raw_bytes = load_image_data(image_path, input_type)
            
    
    def _convert_gps(self, gps):
        def to_deg(value):
            d, m, s = value
            return d[0]/d[1] + (m[0]/m[1])/60 + (s[0]/s[1])/3600

        lat = to_deg(gps["GPSLatitude"])
        if gps.get("GPSLatitudeRef") == "S":
            lat = -lat

        lon = to_deg(gps["GPSLongitude"])
        if gps.get("GPSLongitudeRef") == "W":
            lon = -lon

        return {"lat": lat, "lon": lon}

    # -------------------------
    # Detection
    # -------------------------
    def _detect_mime_type(self) -> str:
        mime = magic.Magic(mime=True)
        return mime.from_buffer(self.raw_bytes)

    def _detect_color_space(self) -> Optional[str]:
        if not self.image:
            return None

        # 1. Check EXIF
        try:
            exif = self.image.getexif()
            if exif:
                # 0xA001 is ColorSpace tag
                cs = exif.get(0xA001)
                if cs == 1: return "sRGB"
                if cs == 2: return "Adobe RGB"
                if cs == 65535: return "Uncalibrated"
        except Exception:
            pass

        # 2. Check PIL info (PNG sRGB chunk or ICC Profile)
        if "srgb" in self.image.info:
            return "sRGB"
        
        if "icc_profile" in self.image.info:
            return "ICC Profile"

        return None

    def _detect_metadata_types(self) -> dict:
        meta = {
            "exif": False,
            "iptc": False,
            "xmp": False,
            "png_text": False
        }

        # EXIF
        if self.image and hasattr(self.image, "getexif"):
            try:
                if self.image.getexif():
                    meta["exif"] = True
            except Exception:
                pass

        # PNG text chunks
        if self.image and self.image.format == "PNG" and self.image.info:
            meta["png_text"] = True

        # XMP (embedded XML)
        if b"<x:xmpmeta" in self.raw_bytes:
            meta["xmp"] = True

        # IPTC (Photoshop marker)
        if b"Photoshop 3.0" in self.raw_bytes:
            meta["iptc"] = True
        

        return meta

    # -------------------------
    # Metadata parsers
    # -------------------------
    
    def _parse_exif(self) -> dict:
        """
        Extract full EXIF including:
        - IFD0
        - ExifIFD
        - GPSInfo
        - Interoperability IFD
        """
        exif_data = {}

        exif = self.image.getexif()

        # ---- IFD0 + ExifIFD ----
        for tag_id, value in exif.items():
            tag_name = TAGS.get(tag_id, tag_id)

            # GPS handled separately
            if tag_name == "GPSInfo":
                continue

            exif_data[tag_name] = self._serialize(value)

        # add additional tags
        exif_data["additional"] = {}

        # ---- GPS IFD ----
        gps_info = exif.get_ifd(ExifTags.IFD.GPSInfo) #0x8825
        if gps_info:
            gps_data = {}
            for tag_id, value in gps_info.items():
                tag_name = GPSTAGS.get(tag_id, tag_id)
                gps_data[tag_name] = self._serialize(value)
            exif_data["GPSInfo"] = gps_data

        # ---- Exif SubIFD ----
        exif_ifd = exif.get_ifd(ExifTags.IFD.Exif)  # ExifOffset 0x8769
        if exif_ifd:
            exif_sub = {}
            for tag_id, value in exif_ifd.items():
                tag_name = TAGS.get(tag_id, tag_id)
                exif_sub[tag_name] = self._serialize(value)
            exif_data["ExifIFD"] = exif_sub

        # ---- Interoperability IFD ----
        interop_ifd = exif.get_ifd(ExifTags.IFD.IFD1)  # InteropIFD 0xA005
        if interop_ifd:
            interop = {}
            for tag_id, value in interop_ifd.items():
                tag_name = TAGS.get(tag_id, tag_id)
                interop[tag_name] = self._serialize(value)
            exif_data["InteropIFD"] = interop

        # 3. Check for depth map information
        if "depth_images" in self.image.info and self.image.info["depth_images"]:
            logger.info("Depth map found.")
            exif_data["additional"]["contains_depth_map"] = True
            exif_data["additional"]["depth_map"] = self.image.info["depth_images"][0]
            # exif_data["InteropIFD"]["depth_map"] = "depth_map"

            # Usually, the first depth image is the one needed
            # depth_image = image.info["depth_images"][0]
            
            # Save or convert the depth map to a numpy array for processing
            # depth_image.save(output_depth_path)
            # print(f"Depth map saved to {output_depth_path}")
            
            # To convert to numpy for CV applications
            # depth_array = np.array(depth_image)
            # return depth_array
        # else:
        #     print("No depth map found in this HEIC image.")
        #     return None

        return exif_data

    def _parse_iptc(self) -> dict:
        iptc_data = {}

        with open(self.image_path, "rb") as f:
            tags = exifread.process_file(f, details=False)

        for tag, value in tags.items():
            if tag.startswith("IPTC"):
                iptc_data[tag] = str(value)

        return iptc_data

    def _parse_xmp(self) -> dict:
        start = self.raw_bytes.find(b"<x:xmpmeta")
        end = self.raw_bytes.find(b"</x:xmpmeta>")

        if start == -1 or end == -1:
            return {}

        end += len(b"</x:xmpmeta>")
        xmp_bytes = self.raw_bytes[start:end]

        try:
            root = ET.fromstring(xmp_bytes)
            return self._xml_to_dict(root)
        except Exception:
            return {"raw_xmp": xmp_bytes.decode(errors="ignore")}

    def _parse_makernote(self, exif_res: dict) -> dict:
        """
        Extract and map MakerNote tags using MakerNoteHelper for deeper parsing.
        """
        make = exif_res.get("Make", "Unknown")
        return MakerNoteHelper.parse_makernotes(self.image_path, make, self._serialize, exif_res)

    # def _parse_png_text(self) -> dict:
    def _parse_png_text(self):
        raw = dict(self.image.info)
        parsed_xmp = {}

        for key, value in raw.items():
            if isinstance(value, str) and "<x:xmpmeta" in value:
                parsed_xmp[key] = self._parse_xmp_string(value)

        return raw, parsed_xmp

    # def _parse_icc(self):
    #     icc_data = {}
    #     icc_profile_bytes = self.image.info.get("icc_profile")
        

    # -------------------------
    # Utilities
    # -------------------------
    def _xml_to_dict(self, element):
        """
        Recursively convert XML tree to dict.
        """
        data = {}
        for child in element:
            tag = child.tag.split("}")[-1]  # strip namespace
            data[tag] = self._xml_to_dict(child) if list(child) else child.text
        return data

    def _serialize(self, value):
        """
        Convert EXIF values to JSON-safe formats.
        """
        if isinstance(value, bytes):
            return value.decode(errors="ignore")
        if isinstance(value, (list, tuple)):
            return [self._serialize(v) for v in value]
        return value
    
    def _make_json_safe(self, obj):
        """
        Recursively convert any object into JSON-serializable types.
        """
        if isinstance(obj, bytes):
            return obj.decode(errors="ignore")

        if isinstance(obj, dict):
            return {str(k): self._make_json_safe(v) for k, v in obj.items()}

        if isinstance(obj, (list, tuple, set)):
            return [self._make_json_safe(v) for v in obj]

        if isinstance(obj, (int, float, str, bool)) or obj is None:
            if isinstance(obj, float):
                import math
                if math.isnan(obj) or math.isinf(obj):
                    return None
            if isinstance(obj, str) and obj.lower() == "nan":
                return None
            return obj

        # Fallback for unknown objects (e.g. PIL types)
        res = str(obj)
        if res.lower() == "nan":
            return None
        return res
    
    def _parse_xmp_string(self, xmp_str: str) -> dict:
        """
        Parse XMP XML string into structured dict.
        """
        try:
            # Remove xpacket headers
            xmp_str = re.sub(r"<\?xpacket.*?\?>", "", xmp_str, flags=re.DOTALL).strip()

            root = ET.fromstring(xmp_str)

            parsed = {}

            # Traverse RDF descriptions
            for elem in root.iter():
                tag = elem.tag.split("}")[-1]

                # Attributes (THIS is where most XMP data lives)
                if elem.attrib:
                    for k, v in elem.attrib.items():
                        key = k.split("}")[-1]
                        parsed[key] = v

                # Text nodes (rare in XMP)
                if elem.text and elem.text.strip():
                    parsed[tag] = elem.text.strip()

            return parsed
        except Exception as e:
            return {
                "_error": "XMP parse failed",
                "_raw": xmp_str[:1000]
            }


if __name__ == "__main__":
    import json
    import os
    import glob
    from pillow_heif import register_heif_opener

    register_heif_opener()

    # def natural_sort_key(path):
    #     return [int(t) if t.isdigit() else t.lower()
    #             for t in re.split(r'(\d+)', path)]
    # images = sorted(glob.glob("inputs/tests/*.*"),key=natural_sort_key)
    # images = ["inputs/tests/12_ip17pm.HEIC"]
    # images = ["inputs/tests/3_photo.jpg"]
    # images = ["inputs/tests/7_ai.png"]
    # images = ["inputs/tests/9_ai.jpeg"]
    images = ["/Users/rohan/Developer/2026/myaidetector/datasets/real/real (2).HEIC"]

    all_results = {}
    all_parsed_results = {}
    
      
    for image_path in images:
        if not os.path.isfile(image_path):
            logger.warning(f"Skipping {image_path} because it does not exist")
            continue
        detector = ImageMetadataDetector()
        result = detector.detect(image_path, "local_path")
        print(result)
        all_results[image_path] = result

        structured = MetadataParser.parse(result)
        all_parsed_results[image_path] = structured.model_dump(mode="json")
        # print(f"\n--- Metadata for {image_path} ---\n")
        # print(json.dumps(result, indent=2))
        # print(f"\n--- Formatted ---\n")
        # print(structured.model_dump_json(indent=2))

    # write once, after the loop
    with open("outputs/imageMetadataData.json", "w", encoding="utf-8") as f:
        json.dump(all_results, f, ensure_ascii=False, indent=4, sort_keys=True)
    logger.info("All metadata saved to outputs/imageMetadataData.json")
    
    with open("outputs/imageParsedData.json", "w", encoding="utf-8") as f:
        json.dump(all_parsed_results, f, ensure_ascii=False, indent=4)
    logger.info("All parsed metadata saved to outputs/imageParsedData.json")
    print(f"\n--- All parsed metadata saved to outputs/imageParsedData.json ---\n")