#!/opt/scripts/unmanic/.venv/bin/python3
# -*- coding: utf-8 -*-

import os, sys, re, json, sqlite3, argparse, time
from pathlib import Path
from datetime import datetime, timezone
import urllib.request, urllib.parse

def _iso_now():
    return datetime.now(timezone.utc).isoformat(timespec="seconds")

def _log(level, msg):
    print(f"[{_iso_now()}] {level}: {msg}")

# ---------- .env pre-parse (no external deps) ----------
def _load_env_file(path: str) -> bool:
    p = Path(path)
    if not p.exists():
        _log("WARNING", f".env not found path={path}")
        return False
    loaded = 0
    for raw in p.read_text(encoding="utf-8").splitlines():
        line = raw.strip()
        if not line or line.startswith("#") or "=" not in line:
            continue
        k, v = line.split("=", 1)
        k = k.strip()
        v = v.strip().strip('"').strip("'")
        os.environ[k] = v
        loaded += 1
    if "SONARR_API_KEY" not in os.environ and "SONARR_APIKEY" in os.environ:
        os.environ["SONARR_API_KEY"] = os.environ["SONARR_APIKEY"]
    if "RADARR_API_KEY" not in os.environ and "RADARR_APIKEY" in os.environ:
        os.environ["RADARR_API_KEY"] = os.environ["RADARR_APIKEY"]
    _log("INFO", f".env loaded from {path} ({loaded} keys)")
    return True

def _preparse_env(argv):
    env_path = None
    new_argv = []
    i = 0
    while i < len(argv):
        a = argv[i]
        if a == "--env" and i + 1 < len(argv):
            env_path = argv[i + 1]; i += 2; continue
        elif a.startswith("--env="):
            env_path = a.split("=", 1)[1]; i += 1; continue
        new_argv.append(a); i += 1
    if env_path:
        if not _load_env_file(env_path):
            _log("WARNING", f".env not loaded path={env_path}")
    return new_argv

sys.argv = _preparse_env(sys.argv)

# ---------- DB ----------
SCHEMA = """
PRAGMA journal_mode=WAL;

CREATE TABLE IF NOT EXISTS meta (
  key TEXT PRIMARY KEY,
  value TEXT
);

CREATE TABLE IF NOT EXISTS series (
  imdb_id TEXT PRIMARY KEY,
  series_path TEXT NOT NULL,
  last_applied TEXT
);

CREATE TABLE IF NOT EXISTS episode_dates (
  imdb_id TEXT NOT NULL,
  season INTEGER NOT NULL,
  episode INTEGER NOT NULL,
  aired TEXT,
  dateadded TEXT,
  source TEXT,
  PRIMARY KEY (imdb_id, season, episode)
);
"""

EXPECTED_SERIES_COLS = {"imdb_id","series_path","last_applied"}
EXPECTED_EP_COLS     = {"imdb_id","season","episode","aired","dateadded","source"}

def _table_cols(conn, table):
    cur = conn.execute(f"PRAGMA table_info({table});")
    return {r[1] for r in cur.fetchall()}

def _ensure_schema(conn, allow_drop=True):
    conn.executescript(SCHEMA)
    ok = True
    try:
        sc = _table_cols(conn, "series")
        ec = _table_cols(conn, "episode_dates")
        if sc != EXPECTED_SERIES_COLS or ec != EXPECTED_EP_COLS:
            ok = False
    except sqlite3.OperationalError:
        ok = False
    if not ok and allow_drop:
        _log("WARNING","Schema mismatch detected — dropping and recreating tables.")
        conn.executescript("""
        DROP TABLE IF EXISTS episode_dates;
        DROP TABLE IF EXISTS series;
        """)
        conn.executescript(SCHEMA)
        # recheck
        sc = _table_cols(conn, "series")
        ec = _table_cols(conn, "episode_dates")
        if sc != EXPECTED_SERIES_COLS or ec != EXPECTED_EP_COLS:
            raise RuntimeError("Failed to create expected schema.")

def db_connect(db_path: Path) -> sqlite3.Connection:
    db_path.parent.mkdir(parents=True, exist_ok=True)
    conn = sqlite3.connect(str(db_path))
    conn.execute("PRAGMA foreign_keys=ON;")
    _ensure_schema(conn, allow_drop=True)
    return conn

# ---------- HTTP/helpers ----------
def http_get_json(url, headers=None, timeout=45):
    req = urllib.request.Request(url, headers=headers or {})
    with urllib.request.urlopen(req, timeout=timeout) as resp:
        data = resp.read()
        if not data: return None
        try: return json.loads(data.decode("utf-8"))
        except Exception: return None

