"""Core functions for PodangelEX audio processing and configuration."""

import os
import json
import re
import signal
import subprocess
import tempfile
import time
from collections import namedtuple
from pathlib import Path
from importlib import resources

import whisper_timestamped
from detoxify import Detoxify
import multiprocessing

worker_model = None
global_pool = None
shutdown_in_progress = False
detox_model = None

WordToMute = namedtuple('WordToMute', ['word', 'start', 'end'])
SegmentToMute = namedtuple('SegmentToMute', ['start', 'end'])


def get_config_dir() -> Path:
    """Get configuration directory - respects PODANGELEX_HOME env var."""
    if env_home := os.getenv('PODANGELEX_HOME'):
        return Path(env_home).expanduser()
    return Path.home() / '.podangelex'


def get_default_workspace() -> Path:
    """Get default workspace directory with smart detection."""
    # Check for environment variable override
    if env_workspace := os.getenv('PODANGELEX_WORKSPACE'):
        return Path(env_workspace).expanduser()
    
    # Development: use project-local workspace if it exists
    script_dir = Path(__file__).parent.parent.parent.parent  # Up to package/ root
    dev_workspace = script_dir / 'podangelex_workspace'
    if dev_workspace.exists():
        return dev_workspace
    
    # Production: use ~/podangelex_data
    return Path.home() / 'podangelex_data'


def initialize_workspace(config_dir: Path = None) -> Path:
    """Initialize workspace directories and create default config if missing."""
    if config_dir is None:
        config_dir = get_config_dir()
    
    config_dir = Path(config_dir)
    config_dir.mkdir(parents=True, exist_ok=True)
    
    config_path = config_dir / 'config.json'
    
    # Determine workspace location
    workspace_path = get_default_workspace()
    
    # Create workspace directories
    for subdir in ['Input', 'Output', '.bridge']:
        (workspace_path / subdir).mkdir(parents=True, exist_ok=True)
    
    # Create default config if missing
    if not config_path.exists():
        default_config = {
            "file_path": str(workspace_path),
            "worker_count": "1",
            "model_size": "small",
            "t": 0.5, "st": 0.5, "o": 0.5, "th": 0.5, "i": 0.5, "id": 0.5
        }
        write_config(str(config_path), default_config)
    
    return workspace_path


def get_swears_path() -> Path:
    """Get path to swears.json bundled with package."""
    try:
        # Try to get from package resources (when installed)
        if hasattr(resources, 'files'):  # Python 3.9+
            swears_file = resources.files('podangelex_JustAnotherCoderTheThird') / 'swears.json'
            return str(swears_file)
        else:
            # Fallback for older Python
            import pkg_resources
            return pkg_resources.resource_filename('podangelex_JustAnotherCoderTheThird', 'swears.json')
    except Exception:
        # Fallback: look in same directory as this file
        fallback_path = Path(__file__).parent / 'swears.json'
        if fallback_path.exists():
            return str(fallback_path)
        return None


def read_config(file_path):
    """Read configuration from JSON file."""
    try:
        with open(file_path, 'r') as f:
            return json.load(f)
    except FileNotFoundError:
        return None


def read_json(file_path):
    """Read JSON file."""
    try:
        with open(file_path, 'r') as f:
            return json.load(f)
    except FileNotFoundError:
        return None


def write_config(file_path, content):
    """Write configuration to JSON file."""
    try:
        Path(file_path).parent.mkdir(parents=True, exist_ok=True)
        with open(file_path, 'w') as outfile:
            json.dump(content, outfile, indent=4, sort_keys=True)
    except Exception as e:
        print(f"Error writing config: {e}")


def size_finder(size_input: str, config_dict: dict, config_dir: Path) -> int | str:
    """Map user input to model name and VRAM cost."""
    models = {
        "1": ("tiny", 1),
        "2": ("base", 1),
        "3": ("small", 2),
        "4": ("medium", 5),
        "5": ("large", 10),
        "6": ("turbo", 6),
    }
    
    key = size_input.lower().strip()
    if key not in models:
        print("Please select a valid number")
        return "F"
    
    model_name, vram_cost = models[key]
    config_dict["model_size"] = model_name
    write_config(str(config_dir / "config.json"), config_dict)
    return vram_cost


def category_finder(input_str: str, config_dict: dict, value: float, config_dir: Path) -> None:
    """Update a toxicity category threshold."""
    valid_categories = {"t", "st", "th", "o", "id", "i"}
    key = input_str.lower().strip()
    
    if key not in valid_categories:
        print("Please select a valid category")
        return
    
    config_dict[key] = value
    write_config(str(config_dir / "config.json"), config_dict)


