import dspy
import os
from wenbi.utils import segment

def configure_lm(model_string, **kwargs):
    if not model_string:
        model_string = "ollama/qwen3"

    parts = model_string.strip().split("/")
    provider = parts[0].lower() if parts else ""

    config = kwargs
    if provider == "ollama":
        config.update({
            "base_url": "http://localhost:11434",
            "model": model_string,
        })
        lm = dspy.LM(**config)
    elif provider == "openai":
        api_key = os.getenv("OPENAI_API_KEY")
        if not api_key:
            raise ValueError("OPENAI_API_KEY environment variable not set.")
        config.update({
            "api_base": "https://api.openai.com/v1",
            "model": model_string.replace("openai/", ""),
            "api_key": api_key,
        })
        lm = dspy.OpenAI(**config)
    elif provider == "gemini":
        api_key = os.getenv("GOOGLE_API_KEY") or os.getenv("GOOGLE_API_KEY_JSON")
        if not api_key:
            raise ValueError("GOOGLE_API_KEY or GOOGLE_API_KEY_JSON environment variable not set.")

        # Extract the actual model name (e.g., "gemini-2.5-flash" from "gemini/gemini-2.5-flash")
        model_name = model_string.split("/", 1)[1] if "/" in model_string else model_string

        # Use the correct format for LiteLLM Gemini integration
        config.update({
            "model": f"gemini/{model_name}",
            "api_key": api_key,
        })
        lm = dspy.LM(**config)
    else:
        config.update({"model": model_string})
        lm = dspy.LM(**config)

    dspy.configure(lm=lm)
    return lm

def translate(
    vtt_path,
    output_dir=None,
    translate_language="Chinese",
    llm="",
    translate_llm="",
    chunk_length=8,
    max_tokens=50000,
    timeout=3600,
    temperature=0.1,
    cite_timestamps=False,
):
    segmented_text = segment(vtt_path, sentence_count=chunk_length, cite_timestamps=cite_timestamps)
    
    # Use translate_llm if provided, otherwise fall back to llm
    model_id = translate_llm or llm or "ollama/qwen3"
    configure_lm(model_id, max_tokens=max_tokens, timeout_s=timeout, temperature=temperature)

    class Translate(dspy.Signature):
        english_text = dspy.InputField(desc="English text to translate")
        translated_text = dspy.OutputField(
            desc=f"Translation into {translate_language}"
        )

    translator = dspy.ChainOfThought(Translate)
    translated_pairs = []

    if cite_timestamps:
        # Handle timestamp sections specially
        sections = segmented_text.split("\n\n")
        
        for section in sections:
            if section.strip():
                # Check if this section starts with a timestamp header
                if section.startswith("### **") and "**" in section:
                    # Extract timestamp header and content
                    lines = section.split("\n")
                    timestamp_header = lines[0] if lines else ""
                    content = "\n".join(lines[1:]).strip() if len(lines) > 1 else ""
                    
                    if content:
                        # Translate only the content, not the timestamp header
                        response = translator(english_text=content)
                        translated_pairs.append(
                            f"{timestamp_header}\n\n# Original\n{content}\n\n# {translate_language}\n{response.translated_text}\n\n---\n"
                        )
                    else:
                        translated_pairs.append(f"{timestamp_header}\n\n---\n")
                else:
                    # Regular paragraph without timestamp header
                    response = translator(english_text=section)
                    translated_pairs.append(
                        f"# Original\n{section}\n\n# {translate_language}\n{response.translated_text}\n\n---\n"
                    )
        
        markdown_content = "\n".join(translated_pairs)
    else:
        # Original behavior for non-timestamp processing
        paragraphs = segmented_text.split("\n\n")
        for para in paragraphs:
            if para.strip():
                response = translator(english_text=para)
                translated_pairs.append(
                    f"# English\n{para}\n\n# {translate_language}\n{response.translated_text}\n\n---\n"
                )
        
        markdown_content = "\n".join(translated_pairs)
    
    output_file = os.path.splitext(vtt_path)[0] + "_bilingual.md"
    with open(output_file, "w", encoding="utf-8") as f:
        f.write(markdown_content)

    return output_file


