# summarization.py
"""
Transcript summarization module with LLM.
Provides a robust function for summarizing long texts using
intelligent chunking and local language models.

Hybrid version: uses LangChain for text splitting and prompts,
but llama_cpp directly for LLM calls (better performance).
"""

import time
from functools import lru_cache
from typing import Iterator
import os
import multiprocessing

from llama_cpp import Llama
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.prompts import PromptTemplate

from .utils import available_gguf_llms, num_vcpus, s2tw_converter

# Detection of available logical cores
detected_cpus = multiprocessing.cpu_count()
is_hf_spaces = os.environ.get('SPACE_ID') is not None
print(f"Detected vCPUs: {detected_cpus}, Effective vCPUs: {num_vcpus}" + (" (HF Spaces limited)" if is_hf_spaces else ""))


@lru_cache(maxsize=1)
def get_llm(selected_gguf_model: str) -> Llama:
    """Cache and return the LLM model"""
    repo_id, filename = available_gguf_llms[selected_gguf_model]
    return Llama.from_pretrained(
        repo_id=repo_id,
        filename=filename,
        verbose=False,
        n_ctx=4096,
        n_threads=num_vcpus,
        repeat_penalty=2.0,  # Strong penalty to prevent repetition
        temperature=0.05,    # Very low temperature for deterministic output
        top_p=0.8,           # More restrictive nucleus sampling
        top_k=30,            # Smaller top_k for more focused output
    )


def get_language_instruction() -> str:
    """Get universal language instruction for system prompts"""
    return "IMPORTANT: You MUST respond in the EXACT SAME LANGUAGE as the input text. If the input is in French, respond in French. If the input is in Chinese, respond in Chinese. If the input is in Spanish, respond in Spanish. Do not translate to English. Maintain the original language throughout your entire response."


def remove_repetition(text: str, min_length: int = 50) -> str:
    """
    Remove repetitive patterns from generated text.
    Truncates at the first sign of repetition (duplicate 'Healthcare' pattern).
    """
    if len(text) < min_length:
        return text
    
    # Find the first occurrence of 'Healthcare' after the initial good content
    healthcare_pos = text.find('Healthcare')
    if healthcare_pos > 100:  # Make sure we're not truncating too early
        # Look for the pattern where 'Healthcare' appears twice in a row (with possible punctuation)
        double_healthcare = text.find('HealthcareHealthcare', healthcare_pos)
        if double_healthcare > 0:
            # Truncate just before the repetition starts
            return text[:double_healthcare].strip()
        
        # If we find single 'Healthcare' that's not at the beginning, it might be the start of repetition
        if healthcare_pos > 200:  # Reasonable position for good content
            return text[:healthcare_pos].strip()
    
    # Fallback: if text is too long, truncate to reasonable length
    if len(text) > 1000:
        truncated = text[:900]
        last_period = truncated.rfind('.')
        if last_period > 400:
            return truncated[:last_period + 1].strip()
        return truncated.strip()
    
    return text
    """Get universal language instruction for system prompts"""
    return "IMPORTANT: You MUST respond in the EXACT SAME LANGUAGE as the input text. If the input is in French, respond in French. If the input is in Chinese, respond in Chinese. If the input is in Spanish, respond in Spanish. Do not translate to English. Maintain the original language throughout your entire response."


def get_summarization_instructions() -> str:
    """Get comprehensive summarization instructions to prevent common issues"""
    return """You are an expert transcript summarizer. Create clear, concise summaries that capture key points without ANY repetition.

CRITICAL RULES - NEVER DO THESE:
- NEVER repeat words, phrases, or sentences
- NEVER start with "Here is a summary", "Okay", "Voici un résumé", or similar
- NEVER copy text directly from the input
- NEVER repeat the same ideas multiple times

REQUIRED BEHAVIOR:
- Create NEW, ORIGINAL content that summarizes the main ideas
- Keep summaries concise (aim for 25-35% of original length)
- Focus on 2-3 key points maximum
- Use natural, flowing language
- Be direct and to the point
- Maintain factual accuracy"""


def create_text_splitter(chunk_size: int = 4000, chunk_overlap: int = 200) -> RecursiveCharacterTextSplitter:
    """Create a text splitter with intelligent separators"""
    return RecursiveCharacterTextSplitter(
        chunk_size=chunk_size,
        chunk_overlap=chunk_overlap,
        separators=["\n\n", "\n", ". ", " ", ""],
        length_function=len,
    )


def create_chunk_summary_prompt() -> PromptTemplate:
    """Prompt for summarizing an individual chunk"""
    template = """Summarize this part of the transcript while keeping the key points and important information.

Transcript:
{text}

Concise summary:"""
    return PromptTemplate(template=template, input_variables=["text"])


def create_final_summary_prompt() -> PromptTemplate:
    """Prompt for creating the final summary from partial summaries"""
    template = """Here are the summaries of different parts of a transcript.
Create a coherent and synthetic summary of the whole.

{user_prompt}

Partial summaries:
{partial_summaries}

Final summary:"""
    return PromptTemplate(
        template=template,
        input_variables=["user_prompt", "partial_summaries"]
    )