def to_utc_iso(s):
    if not s: return None
    try:
        if len(s)==10 and s[4]=="-" and s[7]=="-":
            return datetime.fromisoformat(s).replace(tzinfo=timezone.utc).isoformat(timespec="seconds")
        s2 = s.replace("Z","+00:00")
        dt = datetime.fromisoformat(s2)
        if dt.tzinfo is None: dt = dt.replace(tzinfo=timezone.utc)
        else: dt = dt.astimezone(timezone.utc)
        return dt.isoformat(timespec="seconds")
    except Exception:
        return None

def iso_from_epoch(epoch):
    return datetime.fromtimestamp(epoch, tz=timezone.utc).isoformat(timespec="seconds")

def file_mtime_iso(p: Path):
    try: return iso_from_epoch(p.stat().st_mtime)
    except Exception: return None

# ---------- clients ----------
class SonarrClient:
    def __init__(self):
        self.base = os.environ.get("SONARR_URL","").rstrip("/")
        self.key  = os.environ.get("SONARR_API_KEY","")
        self.timeout = int(os.environ.get("SONARR_TIMEOUT","45") or "45")
        self.retries = int(os.environ.get("SONARR_RETRIES","3") or "3")
        self.enabled = bool(self.base and self.key)
    def _get(self, path, params=None):
        if not self.enabled: return None
        q = f"{self.base}/api/v3{path}"
        if params: q += "?"+urllib.parse.urlencode(params)
        headers = {"X-Api-Key": self.key}
        for _ in range(self.retries):
            j = http_get_json(q, headers=headers, timeout=self.timeout)
            if j is not None: return j
            time.sleep(0.3)
        return None
    def series_by_imdb(self, imdb_id):
        j = self._get("/series/lookup", {"term": f"imdbid:{imdb_id}"})
        if not j: return None
        for s in j:
            if (s.get("imdbId") or "").lower() == imdb_id.lower():
                return s
        return None
    def episodes_for_series(self, series_id):
        return self._get("/episode", {"seriesId": series_id}) or []
    def episode_file(self, episode_file_id):
        return self._get(f"/episodefile/{episode_file_id}")

class TMDBClient:
    def __init__(self):
        self.api_key = os.environ.get("TMDB_API_KEY","")
        self.timeout = int(os.environ.get("TMDB_TIMEOUT","45") or "45")
        self.retries = int(os.environ.get("TMDB_RETRIES","3") or "3")
        self.enabled = bool(self.api_key)
    def _get(self, path, params=None):
        if not self.enabled: return None
        base = "https://api.themoviedb.org/3"
        params = {**(params or {}), "api_key": self.api_key}
        url = f"{base}{path}?{urllib.parse.urlencode(params)}"
        for _ in range(self.retries):
            j = http_get_json(url, timeout=self.timeout)
            if j is not None: return j
            time.sleep(0.3)
        return None
    def tv_id_from_imdb(self, imdb_id):
        j = self._get(f"/find/{imdb_id}", {"external_source":"imdb_id"})
        if not j: return None
        tv = (j.get("tv_results") or [])
        if tv: return tv[0].get("id")
        return None
    def season_airdates(self, tv_id, season_number):
        j = self._get(f"/tv/{tv_id}/season/{season_number}")
        res = {}
        if not j: return res
        for ep in (j.get("episodes") or []):
            num = ep.get("episode_number")
            air = ep.get("air_date")
            if isinstance(num,int) and air:
                res[num] = air
        return res

class OMDbClient:
    def __init__(self):
        self.api_key = os.environ.get("OMDB_API_KEY","")
        self.enabled = bool(self.api_key)
    def season_airdates(self, imdb_id, season_number):
        if not self.enabled: return {}
        url = "http://www.omdbapi.com/?" + urllib.parse.urlencode(
            {"i":imdb_id,"Season":str(season_number),"apikey":self.api_key}
        )
        j = http_get_json(url, timeout=45)
        res = {}
        if not j or j.get("Response")!="True": return res
        for ep in j.get("Episodes",[]):
            try: num = int(ep.get("Episode"))
            except Exception: continue
            rel = ep.get("Released")
            if rel and rel!="N/A": res[num] = rel
        return res

# ---------- NFO / FS ----------
LONG_NFO_RE = re.compile(r".*-S\d{2}E\d{2}-.*\.nfo$", re.IGNORECASE)
SHORT_NFO_RE = re.compile(r"^S(\d{2})E(\d{2})\.nfo$", re.IGNORECASE)
IMDB_TAG_RE  = re.compile(r"\[imdb-(tt\d+)\]", re.IGNORECASE)

