#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import argparse
import os
from pathlib import Path
from datetime import datetime, timezone, timedelta
from typing import Optional, Tuple, List

from PIL import Image, ExifTags
import xml.etree.ElementTree as ET

# Mapy tagów EXIF
TAGS = {v: k for k, v in ExifTags.TAGS.items()}
GPSTAGS = {v: k for k, v in ExifTags.GPSTAGS.items()}

# Akceptowane rozszerzenia zdjęć
IMAGE_EXTS = {".jpg", ".jpeg", ".JPG", ".JPEG"}

def dms_to_decimal(dms, ref) -> Optional[float]:
    """
    Konwersja DMS (na ułamki) -> stopnie dziesiętne.
    dms: np. ((52,1), (13,1), (1234,100))  => 52° 13' 12.34"
    ref: 'N'/'S' lub 'E'/'W'
    """
    try:
        deg = dms[0][0] / dms[0][1]
        minutes = dms[1][0] / dms[1][1]
        seconds = dms[2][0] / dms[2][1]
        dec = deg + minutes / 60.0 + seconds / 3600.0
        if ref in ("S", "W"):
            dec = -dec
        return dec
    except Exception:
        return None

def parse_exif_datetime(raw_dt: Optional[str], subsec: Optional[str]) -> Optional[datetime]:
    """
    Parsuje DateTimeOriginal w formacie 'YYYY:MM:DD HH:MM:SS' i składa z sub-sekundami.
    Zwraca naive datetime (bez strefy) – później ewentualnie przeliczane wg --tz.
    """
    if not raw_dt:
        return None
    try:
        base = datetime.strptime(raw_dt, "%Y:%m:%d %H:%M:%S")
        if subsec and subsec.isdigit():
            # Subsec może mieć różną długość – skracamy/padding do mikrosekund
            micro = int((subsec + "000000")[:6])
            base = base.replace(microsecond=micro)
        return base
    except Exception:
        return None

def apply_tz_and_to_utc(dt: datetime, tz_offset: Optional[str]) -> Optional[datetime]:
    """
    Jeśli podano --tz w formacie +HH:MM lub -HH:MM, interpretuj naive dt w tej strefie,
    a następnie zwróć dt w UTC.
    Jeśli --tz nie podano, zwróć None (żeby pominąć <time> albo zapisać bez gwarancji poprawnej TZ).
    """
    if tz_offset is None:
        return None
    try:
        sign = 1 if tz_offset.startswith("+") else -1
        hh, mm = tz_offset[1:].split(":")
        offset = timedelta(hours=int(hh), minutes=int(mm))
        tzinfo = timezone(sign * offset)
        aware = dt.replace(tzinfo=tzinfo)
        return aware.astimezone(timezone.utc)
    except Exception:
        return None

def extract_exif_gps_and_time(img_path: Path) -> Tuple[Optional[float], Optional[float], Optional[float], Optional[datetime], Optional[datetime]]:
    """
    Zwraca: (lat, lon, ele, naive_time, utc_time_if_available)
    """
    try:
        with Image.open(img_path) as im:
            exif = im._getexif() or {}
    except Exception:
        return None, None, None, None, None

    # Zamiana kluczy numerycznych na nazwy
    exif_readable = {}
    for tag, val in exif.items():
        name = ExifTags.TAGS.get(tag, tag)
        exif_readable[name] = val

    gps_info = exif_readable.get("GPSInfo")
    lat = lon = ele = None
    if isinstance(gps_info, dict):
        gps = {}
        for k, v in gps_info.items():
            gps_name = ExifTags.GPSTAGS.get(k, k)
            gps[gps_name] = v

        lat = dms_to_decimal(gps.get("GPSLatitude"), gps.get("GPSLatitudeRef", "N"))
        lon = dms_to_decimal(gps.get("GPSLongitude"), gps.get("GPSLongitudeRef", "E"))

        # Wysokość (opcjonalna)
        if gps.get("GPSAltitude") and isinstance(gps["GPSAltitude"], tuple):
            try:
                ele = gps["GPSAltitude"][0] / gps["GPSAltitude"][1]
                if gps.get("GPSAltitudeRef") == 1:
                    ele = -ele
            except Exception:
                ele = None

    # Czas
    raw_dt = exif_readable.get("DateTimeOriginal") or exif_readable.get("DateTime") or exif_readable.get("DateTimeDigitized")
    # Subsekundy bywają w różnych tagach:
    subsec = exif_readable.get("SubsecTimeOriginal") or exif_readable.get("SubsecTime") or exif_readable.get("SubSecTimeOriginal")

    naive_dt = parse_exif_datetime(raw_dt, subsec)
    return lat, lon, ele, naive_dt, None  # utc_time wypełnimy później wg --tz

