import os
import requests
import json
from pathlib import Path
import logging
import pickle
from requests.adapters import HTTPAdapter
from urllib3.util.retry import Retry
import time

# Base directories
source_dir = "/mnt/data/Anime"
target_dir = "/mnt/data/anime-jellyfin"

# Configure logging
log_dir = target_dir
logging.basicConfig(
    filename=os.path.join(log_dir, "symlink_creation.log"),
    level=logging.INFO,
    format="%(asctime)s - %(levelname)s - %(message)s"
)

# Cache file to track processed directories and files
CACHE_FILE = os.path.join(target_dir, "processed_dirs.pkl")

# Ollama API endpoint
OLLAMA_URL = "http://localhost:11434/api/generate"

# Supported media file extensions
MEDIA_EXTENSIONS = {'.mkv', '.mp4', '.avi', '.ts', '.m4v', '.mpeg', '.mpg', '.srt', '.ass'}

# Function to create a session with retry logic
def create_requests_session():
    session = requests.Session()
    retries = Retry(total=3, backoff_factor=1, status_forcelist=[429, 500, 502, 503, 504])
    session.mount("http://", HTTPAdapter(max_retries=retries))
    return session

# Function to check if Ollama is reachable
def check_ollama_availability():
    session = create_requests_session()
    try:
        response = session.post(OLLAMA_URL, json={
            "model": "gemma3-27b-qat",
            "prompt": "Test connectivity",
            "stream": False
        }, timeout=30)
        response.raise_for_status()
        logging.info("Successfully connected to Ollama.")
        print("Successfully connected to Ollama.")
        return True
    except Exception as e:
        logging.error(f"Cannot connect to Ollama at {OLLAMA_URL}: {e}")
        print(f"Error: Cannot connect to Ollama at {OLLAMA_URL}: {e}")
        return False

# Function to load processed directories from cache
def load_processed_dirs():
    if os.path.exists(CACHE_FILE):
        with open(CACHE_FILE, "rb") as f:
            return pickle.load(f)
    return {}

# Function to save processed directories to cache
def save_processed_dirs(processed_dirs):
    with open(CACHE_FILE, "wb") as f:
        pickle.dump(processed_dirs, f)

# Function to clean up broken symlinks and empty directories
def cleanup_broken_symlinks_and_empty_dirs(target_dir):
    broken_symlinks = 0
    empty_dirs_removed = 0
    logging.info("Starting cleanup of broken symlinks and empty directories...")
    print("Starting cleanup of broken symlinks and empty directories...")

    for root, dirs, files in os.walk(target_dir, topdown=False):
        if root.startswith(log_dir):
            continue
        for name in files:
            path = os.path.join(root, name)
            if os.path.islink(path) and not os.path.exists(path):
                try:
                    os.remove(path)
                    logging.info(f"Removed broken symlink: {path}")
                    print(f"Removed broken symlink: {path}")
                    broken_symlinks += 1
                except OSError as e:
                    logging.error(f"Error removing broken symlink {path}: {e}")
                    print(f"Error removing broken symlink {path}: {e}")
        for name in dirs:
            path = os.path.join(root, name)
            if not os.path.islink(path) and not os.listdir(path):
                try:
                    os.rmdir(path)
                    logging.info(f"Removed empty directory: {path}")
                    print(f"Removed empty directory: {path}")
                    empty_dirs_removed += 1
                except OSError as e:
                    logging.error(f"Error removing empty directory {path}: {e}")
                    print(f"Error removing empty directory {path}: {e}")

    logging.info(f"Cleanup completed. Removed {broken_symlinks} broken symlinks, {empty_dirs_removed} empty directories.")
    print(f"Cleanup completed. Removed {broken_symlinks} broken symlinks, {empty_dirs_removed} empty directories.")
    return broken_symlinks, empty_dirs_removed

# Function to query Ollama for parsing directory names
def parse_directory_name(dir_name):
    session = create_requests_session()
    prompt = f"""
    You are an expert at parsing anime directory names. Given the directory name "{dir_name}", extract:
    1. The show title (the main name of the show).
    2. The season number (e.g., 1, 2, 3, etc.). If no season is specified, assume season 1.
    3. Any extra information like "(Dub)" or story arc names (e.g., "Shoal of Time").

    Return the result as a JSON object with keys: "show_name", "season", "extras".
    Examples:
    - "100-man no Inochi no Ue ni Ore wa Tatteiru 2nd Season" -> {{"show_name": "100-man no Inochi no Ue ni Ore wa Tatteiru", "season": 2, "extras": ""}}
    - "91 Days - Shoal of TimeAll Our YesterdaysTomorrow and Tomorrow" -> {{"show_name": "91 Days", "season": 1, "extras": "Shoal of TimeAll Our YesterdaysTomorrow and Tomorrow"}}
    - "Accel World" -> {{"show_name": "Accel World", "season": 1, "extras": ""}}
    - "3-gatsu no Lion (Dub)" -> {{"show_name": "3-gatsu no Lion", "season": 1, "extras": "(Dub)"}}
    - ".hackSign" -> {{"show_name": ".hackSign", "season": 1, "extras": ""}}
    - "Shuffle!" -> {{"show_name": "Shuffle", "season": 1, "extras": ""}}
    """

    max_retries = 3
    for attempt in range(max_retries):
        try:
            response = session.post(OLLAMA_URL, json={
                "model": "gemma3-27b-qat",
                "prompt": prompt,
                "format": "json",
                "stream": False
            }, timeout=30)
            response.raise_for_status()
            result = json.loads(response.json().get("response", "{}"))
            show_name = result.get("show_name", dir_name)
            season = result.get("season", 1)
            extras = result.get("extras", "")

            try:
                season = int(season)
            except ValueError:
                logging.warning(f"Non-numeric season '{season}' for {dir_name}. Falling back to show_name='{dir_name}', season=1, extras=''")
                print(f"Warning: Non-numeric season '{season}' for {dir_name}. Falling back to show_name='{dir_name}', season=1, extras=''")
                show_name, season, extras = dir_name, 1, ""

            return show_name, season, extras
        except requests.exceptions.HTTPError as e:
            logging.error(f"Attempt {attempt + 1}/{max_retries} failed parsing {dir_name}: {e}. Response: {response.text}")
            print(f"Attempt {attempt + 1}/{max_retries} failed parsing {dir_name}: {e}. Response: {response.text}")
            if attempt < max_retries - 1:
                time.sleep(2)
                continue
            logging.error(f"Failed to parse {dir_name} after {max_retries} attempts. Exiting.")
            print(f"Error: Failed to parse {dir_name} after {max_retries} attempts. Exiting.")
            exit(1)
        except Exception as e:
            logging.error(f"Attempt {attempt + 1}/{max_retries} failed parsing {dir_name}: {e}")
            print(f"Attempt {attempt + 1}/{max_retries} failed parsing {dir_name}: {e}")
            if attempt < max_retries - 1:
                time.sleep(2)
                continue
            logging.error(f"Failed to parse {dir_name} after {max_retries} attempts. Exiting.")
            print(f"Error: Failed to parse {dir_name} after {max_retries} attempts. Exiting.")
            exit(1)