def severity_tweaker(input_str: str, config_dict: dict, config_dir: Path) -> None:
    """Parse and apply severity threshold changes."""
    try:
        category, value = input_str.split("-")
        float_value = float(value)
        
        if not 0 <= float_value <= 1:
            print("Number must be between 0 and 1")
            return
        
        category_finder(category, config_dict, float_value, config_dir)
    except (ValueError, IndexError):
        print("Invalid format. Use: category-value (e.g., st-0.5)")


def config_menu(config_dir: Path = None) -> None:
    """Interactive configuration menu."""
    if config_dir is None:
        config_dir = get_config_dir()
    
    config_dir = Path(config_dir)
    config_path = config_dir / "config.json"
    
    config_data = {
        "file_path": str(get_default_workspace()),
        "worker_count": "1",
        "model_size": "small",
        "t": 0.5, "st": 0.5, "o": 0.5, "th": 0.5, "i": 0.5, "id": 0.5
    }
    
    try:
        with open(str(config_path), 'r') as f:
            config_data = json.load(f)
    except FileNotFoundError:
        pass
    
    print("\n\nSelect model size:\n 1: tiny (1GB)\n 2: base (1GB)\n 3: small (2GB)\n"
          " 4: medium (5GB)\n 5: large (10GB)\n 6: turbo (6GB)")
    while True:
        model_size = input("\nInput a number: ")
        vram_cost = size_finder(model_size, config_data, config_dir)
        if vram_cost != "F":
            break
    
    if input("\nChange file paths from default? (y/N): ").lower().strip() == "y":
        new_dir = input("Enter path: ").strip()
    else:
        new_dir = str(get_default_workspace())
    
    if not os.path.isabs(new_dir):
        new_dir = str(Path.home() / new_dir)
    
    config_data["file_path"] = new_dir
    write_config(str(config_path), config_data)
    
    while True:
        try:
            worker_count = input("\nNumber of workers: ").strip()
            int_count = int(worker_count)
            vram_total = int_count * vram_cost
            if input(f"\nWarning: Uses {vram_total}GB VRAM. Continue? (Y/n): ").lower().strip() != "n":
                config_data["worker_count"] = worker_count
                write_config(str(config_path), config_data)
                break
        except ValueError:
            print("Please enter a valid number")
    
    if input("\nConfigure toxicity thresholds? (y/N): ").lower().strip() == "y":
        print("\nSet thresholds (category-value, e.g., st-0.5):\n"
              "t=toxicity, st=severe, o=obscene, th=threats, i=insults, id=identity")
        while True:
            severity_tweaker(input("Enter setting (or press Enter to skip): "), config_data, config_dir)
            if input("Continue? (y/N): ").lower().strip() != "y":
                break
    
    # Create workspace directories
    for subdir in ['Input', 'Output', '.bridge']:
        os.makedirs(os.path.join(new_dir, subdir), exist_ok=True)
    
    write_config(str(config_path), config_data)
    print("\nConfiguration complete!")


def get_audio_duration(file_path: str) -> float | None:
    """Get audio duration using ffprobe."""
    cmd = [
        'ffprobe', '-v', 'error', '-show_entries', 'format=duration',
        '-of', 'default=noprint_wrappers=1:nokey=1', file_path
    ]
    try:
        result = subprocess.run(cmd, capture_output=True, text=True, check=True)
        return float(result.stdout.strip())
    except subprocess.CalledProcessError as e:
        print(f"Error: ffprobe failed for {file_path} - {e.stderr}")
        print(f"Make sure ffprobe is installed and the file is a valid audio format")
        return None
    except ValueError:
        print(f"Error: Could not parse duration from ffprobe output for {file_path}")
        return None


def worker_initializer(model_size: str) -> None:
    """Initialize Whisper model in worker process."""
    global worker_model
    worker_model = whisper_timestamped.load_model(model_size)


def process_file(filename: str, input_path: str) -> dict:
    """Transcribe audio file using Whisper."""
    global worker_model
    file_path = os.path.join(input_path, filename)
    transcribed_file = whisper_timestamped.transcribe(worker_model, file_path)
    print(f"Transcribed: {filename}")
    return transcribed_file


def cleanup_workers() -> None:
    """Clean up worker pool."""
    global global_pool
    if global_pool:
        print("Shutting down worker pool...")
        global_pool.terminate()
        global_pool.join()
        global_pool = None


def signal_handler(signum, frame):
    """Handle shutdown signals."""
    global shutdown_in_progress
    if shutdown_in_progress:
        return
    shutdown_in_progress = True
    print(f"\nReceived signal {signum}, shutting down...")
    cleanup_workers()
    exit()


signal.signal(signal.SIGINT, signal_handler)