def collect_images(input_path: Path) -> List[Path]:
    files = []
    if input_path.is_dir():
        for root, _, fns in os.walk(input_path):
            for fn in fns:
                p = Path(root) / fn
                if p.suffix in IMAGE_EXTS:
                    files.append(p)
    else:
        if input_path.suffix in IMAGE_EXTS:
            files.append(input_path)
    return files

def generate_gpx(points: List[dict], creator: str = "exif2gpx") -> ET.ElementTree:
    ET.register_namespace("", "http://www.topografix.com/GPX/1/1")
    ET.register_namespace("xsi", "http://www.w3.org/2001/XMLSchema-instance")
    gpx = ET.Element(
        "gpx",
        {
            "version": "1.1",
            "creator": creator,
            "xmlns": "http://www.topografix.com/GPX/1/1",
            "xmlns:xsi": "http://www.w3.org/2001/XMLSchema-instance",
            "xsi:schemaLocation": "http://www.topografix.com/GPX/1/1 http://www.topografix.com/GPX/1/1/gpx.xsd",
        },
    )
    trk = ET.SubElement(gpx, "trk")
    name = ET.SubElement(trk, "name")
    name.text = "Track from EXIF"
    seg = ET.SubElement(trk, "trkseg")

    for p in points:
        trkpt = ET.SubElement(seg, "trkpt", {"lat": f"{p['lat']:.7f}", "lon": f"{p['lon']:.7f}"})
        if p.get("ele") is not None:
            ele = ET.SubElement(trkpt, "ele")
            ele.text = f"{p['ele']:.2f}"
        if p.get("time"):
            # GPX wymaga ISO 8601 w UTC z 'Z'
            t = ET.SubElement(trkpt, "time")
            t.text = p["time"].strftime("%Y-%m-%dT%H:%M:%S.%fZ").rstrip("0").rstrip(".") + "Z"
        # Dodatkowo można dodać nazwę pliku jako <desc>
        desc = ET.SubElement(trkpt, "desc")
        desc.text = p.get("source", "")

    return ET.ElementTree(gpx)

def main():
    ap = argparse.ArgumentParser(description="Z EXIF zdjęć (JPG) generuje GPX z trackiem.")
    ap.add_argument("input", help="Plik JPG lub folder ze zdjęciami")
    ap.add_argument("output", help="Ścieżka wyjściowa pliku GPX")
    ap.add_argument("--tz", help="Strefa czasowa EXIF w formacie +HH:MM lub -HH:MM (np. +02:00)", default=None)
    args = ap.parse_args()

    src = Path(args.input)
    out = Path(args.output)

    files = collect_images(src)
    if not files:
        print("Brak zdjęć JPG/JPEG w podanej lokalizacji.")
        return

    points = []
    for fp in files:
        lat, lon, ele, naive_dt, _ = extract_exif_gps_and_time(fp)
        if lat is None or lon is None:
            continue

        # Wyznacz czas w UTC jeśli podano --tz
        time_utc = None
        if naive_dt is not None and args.tz:
            time_utc = apply_tz_and_to_utc(naive_dt, args.tz)

        points.append({
            "lat": lat,
            "lon": lon,
            "ele": ele,
            "naive_time": naive_dt,
            "time": time_utc,     # tylko jeśli mamy --tz
            "source": fp.name,
        })

    if not points:
        print("Nie znaleziono zdjęć z danymi GPS.")
        return

    # Sortowanie: najpierw po czasie jeśli dostępny; potem po nazwie pliku
    def sort_key(p):
        # preferuj p['time'] (UTC) gdy jest; w przeciwnym razie naive_time; w ostateczności nazwa
        return (
            p["time"] if p["time"] is not None else
            p["naive_time"] if p["naive_time"] is not None else
            datetime.max
        , p["source"])

    points.sort(key=sort_key)

    tree = generate_gpx(points)
    out.parent.mkdir(parents=True, exist_ok=True)
    tree.write(out, encoding="utf-8", xml_declaration=True)
    print(f"Zapisano GPX: {out}")

if __name__ == "__main__":
    main()