def summarize_chunk(llm: Llama, text: str, prompt_template: PromptTemplate) -> str:
    """Summarize an individual chunk using LangChain for the prompt"""
    try:
        # Use LangChain to format the prompt
        formatted_prompt = prompt_template.format(text=text)

        response = llm.create_chat_completion(
            messages=[
                {"role": "system", "content": f"{get_summarization_instructions()} {get_language_instruction()}"},
                {"role": "user", "content": formatted_prompt}
            ],
            stream=False,
        )
        summary = response['choices'][0]['message']['content']
        # Remove repetition from chunk summary
        cleaned_summary = remove_repetition(summary)
        return s2tw_converter.convert(cleaned_summary)
    except Exception as e:
        print(f"Error during chunk summarization: {e}")
        return f"[Summarization error: {str(e)}]"


def summarize_transcript_langchain(transcript: str, selected_gguf_model: str, prompt_input: str) -> Iterator[str]:
    """
    Hybrid LangChain + llama_cpp version of transcript summarization.

    LangChain advantages used:
    - RecursiveCharacterTextSplitter: intelligent chunking with natural separators
    - PromptTemplate: clean template management
    - More readable and maintainable code

    Keeps llama_cpp for LLM calls (better performance).
    """
    if not transcript or not transcript.strip():
        yield "The transcript is empty."
        return

    try:
        # Component initialization
        llm = get_llm(selected_gguf_model)
        text_splitter = create_text_splitter()
        chunk_prompt = create_chunk_summary_prompt()
        final_prompt = create_final_summary_prompt()

        # Token estimation
        transcript_tokens = len(llm.tokenize(transcript.encode('utf-8')))

        # Direct summary if text is short
        if transcript_tokens <= 2000:
            print(f"[summarize_transcript] Direct summary: {transcript_tokens} tokens")

            # Direct completion without streaming to debug
            completion = llm.create_chat_completion(
                messages=[
                    {"role": "system", "content": f"{get_summarization_instructions()} {get_language_instruction()}"},
                    {"role": "user", "content": f"{prompt_input}\n\n{transcript}"}
                ],
                stream=False,  # Disable streaming to get complete response
            )
            
            full_response = completion['choices'][0]['message']['content']
            # Remove repetition from the response
            cleaned_response = remove_repetition(full_response)
            yield s2tw_converter.convert(cleaned_response)

        # Chunking with LangChain for long texts
        chunks = text_splitter.split_text(transcript)
        print(f"[summarize_transcript] Text divided into {len(chunks)} chunks")

        # Summary of each chunk
        partial_summaries = []
        for i, chunk in enumerate(chunks, 1):
            print(f"Summarizing chunk {i}/{len(chunks)}")
            summary = summarize_chunk(llm, chunk, chunk_prompt)
            partial_summaries.append(summary)

        # Combination and final summary
        combined_summaries = "\n\n".join(partial_summaries)

        # Check combination size
        combined_tokens = len(llm.tokenize(combined_summaries.encode('utf-8')))

        if combined_tokens <= 3500:  # Leave some margin
            print(f"[summarize_transcript] Final summary of {len(partial_summaries)} partial summaries")

            # Use LangChain to format the final prompt
            final_prompt_formatted = final_prompt.format(
                user_prompt=prompt_input,
                partial_summaries=combined_summaries
            )

            # Use non-streaming for final summary to enable repetition removal
            completion = llm.create_chat_completion(
                messages=[
                    {"role": "system", "content": f"{get_summarization_instructions()} {get_language_instruction()}"},
                    {"role": "user", "content": final_prompt_formatted}
                ],
                stream=False,
            )
            
            full_response = completion['choices'][0]['message']['content']
            # Remove repetition from the final response
            cleaned_response = remove_repetition(full_response)
            yield s2tw_converter.convert(cleaned_response)
        else:
            print(f"[summarize_transcript] Combination too long ({combined_tokens} tokens), simplified summary")
            # Fallback: direct summary of the combination
            stream = llm.create_chat_completion(
                messages=[
                    {"role": "system", "content": f"{get_summarization_instructions()} {get_language_instruction()}"},
                    {"role": "user", "content": f"{prompt_input}\n\n{combined_summaries}"}
                ],
                stream=True,
            )

            full_response = ""
            for chunk in stream:
                delta = chunk['choices'][0]['delta']
                if 'content' in delta:
                    full_response += delta['content']
                    yield s2tw_converter.convert(full_response)

    except Exception as e:
        print(f"General error during summarization: {e}")
        yield f"[Error during summarization: {str(e)}]"


def create_title_prompt() -> PromptTemplate:
    """Prompt for generating a document title"""
    template = """Generate a single, concise title for this transcript that captures the main topic or theme. Keep it under 10 words.

Transcript:
{text}

Title:"""
    return PromptTemplate(template=template, input_variables=["text"])