def remove_long_nfos(season_dir: Path) -> int:
    count = 0
    for f in season_dir.glob("*.nfo"):
        if LONG_NFO_RE.match(f.name):
            try: f.unlink(); count += 1
            except Exception: pass
    return count

def resolve_season_dir(series_dir: Path, season_num: int) -> Path:
    # prefer existing folders
    by_number = series_dir / f"Season {season_num:02d}"
    if by_number.exists(): return by_number
    if season_num == 0:
        specials = series_dir / "Specials"
        if specials.exists(): return specials
    # fallback: create standard folder
    try:
        (by_number).mkdir(parents=True, exist_ok=True)
        return by_number
    except Exception:
        # last ditch: specials for season 0
        if season_num == 0:
            specials = series_dir / "Specials"
            specials.mkdir(parents=True, exist_ok=True)
            return specials
        return by_number

def write_episode_nfo(series_dir: Path, season_num: int, ep_num: int, aired: str, dateadded: str, source: str):
    season_dir = resolve_season_dir(series_dir, season_num)
    name = f"S{season_num:02d}E{ep_num:02d}.nfo"
    dst = season_dir / name
    xml = [
        '<?xml version="1.0" encoding="UTF-8" standalone="yes"?>',
        "<episodedetails>",
        f"  <season>{season_num}</season>",
        f"  <episode>{ep_num}</episode>",
        f"  <aired>{(aired or '').split('T')[0]}</aired>",
    ]
    if dateadded:
        xml.append(f"  <dateadded>{dateadded}</dateadded>")
        xml.append(f"  <!-- source: {source or ''} -->")
    xml.append("</episodedetails>\n")
    try:
        dst.write_text("\n".join(xml), encoding="utf-8")
    except Exception as e:
        _log("ERROR", f"Failed writing {dst}: {e}")

def ensure_season_nfo(series_dir: Path, season_num: int):
    season_dir = resolve_season_dir(series_dir, season_num)
    p = season_dir / "season.nfo"
    if p.exists(): return
    xml = f"""<?xml version="1.0" encoding="UTF-8" standalone="yes"?>
<season>
  <seasonnumber>{season_num}</seasonnumber>
</season>
"""
    try: p.write_text(xml, encoding="utf-8")
    except Exception as e: _log("WARNING", f"Could not write season.nfo in {season_dir}: {e}")

def maybe_tvshow_nfo(series_dir: Path, imdb_id: str):
    p = series_dir / "tvshow.nfo"
    if p.exists(): return
    xml = f"""<?xml version="1.0" encoding="UTF-8" standalone="yes"?>
<tvshow>
  <id>{imdb_id}</id>
  <uniqueid type="imdb" default="true">{imdb_id}</uniqueid>
</tvshow>
"""
    try: p.write_text(xml, encoding="utf-8")
    except Exception as e: _log("WARNING", f"Could not create tvshow.nfo in {series_dir}: {e}")

def set_dir_mtime(path: Path, iso_timestamp: str):
    try:
        ts = datetime.fromisoformat(iso_timestamp.replace("Z","+00:00")).timestamp()
        os.utime(path, (ts, ts), follow_symlinks=False)
    except Exception as e:
        _log("WARNING", f"Failed to set mtime on {path}: {e}")

def set_media_mtimes(season_dir: Path, iso_timestamp: str):
    try: ts = datetime.fromisoformat(iso_timestamp.replace("Z","+00:00")).timestamp()
    except Exception as e:
        _log("WARNING", f"Bad timestamp for media mtimes in {season_dir}: {e}")
        return
    exts = (".mkv",".mp4",".avi",".mov",".m4v",".srt",".ass",".vtt")
    for f in season_dir.iterdir():
        if f.suffix.lower() in exts:
            try: os.utime(f, (ts, ts), follow_symlinks=False)
            except Exception: pass

# ---------- Core ----------
def parse_imdb_from_path(p: Path):
    m = IMDB_TAG_RE.search(p.name)
    return m.group(1).lower() if m else None

def is_series_dir(p: Path) -> bool:
    return p.is_dir() and parse_imdb_from_path(p) is not None