# Clean up broken symlinks and empty directories
broken_symlinks_removed, empty_dirs_removed = cleanup_broken_symlinks_and_empty_dirs(target_dir)

# Check Ollama availability
if not check_ollama_availability():
    logging.error("Aborting hardlink creation due to Ollama being unreachable.")
    print("Error: Aborting hardlink creation due to Ollama being unreachable.")
    exit(1)

# Load processed directories
processed_dirs = load_processed_dirs()

# Get total directories for progress
dir_list = [d for d in os.listdir(source_dir) if os.path.isdir(os.path.join(source_dir, d))]
total_dirs = len(dir_list)
dirs_processed = 0

# Walk through source directory
new_dirs_processed = 0
new_files_linked = 0
for dir_name in dir_list:
    source_path = os.path.join(source_dir, dir_name)
    dirs_processed += 1
    progress = (dirs_processed / total_dirs) * 100
    logging.info(f"Processing {dir_name} ({dirs_processed}/{total_dirs}, {progress:.1f}%)")
    print(f"Processing {dir_name} ({dirs_processed}/{total_dirs}, {progress:.1f}%)")

    if os.path.isdir(source_path):
        # Check if directory is in cache
        if dir_name in processed_dirs:
            show_name = processed_dirs[dir_name]["show_name"]
            season = processed_dirs[dir_name]["season"]
            extras = processed_dirs[dir_name]["extras"]
            cached_files = processed_dirs[dir_name]["files"]
        else:
            # Parse directory name using Ollama
            show_name, season, extras = parse_directory_name(dir_name)
            cached_files = set()
            processed_dirs[dir_name] = {
                "show_name": show_name,
                "season": season,
                "extras": extras,
                "files": cached_files
            }
            new_dirs_processed += 1

        # Sanitize show name for filesystem compatibility
        show_name = show_name.replace("/", "-").replace(":", "-")

        # Handle dubbed versions
        jellyfin_show_name = show_name
        if "(Dub)" in extras:
            jellyfin_show_name += " (Dub)"

        # Create Jellyfin-compatible directory structure
        jellyfin_show_path = os.path.join(target_dir, jellyfin_show_name, f"Season{season}")
        os.makedirs(jellyfin_show_path, exist_ok=True)

        # Create hard links for new media files
        for file_name in os.listdir(source_path):
            if os.path.splitext(file_name)[1].lower() in MEDIA_EXTENSIONS and file_name not in cached_files:
                source_file = os.path.join(source_path, file_name)
                if os.path.isfile(source_file):
                    hardlink_path = os.path.join(jellyfin_show_path, file_name)
                    try:
                        # Check if hardlink already exists
                        if os.path.exists(hardlink_path):
                            logging.info(f"Hardlink exists: {hardlink_path}")
                            print(f"Hardlink exists: {hardlink_path}")
                            processed_dirs[dir_name]["files"].add(file_name)
                            continue
                        # Create hard link
                        os.link(source_file, hardlink_path)
                        logging.info(f"Created hardlink: {hardlink_path} -> {source_file}")
                        print(f"Created hardlink: {hardlink_path} -> {source_file}")
                        new_files_linked += 1
                        processed_dirs[dir_name]["files"].add(file_name)
                    except OSError as e:
                        logging.error(f"Error creating hardlink for {file_name} in {dir_name}: {e}")
                        print(f"Error creating hardlink for {file_name} in {dir_name}: {e}")
                    except FileNotFoundError as e:
                        logging.error(f"Target file not found for {file_name} in {dir_name}: {e}")
                        print(f"Target file not found for {file_name} in {dir_name}: {e}")

# Save processed directories
save_processed_dirs(processed_dirs)

logging.info(f"Hardlink creation completed! Processed {new_dirs_processed} new directories, linked {new_files_linked} new files. Removed {broken_symlinks_removed} broken symlinks, {empty_dirs_removed} empty directories.")
print(f"Hardlink creation completed! Processed {new_dirs_processed} new directories, linked {new_files_linked} new files. Removed {broken_symlinks_removed} broken symlinks, {empty_dirs_removed} empty directories.")