def generate_title(transcript: str, selected_gguf_model: str) -> str:
    """
    Generate a title for the transcript using the selected LLM.
    Returns a concise title that captures the main topic.
    """
    if not transcript or not transcript.strip():
        return "Untitled Document"

    try:
        # Get the LLM
        llm = get_llm(selected_gguf_model)
        title_prompt = create_title_prompt()

        # Use first 2000 tokens for title generation to avoid excessive context
        tokens = llm.tokenize(transcript.encode('utf-8'))
        if len(tokens) > 2000:
            # Truncate to first 2000 tokens and decode back to text
            truncated_tokens = tokens[:2000]
            truncated_text = llm.detokenize(truncated_tokens).decode('utf-8')
        else:
            truncated_text = transcript

        # Format the prompt
        formatted_prompt = title_prompt.format(text=truncated_text)

        # Generate title
        response = llm.create_chat_completion(
            messages=[
                {"role": "system", "content": f"You are an expert at creating single, concise titles for documents and transcripts. Always provide exactly one title, nothing else. {get_language_instruction()}"},
                {"role": "user", "content": formatted_prompt}
            ],
            stream=False,
            max_tokens=20,  # Very short for titles
        )

        title = response['choices'][0]['message']['content'].strip()
        # Clean up the title (remove quotes, extra whitespace)
        title = title.strip('"\'').strip()
        # Apply zh-cn to zh-tw conversion
        title = s2tw_converter.convert(title)
        return title if title else "Untitled Document"

    except Exception as e:
        print(f"Error generating title: {e}")
        return "Untitled Document"


def create_speaker_name_detection_prompt() -> PromptTemplate:
    """Prompt for detecting speaker names from their utterances"""
    template = """Analyze the following utterances from a single speaker and suggest a name for this speaker. Look for:

1. Self-introductions or self-references
2. Names mentioned in context
3. Speech patterns, vocabulary, and topics that might indicate identity
4. Professional titles, roles, or relationships mentioned

Utterances from this speaker:
{text}

Based on the content, suggest a name for this speaker. Consider:
- If the speaker introduces themselves, use that name
- If the speaker is addressed by others, use that name
- If the content suggests a clear identity (e.g., "I'm Dr. Smith", "As CEO", "My name is John")
- If no clear name is evident, suggest "Unknown"

Provide your answer in this exact format:
NAME: [suggested name]
CONFIDENCE: [high/medium/low]
REASON: [brief explanation]

If confidence is "low", the name should not be used."""
    return PromptTemplate(template=template, input_variables=["text"])


def detect_speaker_names(utterances: list, selected_gguf_model: str) -> dict:
    """
    Detect speaker names from diarized utterances using LLM analysis.

    Args:
        utterances: List of utterance dicts with 'speaker', 'text', 'start', 'end' keys
        selected_gguf_model: The LLM model to use for analysis

    Returns:
        Dict mapping speaker_id to detected name info:
        {
            speaker_id: {
                'name': str,
                'confidence': str,  # 'high', 'medium', 'low'
                'reason': str
            }
        }
    """
    if not utterances:
        return {}

    # Group utterances by speaker
    speaker_utterances = {}
    for utt in utterances:
        speaker_id = utt.get('speaker')
        if speaker_id is not None:
            if speaker_id not in speaker_utterances:
                speaker_utterances[speaker_id] = []
            speaker_utterances[speaker_id].append(utt['text'])

    if not speaker_utterances:
        return {}

    try:
        llm = get_llm(selected_gguf_model)
        prompt = create_speaker_name_detection_prompt()

        speaker_names = {}

        for speaker_id, texts in speaker_utterances.items():
            # Combine all utterances for this speaker (limit to reasonable length)
            combined_text = ' '.join(texts)
            if len(combined_text) > 4000:  # Limit context
                combined_text = combined_text[:4000] + '...'

            # Format prompt
            formatted_prompt = prompt.format(text=combined_text)

            # Get LLM response
            response = llm.create_chat_completion(
                messages=[
                    {"role": "system", "content": f"You are an expert at analyzing speech patterns and identifying speaker identities from transcripts. Be precise and only suggest names when you have clear evidence. {get_language_instruction()}"},
                    {"role": "user", "content": formatted_prompt}
                ],
                stream=False,
                max_tokens=100,
            )

            result_text = response['choices'][0]['message']['content'].strip()

            # Parse the response
            name = "Unknown"
            confidence = "low"
            reason = "No clear identification found"

            lines = result_text.split('\n')
            for line in lines:
                if line.startswith('NAME:'):
                    name = line.replace('NAME:', '').strip()
                elif line.startswith('CONFIDENCE:'):
                    confidence = line.replace('CONFIDENCE:', '').strip().lower()
                elif line.startswith('REASON:'):
                    reason = line.replace('REASON:', '').strip()

            # Only include high confidence detections
            if confidence == 'high' and name != "Unknown":
                speaker_names[speaker_id] = {
                    'name': name,
                    'confidence': confidence,
                    'reason': reason
                }

        return speaker_names

    except Exception as e:
        print(f"Error detecting speaker names: {e}")
        return {}


# Alias pour maintenir la compatibilité
summarize_transcript = summarize_transcript_langchain