#!/usr/bin/env python3
import os
import glob
import logging
import json
import os,sys,spacy
from multiprocessing import Process
import speech_recognition as sr
from abstract_utilities import *
from moviepy.editor import *
import moviepy.editor as mp
from pydub import AudioSegment
from datetime import timedelta
from .video_utils import derive_all_video_meta
from collections import Counter
logger = get_logFile('vid_to_aud')
nlp = spacy.load("en_core_web_sm")
# Set up logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler('vid_to_aud.log'),
        logging.StreamHandler()
    ]
)

# Initialize recognizer
r = sr.Recognizer()

def format_timestamp(ms):
    """Convert milliseconds to a formatted timestamp (HH:MM:SS.mmm)."""
    td = timedelta(milliseconds=ms)
    hours, remainder = divmod(td.seconds, 3600)
    minutes, seconds = divmod(remainder, 60)
    milliseconds = td.microseconds // 1000
    return f"{hours:02d}:{minutes:02d}:{seconds:02d}.{milliseconds:03d}"

def extract_audio_from_video(video_path, audio_path):
    """Extract audio from a video file using moviepy."""
    try:
        logging.info(f"Extracting audio from {video_path} to {audio_path}")
        video = mp.VideoFileClip(video_path)
        video.audio.write_audiofile(audio_path)
        video.close()
        if not os.path.exists(audio_path):
            raise FileNotFoundError(f"Audio file {audio_path} was not created.")
        logging.info(f"Audio extracted successfully: {audio_path}")
        return audio_path
    except Exception as e:
        logging.error(f"Error extracting audio from {video_path}: {e}")
        return None
def create_key_value(json_obj,key,value):
    if key not in json_obj:
        json_obj[key]=value
    return json_obj
def check_if_transcribed_block(audio_text, start_time, end_time):
    """
    Checks if a given time block is already fully transcribed.
    
    If fully transcribed: returns True.
    If partially transcribed or missing: returns (updated_start_time, updated_end_time).
    """
    block_covered = False
    for entry in audio_text:
        entry_start = parse_timestamp(entry.get("start_time", "0:00"))
        entry_end = parse_timestamp(entry.get("end_time", "0:00"))
        entry_text = entry.get("text", "").strip()
        
        if entry_start <= start_time and entry_end >= end_time and entry_text:
            block_covered = True
            break
    
    if block_covered:
        return True
    else:
        # Figure out what part was covered, and return updated range
        uncovered_start = start_time
        uncovered_end = end_time

        for entry in audio_text:
            entry_start = parse_timestamp(entry.get("start_time", "0:00"))
            entry_end = parse_timestamp(entry.get("end_time", "0:00"))
            entry_text = entry.get("text", "").strip()

            # Skip if entry has no content
            if not entry_text:
                continue

            # Clip off already transcribed part at the beginning
            if entry_start <= uncovered_start < entry_end:
                uncovered_start = entry_end

            # Clip off from the end
            if entry_start < uncovered_end <= entry_end:
                uncovered_end = entry_start

        return (uncovered_start, uncovered_end)
def transcribe_audio_file(audio_path, json_data, chunk_length_ms=60000,summarizer=None):
    """Transcribe audio file in chunks and save as text and time-blocked JSON."""
    try:
        logging.info(f"Transcribing audio: {audio_path}")
        audio = AudioSegment.from_wav(audio_path)
        chunks = [audio[i:i + chunk_length_ms] for i in range(0, len(audio), chunk_length_ms)]
        full_text = ""
        json_data = create_key_value(json_data,'audio_text',[])
        file_path = json_data['info_path']
        for i, chunk in enumerate(chunks):
            start_time = i * chunk_length_ms
            end_time = min((i + 1) * chunk_length_ms, len(audio))

            result = check_if_transcribed_block(audio_text, start_time, end_time)

            if result is True:
                logging.info(f"Chunk {i+1} already done.")
                continue
            elif isinstance(result, tuple):
                updated_start, updated_end = result

                # Adjust actual chunk
                actual_chunk = audio[updated_start:updated_end]
                chunk_path = f"chunk_{i}_partial.wav"
                actual_chunk.export(chunk_path, format="wav")
                
                logging.info(f"Processing updated chunk {i+1}: {format_timestamp(updated_start)} - {format_timestamp(updated_end)}")

                with sr.AudioFile(chunk_path) as source:
                    r.adjust_for_ambient_noise(source)
                    audio_data = r.record(source)
                    try:
                        text = r.recognize_google(audio_data)
                    except sr.UnknownValueError:
                        text = ""
                    except sr.RequestError as e:
                        logging.error(f"API error: {e}")
                        text = ""

                json_data['audio_text'].append({
                    "start_time": format_timestamp(updated_start),
                    "end_time": format_timestamp(updated_end),
                    "text": text
                })
                safe_dump_to_file(data=json_data,file_path=file_path) 
                os.remove(chunk_path)

        # Save plain text
        if full_text:
            full_text = full_text.strip()
        text_result= get_voice(full_text, text=full_text)
        json_data = create_key_value(json_data,'full_text',text_result)
        json_data = create_key_value(json_data,'summary',None)
        json_data = create_key_value(json_data,'keywords',extract_keywords_nlp(text_result))
        if summarizer:
            try:
                json_data['summary'] = summarizer(text_result, max_length=160, min_length=40)
            except Exception as e:
                logging.error(f"Error getting summary for {json_data['filename']}: {e}")
        
        return json_data
    except Exception as e:
        logging.error(f"Error transcribing {audio_path}: {e}")
        return json_data