def rewrite(
    file_path,
    output_dir=None,
    llm="",
    rewrite_llm="",
    rewrite_lang="Chinese",
    chunk_length=8,
    max_tokens=50000,
    timeout=3600,
    temperature=0.1,
    cite_timestamps=False,
):
    segmented_text = segment(file_path, sentence_count=chunk_length, cite_timestamps=cite_timestamps)
    
    # Use rewrite_llm if provided, otherwise fall back to llm
    model_id = rewrite_llm or llm or "ollama/qwen3"
    configure_lm(model_id, max_tokens=max_tokens, timeout_s=timeout, temperature=temperature)

    if cite_timestamps:
        # Handle timestamp sections specially
        rewritten_paragraphs = []
        sections = segmented_text.split("\n\n")
        
        for section in sections:
            if section.strip():
                # Check if this section starts with a timestamp header
                if section.startswith("### **") and "**" in section:
                    # Extract timestamp header and content
                    lines = section.split("\n")
                    timestamp_header = lines[0] if lines else ""
                    content = "\n".join(lines[1:]).strip() if len(lines) > 1 else ""
                    
                    if content:
                        # Rewrite only the content, not the timestamp header
                        class ParaRewrite(dspy.Signature):
                            """
                            Rewrite this text in {rewrite_lang}, add punctuation, grammar corrected, proofread, converting from spoken to written form
                            while preserving the meaning. Ensure the rewritten text is at least 95% of the original length.
                            """
                            text: str = dspy.InputField(
                                desc=f"Spoken text to rewrite in {rewrite_lang}"
                            )
                            rewritten: str = dspy.OutputField(
                                desc=f"Rewritten paragraph in {rewrite_lang}"
                            )
                        rewrite = dspy.ChainOfThought(ParaRewrite)
                        response = rewrite(text=content)
                        rewritten_paragraphs.append(f"{timestamp_header}\n\n{response.rewritten}")
                    else:
                        rewritten_paragraphs.append(timestamp_header)
                else:
                    # Regular paragraph without timestamp header
                    class ParaRewrite(dspy.Signature):
                        """
                        Rewrite this text in {rewrite_lang}, add punctuation, grammar corrected, proofread, converting from spoken to written form
                        while preserving the meaning. Ensure the rewritten text is at least 95% of the original length.
                        """
                        text: str = dspy.InputField(
                            desc=f"Spoken text to rewrite in {rewrite_lang}"
                        )
                        rewritten: str = dspy.OutputField(
                            desc=f"Rewritten paragraph in {rewrite_lang}"
                        )
                    rewrite = dspy.ChainOfThought(ParaRewrite)
                    response = rewrite(text=section)
                    rewritten_paragraphs.append(response.rewritten)
        
        rewritten_text = "\n\n".join(rewritten_paragraphs)
    else:
        # Original behavior for non-timestamp processing
        paragraphs = segmented_text.split("\n\n")
        rewritten_paragraphs = []
        for para in paragraphs:
            if para.strip():
                class ParaRewrite(dspy.Signature):
                    """
                    Rewrite this text in {rewrite_lang}, add punctuation, grammar corrected, proofread, converting from spoken to written form
                    while preserving the meaning. Ensure the rewritten text is at least 95% of the original length.
                    """
                    text: str = dspy.InputField(
                        desc=f"Spoken text to rewrite in {rewrite_lang}"
                    )
                    rewritten: str = dspy.OutputField(
                        desc=f"Rewritten paragraph in {rewrite_lang}"
                    )
                rewrite = dspy.ChainOfThought(ParaRewrite)
                response = rewrite(text=para)
                rewritten_paragraphs.append(response.rewritten)
        
        rewritten_text = "\n\n".join(rewritten_paragraphs)
    
    if output_dir:
        base_name = os.path.splitext(os.path.basename(file_path))[0]
        out_file = os.path.join(output_dir, f"{base_name}_rewritten.md")
        with open(out_file, "w", encoding="utf-8") as f:
            f.write(rewritten_text)
    else:
        out_file = None

    return rewritten_text