def clean_text(text: str, words_to_remove: list) -> str:
    """Remove words from text while preserving context."""
    cleaned = text
    for word in words_to_remove:
        cleaned = re.sub(rf'\b{re.escape(word)}\b', '', cleaned, flags=re.IGNORECASE)
    return cleaned.strip()


def audio_cleaner(transcribed_file: dict, bad_words: dict, thresholds: list) -> tuple:
    """Find segments and words to mute based on toxicity."""
    words_to_mute = []
    segments_to_mute = []
    swear_words = {s.lower() for s in bad_words.get('swears', [])}
    
    for segment in transcribed_file.get('segments', []):
        for word in segment.get('words', []):
            word_text = word.get('text', '').lower().strip('.,!?')
            if word_text in swear_words:
                words_to_mute.append(WordToMute(
                    word['text'],
                    float(word['start']),
                    float(word['end'])
                ))

    global detox_model
    if detox_model is None:
        print("Warning: detox_model not initialized")
        return words_to_mute, segments_to_mute
    
    swear_set = {m.word.lower().strip('.,!?') for m in words_to_mute}
    
    for segment in transcribed_file.get('segments', []):
        cleaned_text = clean_text(segment.get('text', ''), list(swear_set))
        
        if not cleaned_text.strip():
            continue
        
        ratings = detox_model.predict(cleaned_text)
        toxic = any([
            ratings.get('toxicity', 0) > thresholds[0],
            ratings.get('severe_toxicity', 0) > thresholds[1],
            ratings.get('obscene', 0) > thresholds[2],
            ratings.get('threat', 0) > thresholds[5],
            ratings.get('insult', 0) > thresholds[4],
            ratings.get('identity_attack', 0) > thresholds[3],
        ])
        
        if toxic:
            segments_to_mute.append(SegmentToMute(
                float(segment['start']),
                float(segment['end'])
            ))
    return words_to_mute, segments_to_mute


def mute_words(input_file: str, output_file: str, words_to_mute: list, segments_to_mute: list) -> bool:
    """Extract clean audio segments and concatenate them."""
    if not words_to_mute and not segments_to_mute:
        try:
            subprocess.run(['cp', input_file, output_file], check=True)
            return True
        except subprocess.CalledProcessError as e:
            print(f"Error: Failed to copy clean audio file: {e}")
            return False

    all_segments = [(float(w.start), float(w.end)) for w in words_to_mute]
    all_segments.extend([(float(s.start), float(s.end)) for s in segments_to_mute])
    
    if not all_segments:
        try:
            subprocess.run(['cp', input_file, output_file], check=True)
            return True
        except subprocess.CalledProcessError as e:
            print(f"Error: Failed to copy audio file: {e}")
            return False

    all_segments.sort()
    merged = []
    for start, end in all_segments:
        if merged and start <= merged[-1][1]:
            merged[-1] = (merged[-1][0], max(merged[-1][1], end))
        else:
            merged.append((start, end))

    duration = get_audio_duration(input_file)
    if duration is None:
        print(f"Error: Could not process {input_file} - unable to get audio duration")
        return False

    segments_to_keep = []
    current_pos = 0.0
    
    for start, end in merged:
        if current_pos < start:
            segments_to_keep.append((current_pos, start))
        current_pos = max(current_pos, end)
    
    if current_pos < duration:
        segments_to_keep.append((current_pos, duration))
    
    epsilon = 0.001
    segments_to_keep = [(max(0, start), min(duration, end)) for start, end in segments_to_keep]
    segments_to_keep = [(start, end) for start, end in segments_to_keep if end - start > epsilon]
    
    if not segments_to_keep:
        print("All audio removed, creating silent file")
        try:
            result = subprocess.run(['ffmpeg', '-y', '-f', 'lavfi', '-i', 'anullsrc=r=44100:cl=mono', '-t', '0.1', output_file], capture_output=True, text=True, check=True)
            if os.path.exists(output_file) and os.path.getsize(output_file) > 0:
                return True
            else:
                print(f"Error: Failed to create silent audio file for {output_file}")
                return False
        except subprocess.CalledProcessError as e:
            print(f"Error: ffmpeg failed to create silent file: {e.stderr}")
            return False

    with tempfile.TemporaryDirectory() as temp_dir:
        file_ext = Path(input_file).suffix
        segment_files = []
        
        try:
            for i, (start, end) in enumerate(segments_to_keep):
                segment_file = os.path.join(temp_dir, f"segment_{i}{file_ext}")
                segment_files.append(segment_file)
                duration_segment = end - start
                cmd = ['ffmpeg', '-y', '-i', input_file, '-ss', str(start), '-t', str(duration_segment), '-vn', '-c:a', 'copy', segment_file]
                result = subprocess.run(cmd, capture_output=True, text=True)
                if result.returncode != 0:
                    print(f"Error: Failed to extract segment {i} from {input_file}:\n{result.stderr}")
                    return False
                if not os.path.exists(segment_file):
                    print(f"Error: Segment file {segment_file} was not created by ffmpeg")
                    return False

            if len(segment_files) == 1:
                try:
                    subprocess.run(['cp', segment_files[0], output_file], check=True)
                except subprocess.CalledProcessError as e:
                    print(f"Error: Failed to copy single segment to output: {e}")
                    return False
            else:
                concat_file = os.path.join(temp_dir, "concat.txt")
                with open(concat_file, 'w') as f:
                    f.write('\n'.join(f"file '{seg}'" for seg in segment_files))
                
                cmd = ['ffmpeg', '-y', '-f', 'concat', '-safe', '0', '-i', concat_file, '-c:a', 'copy', '-map_metadata', '0', output_file]
                result = subprocess.run(cmd, capture_output=True, text=True)
                if result.returncode != 0:
                    print(f"Error: Concatenation failed for {output_file}:\n{result.stderr}")
                    return False
            
            if os.path.exists(output_file) and os.path.getsize(output_file) > 0:
                print(f"Output written to {output_file}")
                return True
            else:
                print(f"Error: Output file {output_file} was not created or is empty")
                return False
        except Exception as e:
            print(f"Error processing audio: {e}")
            return False