def get_voice(voice, text=None):
    """Save transcribed text to a file."""
    text = text or ''
    if voice:
        text = text + '\n' + str(voice) if text else str(voice)
       
    return text
def extract_keywords_nlp(strings, top_n=5):

    # Process with spaCy
    doc = nlp(strings)
    
    # Extract nouns, proper nouns, and entities
    word_counts = Counter()
    for token in doc:
        if token.pos_ in ["NOUN", "PROPN"] and not token.is_stop and len(token.text) > 2:
            word_counts[token.text] += 1
    
    # Extract multi-word entities (e.g., "Elon Musk")
    entity_counts = Counter(ent.text.lower() for ent in doc.ents if len(ent.text.split()) > 1)
    
    # Combine and rank
    combined_counts = word_counts + entity_counts
    top_keywords = [word for word, count in combined_counts.most_common(top_n)]
    
    return top_keywords

def initiate_process(target,*args):
    p = Process(target=target, args=args)
    p.start()
    logging.info(f"Started process for: {args}")
def transcribe_all_video_paths(directory=None,output_dir=None,remove_phrases=None,summarizer=None,get_vid_data=None):
    get_vid_data = get_vid_data or False
    remove_phrases=remove_phrases or []
    directory = directory or os.getcwd()
    output_dir = output_dir if output_dir else make_dirs(directory,'text_dir')
    paths = glob.glob(path_join(directory, '**', '**'), recursive=True)
    paths = [file_path for file_path in paths if confirm_type(file_path,media_types=get_media_types(['video']))]
    video_paths = get_all_file_types(directory=directory,types='video') or get_all_file_types(directory=abs_dirname,types='videos')
    for video_path in video_paths:
        transcribe_audio(video_path,output_dir=output_dir,remove_phrases=remove_phrases,summarizer=summarizer,get_vid_data=get_vid_data)
def transcribe_audio(video_path,output_dir=None,remove_phrases=None,summarizer=None,get_vid_data=False):
    get_vid_data = get_vid_data or False
    remove_phrases=remove_phrases or []
    dirname = os.path.dirname(video_path)
    basename = os.path.basename(video_path)
    filename,ext = os.path.splitext(basename)
    video_directory = make_dirs(output_dir,filename)
    info_path = os.path.join(video_directory,'info.json')
    video_text_path = os.path.join(video_directory,'video_text.json')
    audio_path = os.path.join(video_directory,'audio.wav')
    video_json_path = os.path.join(video_directory,'video_json.json')
    if get_vid_data:
        args_tuple = (video_path, video_directory, video_json_path, [], '', filename,remove_phrases,summarizer)
        initiate_process(derive_all_video_meta, *args_tuple)
    info = {}
    
    info['video_path']=video_path
    info['info_directory']=video_directory
    info['info_path']=info_path
    info['filename']=filename
    info['ext']=ext
    info['audio_path']=audio_path
    info['video_json']=video_json_path
    extract_audio_from_video(video_path=info['video_path'],audio_path=info['audio_path'])
    info = transcribe_audio_file(audio_path=info['audio_path'],json_data=info,summarizer=summarizer)
    safe_dump_to_file(data=info,file_path=info['info_path']) 
   
    