def gather_episode_dates(series_dir: Path, imdb_id: str, args):
    sonarr = SonarrClient()
    tmdb   = TMDBClient()
    omdb   = OMDbClient()

    out = {}

    if sonarr.enabled:
        s = sonarr.series_by_imdb(imdb_id)
        if s:
            series_id = s.get("id")
            eps = sonarr.episodes_for_series(series_id)
            if isinstance(eps, list):
                for ep in eps:
                    sn = ep.get("seasonNumber"); en = ep.get("episodeNumber")
                    if not isinstance(sn,int) or not isinstance(en,int): continue
                    air_utc  = to_utc_iso(ep.get("airDateUtc"))
                    date_add = None; src = None
                    efid = ep.get("episodeFileId")
                    if efid:
                        fobj = sonarr.episode_file(efid)
                        if fobj and fobj.get("dateAdded"):
                            date_add = to_utc_iso(fobj.get("dateAdded")); src = "sonarr:episodeFile.dateAdded"
                    if not date_add and air_utc:
                        date_add = air_utc; src = "sonarr:episode.airDateUtc"
                    if air_utc or date_add:
                        out[(sn,en)] = (air_utc, date_add, src)

    if tmdb.enabled:
        tv_id = tmdb.tv_id_from_imdb(imdb_id)
        if tv_id:
            seasons = sorted({int(p.name.split()[-1]) for p in series_dir.iterdir()
                              if p.is_dir() and (p.name.lower().startswith("season ") or p.name.lower()=="specials")
                              and (p.name.lower()=="specials" or p.name.split()[-1].isdigit())})
            # normalize specials to 0
            seasons = [0 if (series_dir/f"Season {s:02d}").name.lower()=="specials" else s for s in seasons]
            for sn in seasons:
                tm = tmdb.season_airdates(tv_id, sn)
                for en, air in tm.items():
                    key = (sn,en)
                    if key not in out or not out[key][1]:
                        air_iso = to_utc_iso(air)
                        out[key] = (air_iso, air_iso, "tmdb:air_date")

    if omdb.enabled:
        seasons = []
        for p in series_dir.iterdir():
            if not p.is_dir(): continue
            name = p.name.lower()
            if name=="specials": seasons.append(0)
            elif name.startswith("season ") and p.name.split()[-1].isdigit():
                seasons.append(int(p.name.split()[-1]))
        for sn in sorted(set(seasons)):
            m = omdb.season_airdates(imdb_id, sn)
            for en, air in m.items():
                key = (sn,en)
                if key not in out or not out[key][1]:
                    air_iso = to_utc_iso(air)
                    out[key] = (air_iso, air_iso, "omdb:Released")

    # File mtime fallback
    for p in series_dir.iterdir():
        if not p.is_dir(): continue
        name = p.name.lower()
        if name=="specials": sn = 0
        elif name.startswith("season ") and p.name.split()[-1].isdigit():
            sn = int(p.name.split()[-1])
        else:
            continue
        media = [f for f in p.iterdir() if f.is_file() and re.search(r"S(\d{2})E(\d{2})", f.name, re.IGNORECASE)]
        for f in media:
            m = re.search(r"S(\d{2})E(\d{2})", f.name, re.IGNORECASE)
            if not m: continue
            sn2, en2 = int(m.group(1)), int(m.group(2))
            if sn2 != sn: continue
            key = (sn,en2)
            if key not in out or not out[key][1]:
                mt = file_mtime_iso(f)
                out[key] = ((out.get(key) or (None,))[0], mt, "file:mtime")
    return out

