import subprocess
import json
import shutil
from pathlib import Path
from typing import List

def run(cmd: subprocess._CMD) -> subprocess.CompletedProcess[str]:
    """
    Run a subprocess command.

    Args:
        cmd: Command and arguments as a list.

    Raises:
        subprocess.CalledProcessError: If the command exits with a non-zero status.
    """
    try:
        result = subprocess.run(
            cmd,
            check=True,
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT,
            text=True
        )
        return result
    except subprocess.CalledProcessError as e:
        print(e.stdout)
        raise

def get_audio_streams(file: Path) -> List[int]:
    """
    Return a list of available audio stream indices (0-based for ffmpeg mapping).
    
    Example return:
        [0, 1]  -> corresponds to 0:a:0 and 0:a:1
    """
    cmd = [
        "ffprobe",
        "-i", str(file),
        "-select_streams", "a",
        "-show_entries", "stream=index",
        "-of", "json",
        "-loglevel", "error"
    ]

    result = run(cmd)
    data = json.loads(result.stdout)

    # ffprobe gives absolute stream indices, but we want relative audio indices
    streams = data.get("streams", [])

    # Convert to 0-based audio stream numbering
    return list(range(len(streams)))

def get_track_info(file: Path, audioline: int):
    """Gets codec, bitrate, language, and title."""
    cmd = [
        "ffprobe", "-v", "error", "-select_streams", f"a:{audioline}",
        "-show_entries", "stream=codec_name,bit_rate:stream_tags=language,title",
        "-of", "json", str(file)
    ]
    result = subprocess.run(cmd, capture_output=True, text=True)
    data = json.loads(result.stdout)
    if "streams" in data and len(data["streams"]) > 0:
        s = data["streams"][0]
        return {
            "codec": s.get("codec_name", "aac"),
            "bitrate": s.get("bit_rate", "192k"),
            "lang": s.get("tags", {}).get("language", "eng"),
            "title": s.get("tags", {}).get("title", f"Track {audioline}")
        }
    return {"codec": "aac", "bitrate": "192k", "lang": "eng", "title": "Audio"}

def shift_audio(file: Path, audioline: int, offset: float):
    """
    Shifts audio by offset sekonds to the right.
    """
    if not file.is_file(): return
    
    info = get_track_info(file, audioline)
    # mkvmerge uses 0 for video, usually 1+ for audio. 
    # We need the MKV Toolnix Track ID.
    mkv_id = audioline + 1 
    
    temp_mka = file.parent / "temp_shifted.mka"
    output_mkv = file.parent / "temp_output.mkv"

    # --- STEP 1: Physical Shift via Filter ---
    # We convert offset to milliseconds for the 'adelay' filter
    ms = int(abs(offset) * 1000)
    
    if offset > 0:
        # RIGHT SHIFT: Delay start, then trim the end to keep original length
        # 'all' applies delay to all channels
        filter_str = f"adelay={ms}|{ms}|{ms}|{ms}|{ms}|{ms},atrim=0,atrim=duration=base" # Simplified logic below
        # For reliability, we use a simpler pad/trim approach:
        filter_chain = f"adelay={ms}:all=1"
    else:
        # LEFT SHIFT: Skip the beginning
        filter_chain = f"atrim=start={abs(offset)}"

    # Process to MKA
    # We use the same codec and bitrate detected earlier
    process_cmd = [
        "ffmpeg", "-y", "-i", str(file),
        "-map", f"0:a:{audioline}",
        "-af", filter_chain,
        "-c:a", info['codec'],
        "-b:a", str(info['bitrate']),
        str(temp_mka)
    ]
    subprocess.run(process_cmd, check=True)

    # --- STEP 2: Replace in MKV via mkvmerge ---
    # --audio-tracks !ID excludes the old track
    # The new MKA is added at the end
    merge_cmd = [
        "mkvmerge", "-o", str(output_mkv),
        "--audio-tracks", f"!{mkv_id}", str(file),
        "--language", f"0:{info['lang']}",
        "--track-name", f"0:{info['title']}",
        str(temp_mka)
    ]
    
    subprocess.run(merge_cmd, check=True)

    # Cleanup and Replace
    shutil.move(str(output_mkv), str(file))
    temp_mka.unlink()
    print(f"Done: {file.name} updated physically.")

# Example usage:
file_path = Path("/home/honney/mount/Media/tmp/serien/ready/Charlotte/Season 01/[NO CHAPTER]_E12.mkv")
audio_streams = get_audio_streams(file_path)
shift_audio(file_path, audioline=audio_streams[-1], offset=(1/23.98)) # Shifts the 2nd audio track right by 1.5 seconds