from wenbi.utils import (
    transcribe,
    parse_subtitle,
    extract_audio_segment,  # replace video_to_audio with extract_audio_segment
    language_detect,
    download_audio,
)
from wenbi.model import rewrite, translate, process_docx, academic
import os
import sys
import dspy
import logging

# Only use package-relative output dir for web interface
OUTPUT_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "output")

# Ensure project root is in sys.path
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))


def is_video_audio_or_url(file_path, url):
    """Check if input is video, audio, or URL"""
    if url.strip():
        return True
    
    if file_path:
        video_extensions = ('.mp4', '.avi', '.mov', '.mkv', '.flv', '.wmv', '.m4v', '.webm')
        audio_extensions = ('.mp3', '.flac', '.aac', '.ogg', '.m4a', '.opus')
        return file_path.lower().endswith(video_extensions + audio_extensions)
    
    return False


def is_text_file(file_path):
    """Check if input is VTT, markdown, or docx file"""
    if not file_path:
        return False
    
    text_extensions = ('.vtt', '.srt', '.ass', '.ssa', '.sub', '.smi', '.txt', '.md', '.markdown', '.docx')
    return file_path.lower().endswith(text_extensions)


def process_input(
    file_path=None,
    url="",
    transcribe_lang="",  # renamed from 'language'
    llm="",
    multi_language=False,
    lang="Chinese",  # Consolidated language parameter
    output_dir="",
    chunk_length=20,
    max_tokens=50000,
    timeout=3600,
    temperature=0.1,
    transcribe_model="large-v3",
    timestamp=None,
    output_wav="",
    subcommand=None,  # New parameter to specify which subcommand to use
    cite_timestamps=False,  # New parameter for timestamp citation
    verbose=False,  # New parameter for verbose logging
):
    """Process input with logic:
    1. If input is video/audio/URL: convert to WAV -> transcribe to VTT -> process with subcommand
    2. If input is VTT/markdown/docx: directly process with subcommand
    """
    logger = logging.getLogger(__name__)
    
    if verbose:
        logger.debug("=== Starting process_input ===")
        logger.debug(f"Input file: {file_path}")
        logger.debug(f"Input URL: {url}")
        logger.debug(f"Subcommand: {subcommand}")
        logger.debug(f"LLM: {llm}")
    # Use current directory for CLI, package directory for web interface
    out_dir = (
        output_dir
        if output_dir.strip()
        else OUTPUT_DIR
        if __name__ == "__main__"
        else os.getcwd()
    )
    os.makedirs(out_dir, exist_ok=True)

    if not file_path and not url:
        return "Error: No input provided", None, None, None

    # Check if input is video/audio/URL
    if is_video_audio_or_url(file_path, url):
        # Step 1: Convert to WAV
        try:
            if url:
                file_path = download_audio(url.strip(), output_dir=out_dir, timestamp=timestamp, output_wav=output_wav)
            elif file_path:
                # Use extract_audio_segment for all audio/video files except .wav
                if file_path.lower().endswith(('.mp4', '.avi', '.mov', '.mkv', '.flv', '.wmv', '.m4v',
                                             '.mp3', '.flac', '.aac', '.ogg', '.m4a', '.webm', '.opus')):
                    file_path = extract_audio_segment(file_path, timestamp, out_dir, output_wav=output_wav)
                # If .wav, do nothing (already correct format)
        except Exception as e:
            print(f"Error converting to WAV: {e}")
            return "Error: Failed to convert to WAV", None, None, None

        # Step 2: Transcribe to VTT
        try:
            if multi_language:
                from wenbi.mutilang import transcribe_multi_speaker, speaker_vtt

                base_name = os.path.splitext(os.path.basename(file_path))[0]
                transcriptions = transcribe_multi_speaker(
                    file_path, model_size=transcribe_model
                )
                vtt_files = speaker_vtt(
                    transcriptions, output_dir=out_dir, base_filename=base_name
                )
            else:
                lang_code = transcribe_lang if transcribe_lang.strip() else None
                vtt_file, _ = transcribe(
                    file_path,
                    language=lang_code,
                    output_dir=out_dir,
                    model_size=transcribe_model,
                )
                vtt_files = {None: vtt_file}
        except Exception as e:
            print(f"Error during transcription: {e}")
            return "Error: Failed during transcription", None, None, None

        # Step 3: Process VTT files with subcommand
        if subcommand:
            # Process each VTT file with the specified subcommand
            final_outputs = {}
            for speaker, vtt_file in vtt_files.items():
                try:
                    if subcommand == "translate":
                        result = translate(
                            vtt_file,
                            output_dir=out_dir,
                            translate_language=lang,
                            llm=llm,
                            chunk_length=chunk_length,
                            max_tokens=max_tokens,
                            timeout=timeout,
                            temperature=temperature,
                            cite_timestamps=cite_timestamps,
                        )
                    elif subcommand == "rewrite":

                        rewrite_text, rewrite_file = rewrite(
vtt_file,
                            output_dir=out_dir,
                            llm=llm,
                            rewrite_language=lang,
                            chunk_length=chunk_length,
                            max_tokens=max_tokens,
                            timeout=timeout,
                            temperature=temperature,
                            cite_timestamps=cite_timestamps,

                        )

                        result = (rewrite_text, rewrite_file)
                    elif subcommand == "academic":
                        result = academic(
                            vtt_file,
                            output_dir=out_dir,
                            llm=llm,
                            academic_lang=lang,
                            chunk_length=chunk_length,
                            max_tokens=max_tokens,
                            timeout=timeout,
                            temperature=temperature,
                            cite_timestamps=cite_timestamps,
                        )
                    else:
                        result = "Error: Unknown subcommand"
                    
                    final_outputs[speaker if speaker else "output"] = result
                except Exception as e:
                    print(f"Error processing VTT with {subcommand}: {e}")
                    return f"Error: Failed during {subcommand} processing", None, None, None

            if multi_language:
                return final_outputs
            else:
                result = final_outputs["output"]
                base_name = os.path.splitext(os.path.basename(vtt_file))[0]
                # Handle tuple result format for rewrite

                if isinstance(result, tuple) and len(result) == 2:

                    text_content, file_path = result

                    return text_content, file_path, None, base_name

                else:

                    return result, result, None, base_name
        else:
            # Original behavior for backward compatibility
            return _process_vtt_original_logic(vtt_files, out_dir, lang, llm, chunk_length, max_tokens, timeout, temperature, multi_language)

    # Check if input is text file (VTT, markdown, docx)
    elif is_text_file(file_path):
        if subcommand:
            # Process text file directly with subcommand
            try:
                if subcommand == "translate":
                    result = translate(
                        file_path,
                        output_dir=out_dir,
                        translate_language=lang,
                        llm=llm,
                        chunk_length=chunk_length,
                        max_tokens=max_tokens,
                        timeout=timeout,
                        temperature=temperature,
                        cite_timestamps=cite_timestamps,
                    )
                elif subcommand == "rewrite":

                    rewrite_text, rewrite_file = rewrite(
file_path,
                        output_dir=out_dir,
                        llm=llm,
                        rewrite_language=lang,
                        chunk_length=chunk_length,
                        max_tokens=max_tokens,
                        timeout=timeout,
                        temperature=temperature,
                        cite_timestamps=cite_timestamps,

                    )

                    result = (rewrite_text, rewrite_file)
                elif subcommand == "academic":
                    result = academic(
                        file_path,
                        output_dir=out_dir,
                        llm=llm,
                        academic_lang=lang,
                        chunk_length=chunk_length,
                        max_tokens=max_tokens,
                        timeout=timeout,
                        temperature=temperature,
                        cite_timestamps=cite_timestamps,
                    )
                else:
                    return "Error: Unknown subcommand", None, None, None
                
                base_name = os.path.splitext(os.path.basename(file_path))[0]
                # Handle tuple result format for rewrite

                if isinstance(result, tuple) and len(result) == 2:

                    text_content, file_path = result

                    return text_content, file_path, None, base_name

                else:

                    return result, result, None, base_name
            except Exception as e:
                print(f"Error processing text file with {subcommand}: {e}")
                return f"Error: Failed during {subcommand} processing", None, None, None
        else:
            # Original behavior for backward compatibility
            return _process_text_file_original_logic(file_path, out_dir, lang, llm, chunk_length, max_tokens, timeout, temperature)
    
    else:
        return "Error: Unsupported file type", None, None, None