def run_program(config_dict: dict, config_dir: Path = None) -> list | None:
    """Main audio processing orchestration."""
    if config_dir is None:
        config_dir = get_config_dir()
    
    model_size = config_dict["model_size"]
    path = config_dict["file_path"]
    
    if not os.path.isabs(path):
        path = str(Path.home() / path)
    
    worker_counts = int(config_dict["worker_count"])
    
    thresholds = [
        config_dict["t"],
        config_dict["st"],
        config_dict["o"],
        config_dict["id"],
        config_dict["i"],
        config_dict["th"]
    ]
    
    # Load swears from package resources
    swears_path = get_swears_path()
    if swears_path and os.path.exists(swears_path):
        bad_words = read_json(swears_path) or {"swears": []}
    else:
        bad_words = {"swears": []}
    
    global detox_model
    detox_model = Detoxify('original')

    input_path = os.path.join(path, 'Input')
    output_path = os.path.join(path, 'Output')

    in_files = {f for f in os.listdir(input_path) if os.path.isfile(os.path.join(input_path, f)) and not f.startswith('.')}
    out_files = {f for f in os.listdir(output_path) if os.path.isfile(os.path.join(output_path, f)) and not f.startswith('.')}
    
    working_list = sorted(in_files - out_files)
    
    if not working_list:
        print("No files to update")
        return None

    global global_pool
    results = []
    
    try:
        with multiprocessing.Pool(processes=worker_counts, initializer=worker_initializer, initargs=(model_size,)) as pool:
            global_pool = pool

            job_to_filename = {}
            for filename in working_list:
                async_result = pool.apply_async(process_file, (filename, input_path))
                job_to_filename[async_result] = filename

            pending = list(job_to_filename.keys())
            processed_files = {"success": [], "failed": []}
            while pending:
                for async_result in list(pending):
                    if async_result.ready():
                        filename = job_to_filename[async_result]
                        try:
                            transcribed_file = async_result.get()
                            results.append(transcribed_file)
                            words_to_mute, segments_to_mute = audio_cleaner(transcribed_file, bad_words, thresholds)
                            success = mute_words(
                                os.path.join(input_path, filename),
                                os.path.join(output_path, filename),
                                words_to_mute,
                                segments_to_mute
                            )
                            if success:
                                processed_files["success"].append(filename)
                            else:
                                processed_files["failed"].append(filename)
                        except Exception as e:
                            print(f"Error processing {filename}: {e}")
                            processed_files["failed"].append(filename)
                        pending.remove(async_result)
                if pending:
                    time.sleep(0.1)
            
            global_pool = None
            
            print(f"\n=== Processing Summary ===")
            print(f"Successfully processed: {len(processed_files['success'])} files")
            if processed_files['success']:
                for f in processed_files['success']:
                    print(f"  ✓ {f}")
            if processed_files['failed']:
                print(f"Failed to process: {len(processed_files['failed'])} files")
                for f in processed_files['failed']:
                    print(f"  ✗ {f}")
            print(f"========================\n")
            
            return results
    except KeyboardInterrupt:
        print("\nKeyboard interrupt detected, shutting down...")
        if global_pool:
            global_pool.terminate()
            global_pool.join()
        raise
    except Exception as e:
        print(f"Error during processing: {e}")
        if global_pool:
            global_pool.terminate()
            global_pool.join()
        raise