def apply_series(series_dir: Path, args, conn: sqlite3.Connection):
    imdb_id = parse_imdb_from_path(series_dir)
    if not imdb_id:
        _log("ERROR", f"Series path does not contain an IMDb id [imdb-tt...] : {series_dir}")
        return

    conn.execute(
        "INSERT INTO series(imdb_id, series_path, last_applied) VALUES(?,?,?) "
        "ON CONFLICT(imdb_id) DO UPDATE SET series_path=excluded.series_path, last_applied=excluded.last_applied",
        (imdb_id, str(series_dir), _iso_now()),
    )
    conn.commit()

    if args.manage_nfo:
        for sd in series_dir.iterdir():
            if not sd.is_dir(): continue
            name = sd.name.lower()
            if not (name=="specials" or name.startswith("season ")): continue
            removed = remove_long_nfos(sd)
            if removed:
                _log("INFO", f"Removed {removed} long-name NFO(s) in {sd}")
        maybe_tvshow_nfo(series_dir, imdb_id)

    dates = gather_episode_dates(series_dir, imdb_id, args)

    seasons_latest = {}
    for (sn,en), (aired, dateadded, source) in sorted(dates.items()):
        if args.manage_nfo:
            write_episode_nfo(series_dir, sn, en, aired, dateadded, source or "")
            # add season.nfo (safe)
            ensure_season_nfo(series_dir, sn)

        conn.execute(
            "INSERT INTO episode_dates(imdb_id, season, episode, aired, dateadded, source) "
            "VALUES(?,?,?,?,?,?) "
            "ON CONFLICT(imdb_id, season, episode) DO UPDATE SET aired=excluded.aired, dateadded=excluded.dateadded, source=excluded.source",
            (imdb_id, sn, en, aired or None, dateadded or None, source or None),
        )

        if dateadded:
            try:
                ts = datetime.fromisoformat(dateadded.replace("Z","+00:00"))
                seasons_latest[sn] = max(ts, seasons_latest.get(sn, ts))
            except Exception:
                pass

    conn.commit()

    if args.fix_dir_mtimes and seasons_latest:
        for sn, ts in seasons_latest.items():
            season_dir = resolve_season_dir(series_dir, sn)
            iso = ts.astimezone(timezone.utc).isoformat(timespec="seconds")
            set_dir_mtime(season_dir, iso)
            _log("INFO", f"Updated mtime on {season_dir} -> {iso}")
            if args.update_video_mtimes:
                set_media_mtimes(season_dir, iso)

def iterate_target_path(path: Path):
    if is_series_dir(path): return [path]
    if path.is_dir():
        out = [p for p in path.iterdir() if is_series_dir(p)]
        return sorted(out)
    return []

# ---------- CLI ----------
def cmd_rebuild_db(args):
    _log("INFO", f"Rebuilding DB at {args.db}")
    conn = sqlite3.connect(str(Path(args.db)))
    conn.execute("PRAGMA foreign_keys=ON;")
    conn.executescript("""
    DROP TABLE IF EXISTS episode_dates;
    DROP TABLE IF EXISTS series;
    """)
    _ensure_schema(conn, allow_drop=False)
    conn.commit(); conn.close()

def cmd_apply(args):
    conn = db_connect(Path(args.db))
    targets = iterate_target_path(Path(args.path))
    if not targets:
        _log("ERROR", f"No series found at {args.path} (expects dirs named like 'Show Name [imdb-ttxxxxxxx]')")
        return
    if not (os.environ.get("SONARR_URL") and (os.environ.get("SONARR_API_KEY") or os.environ.get("SONARR_APIKEY"))):
        _log("INFO", "Sonarr not configured (set SONARR_URL and SONARR_API_KEY)")
    if not os.environ.get("TMDB_API_KEY"):
        _log("INFO", "TMDB not configured (set TMDB_API_KEY)")
    if not os.environ.get("OMDB_API_KEY"):
        _log("INFO", "OMDb not configured (set OMDB_API_KEY)")

    _log("INFO", f"=== APPLY: {targets[0] if len(targets)==1 else args.path} ===")
    for series_dir in targets:
        try:
            apply_series(series_dir, args, conn)
        except KeyboardInterrupt:
            raise
        except Exception as e:
            _log("ERROR", f"Failed applying {series_dir}: {e}")
    conn.close()

def build_parser():
    p = argparse.ArgumentParser(prog="media_date_cache.py")
    p.add_argument("--env", help="Path to .env file (can appear anywhere on the command)", default=None)
    sub = p.add_subparsers(dest="cmd", required=True)

    p_rebuild = sub.add_parser("rebuild-db", help="Drop & recreate schema")
    p_rebuild.add_argument("--db", required=True, help="Path to sqlite DB")
    p_rebuild.set_defaults(func=cmd_rebuild_db)

    p_apply = sub.add_parser("apply", help="Apply dates → NFO + mtimes")
    p_apply.add_argument("--db", required=True, help="Path to sqlite DB")
    p_apply.add_argument("--path", required=True, help="Series dir (with [imdb-tt...]) or a root containing them")
    p_apply.add_argument("--manage-nfo", action="store_true", help="Create short NFOs; remove long-name NFOs; ensure season.nfo/tvshow.nfo")
    p_apply.add_argument("--fix-dir-mtimes", action="store_true", help="Set Season/ Specials folder mtimes to latest episode 'dateadded'")
    p_apply.add_argument("--update-video-mtimes", action="store_true", help="Also set video/subtitle file mtimes (use with care)")
    p_apply.set_defaults(func=cmd_apply)

    return p

def main():
    parser = build_parser()
    args = parser.parse_args()
    args.func(args)

if __name__ == "__main__":
    main()