def _process_vtt_original_logic(vtt_files, out_dir, lang, llm, chunk_length, max_tokens, timeout, temperature, multi_language):
    """Original VTT processing logic for backward compatibility"""
    final_outputs = {}
    
    # Defensive: handle both dict and list (or single file)
    if isinstance(vtt_files, dict):
        vtt_iter = vtt_files.items()
    elif isinstance(vtt_files, list):
        vtt_iter = enumerate(vtt_files)
    else:
        # fallback: treat as single file
        vtt_iter = [(None, vtt_files)]

    for speaker, vtt_file in vtt_iter:
        if not multi_language:
            base_name = os.path.splitext(os.path.basename(vtt_file))[0]
            csv_file = os.path.join(out_dir, f"{base_name}.csv")
            parse_subtitle(vtt_file).to_csv(
                csv_file, index=True, encoding="utf-8"
            )
            print(f"CSV file '{csv_file}' created successfully.")

        detected_lang = language_detect(vtt_file)
        print(f"Detected language for {speaker or 'input'}: {detected_lang}")

        # Determine which LLM function to call based on detected language and target language
        if detected_lang == "zh" and lang.lower() == "chinese":
            # Use rewrite for Chinese to Chinese translation/rewriting
            rewrite_text, rewrite_file = rewrite(
vtt_file,
                output_dir=out_dir,
                llm=llm,
                rewrite_language=lang,
                chunk_length=chunk_length,
                max_tokens=max_tokens,
                timeout=timeout,
                temperature=temperature,

            )

            output = (rewrite_text, rewrite_file)
        else:
            # Use translate for other languages or when target is not Chinese
            output = translate(
                vtt_file,
                output_dir=out_dir,
                translate_language=lang,
                llm=llm,
                chunk_length=chunk_length,
                max_tokens=max_tokens,
                timeout=timeout,
                temperature=temperature,
            )
        final_outputs[speaker if speaker else "output"] = output

    if multi_language:
        return final_outputs
    else:
        result = final_outputs["output"]
        # Handle tuple result format for rewrite

        if isinstance(result, tuple) and len(result) == 2:

            text_content, file_path = result

            return text_content, file_path, csv_file, base_name

        else:

            return result, result, csv_file, base_name


def _process_text_file_original_logic(file_path, out_dir, lang, llm, chunk_length, max_tokens, timeout, temperature):
    """Original text file processing logic for backward compatibility"""
    # If llm is provided and it's intended for academic processing,
    # and the input is a text file, go directly to academic processing.
    if llm and file_path.lower().endswith(('.docx', '.txt', '.md')):
        academic_text = academic(
            file_path,
            output_dir=out_dir,
            llm=llm,
            academic_lang=lang,
            chunk_length=chunk_length,
            max_tokens=max_tokens,
            timeout=timeout,
            temperature=temperature,
        )

        base_name = os.path.splitext(os.path.basename(file_path))[0]
        output_file = os.path.join(out_dir, f"{base_name}_academic.md")
        with open(output_file, "w", encoding="utf-8") as f:
            f.write(academic_text)

        return academic_text, output_file, None, base_name
    
    return "Error: Cannot process this file type without subcommand", None, None, None