def academic(
    file_path,
    output_dir=None,
    llm="",
    academic_lang="English",
    chunk_length=8,
    max_tokens=50000,
    timeout=3600,
    temperature=0.1,
    cite_timestamps=False,
):
    # Read the original text
    segmented_text = segment(file_path, sentence_count=chunk_length, cite_timestamps=cite_timestamps)
    
    model_id = llm if llm else "ollama/qwen3"
    configure_lm(model_id, max_tokens=max_tokens, timeout_s=timeout, temperature=temperature)

    if cite_timestamps:
        # Handle timestamp sections specially
        academic_paragraphs = []
        sections = segmented_text.split("\n\n")
        
        for section in sections:
            if section.strip():
                # Check if this section starts with a timestamp header
                if section.startswith("### **") and "**" in section:
                    # Extract timestamp header and content
                    lines = section.split("\n")
                    timestamp_header = lines[0] if lines else ""
                    content = "\n".join(lines[1:]).strip() if len(lines) > 1 else ""
                    
                    if content:
                        # Rewrite only the content, not the timestamp header
                        class AcademicRewrite(dspy.Signature):
                            """
                            Rewrite this text in formal academic style in {academic_lang}. Focus on:
                            1. Using scholarly vocabulary and formal language
                            2. keep the original meaning intact and 96% of same length of words.
                            3. Do not change the citation format.
                            4. Avoiding colloquialisms and informal expressions
                            5. Ensuring logical flow and academic structure
                            """
                            text: str = dspy.InputField(
                                desc=f"Text to rewrite in academic {academic_lang}"
                            )
                            academic: str = dspy.OutputField(
                                desc=f"Academic rewritten text in {academic_lang}"
                            )
                        academic_rewrite = dspy.ChainOfThought(AcademicRewrite)
                        response = academic_rewrite(text=content)
                        academic_paragraphs.append(f"{timestamp_header}\n\n{response.academic}")
                    else:
                        academic_paragraphs.append(timestamp_header)
                else:
                    # Regular paragraph without timestamp header
                    class AcademicRewrite(dspy.Signature):
                        """
                        Rewrite this text in formal academic style in {academic_lang}. Focus on:
                        1. Using scholarly vocabulary and formal language
                        2. keep the original meaning intact and 96% of same length of words.
                        3. Do not change the citation format.
                        4. Avoiding colloquialisms and informal expressions
                        5. Ensuring logical flow and academic structure
                        """
                        text: str = dspy.InputField(
                            desc=f"Text to rewrite in academic {academic_lang}"
                        )
                        academic: str = dspy.OutputField(
                            desc=f"Academic rewritten text in {academic_lang}"
                        )
                    academic_rewrite = dspy.ChainOfThought(AcademicRewrite)
                    response = academic_rewrite(text=section)
                    academic_paragraphs.append(response.academic)
        
        academic_text = "\n\n".join(academic_paragraphs)
    else:
        # Original behavior for non-timestamp processing
        paragraphs = segmented_text.split("\n\n")
        academic_paragraphs = []
        for para in paragraphs:
            if para.strip():  # Skip empty paragraphs
                class AcademicRewrite(dspy.Signature):
                    """
                    Rewrite this text in formal academic style in {academic_lang}. Focus on:
                    1. Using scholarly vocabulary and formal language
                    2. keep the original meaning intact and 96% of same length of words.
                    3. Do not change the citation format.
                    4. Avoiding colloquialisms and informal expressions
                    5. Ensuring logical flow and academic structure
                    """
                    text: str = dspy.InputField(
                        desc=f"Text to rewrite in academic {academic_lang}"
                    )
                    academic: str = dspy.OutputField(
                        desc=f"Academic rewritten text in {academic_lang}"
                    )

                academic_rewrite = dspy.ChainOfThought(AcademicRewrite)
                response = academic_rewrite(text=para)
                academic_paragraphs.append(response.academic)
        
        academic_text = "\n\n".join(academic_paragraphs)
    
    if output_dir:
        base_name = os.path.splitext(os.path.basename(file_path))[0]
        out_file = os.path.join(output_dir, f"{base_name}_academic.md")
        with open(out_file, "w", encoding="utf-8") as f:
            f.write(academic_text)
        
        # Create comparison file using Redline
        try:
            import redlines
            
            # Create a comparison between original and academic text
            diff = redlines.Redlines(segmented_text, academic_text)
            compare_output = diff.output_markdown
            
            # Save comparison file
            compare_file = os.path.join(output_dir, f"{base_name}_compare.md")
            with open(compare_file, "w", encoding="utf-8") as f:
                f.write("# Academic Writing Comparison\n\n")
                f.write("This document shows the changes made during academic rewriting.\n\n")
                f.write("## Legend:\n")
                f.write("- **Red text with strikethrough**: Deleted content\n")
                f.write("- **Green text**: Added content\n\n")
                f.write("## Comparison:\n\n")
                f.write(compare_output)
            
            print(f"Comparison file created: {compare_file}")
            
        except ImportError:
            print("Warning: redlines library not available, skipping comparison file")
        except Exception as e:
            print(f"Warning: Could not create comparison file: {e}")
    else:
        out_file = None

    return academic_text


def process_docx(
    file_path,
    output_dir=None,
    llm="",
    academic_lang="English",
    chunk_length=8,
    max_tokens=50000,
    timeout=3600,
    temperature=0.1,
):
    academic_text = academic(
        file_path,
        output_dir=output_dir,
        llm=llm,
        academic_lang=academic_lang,
        chunk_length=chunk_length,
        max_tokens=max_tokens,
        timeout=timeout,
        temperature=temperature,
    )

    if output_dir:
        from docx import Document
        new_doc = Document()
        for para in academic_text.split('\n\n'):
            if para.strip():
                new_doc.add_paragraph(para.strip())

        base_name = os.path.splitext(os.path.basename(file_path))[0]
        docx_out = os.path.join(output_dir, f"{base_name}_academic.docx")
        new_doc.save(docx_out)
    else:
        docx_out = None

    return academic_text, docx_out
