"""
Text summarization module with support for chunking long texts using API-based LLMs.
"""

# Handle imports for both running from src and from parent directory
try:
    from model_loader import get_llm_client
    from config import (
        MAX_WORDS_PER_CHUNK, 
        MAX_TOKENS,
        SUMMARY_MAX_LENGTH,
        SUMMARY_MIN_LENGTH,
        LLM_TEMPERATURE
    )
except ImportError:
    from .model_loader import get_llm_client
    from .config import (
        MAX_WORDS_PER_CHUNK, 
        MAX_TOKENS,
        SUMMARY_MAX_LENGTH,
        SUMMARY_MIN_LENGTH,
        LLM_TEMPERATURE
    )


def chunk_text(text, max_words=MAX_WORDS_PER_CHUNK):
    """
    Splits text into chunks of specified word count.
    
    Args:
        text (str): Input text to chunk
        max_words (int): Maximum words per chunk
        
    Returns:
        list: List of text chunks
    """
    words = text.split()
    chunks = []
    for i in range(0, len(words), max_words):
        chunk = " ".join(words[i:i+max_words])
        chunks.append(chunk)
    return chunks


def summarize_with_api(text, max_tokens=MAX_TOKENS):
    """
    Summarizes text using API-based LLMs.
    
    Args:
        text (str): Text to summarize
        max_tokens (int): Maximum tokens for summary
        
    Returns:
        str: Summary text
    """
    client, model_name, provider = get_llm_client()
    
    system_prompt = "You are a professional summarizer. Create a concise, comprehensive summary of the provided text while maintaining key information and context."
    user_prompt = f"Summarize the following text:\n\n{text}"
    
    try:
        if provider == "mistral":
            response = client.chat.complete(
                model=model_name,
                messages=[
                    {"role": "system", "content": system_prompt},
                    {"role": "user", "content": user_prompt}
                ],
                temperature=LLM_TEMPERATURE,
                max_tokens=max_tokens
            )
            summary = response.choices[0].message.content
        else:
            # OpenAI and Groq
            response = client.chat.completions.create(
                model=model_name,
                messages=[
                    {"role": "system", "content": system_prompt},
                    {"role": "user", "content": user_prompt}
                ],
                temperature=LLM_TEMPERATURE,
                max_tokens=max_tokens
            )
            summary = response.choices[0].message.content
        
        # Print word count for all providers
        word_count = len(summary.split())
        print(f"{word_count} words in summary.")
        return summary
        
    except Exception as e:
        print(f"Error in API summarization: {e}")
        raise


def summarize_text(transcription, max_tokens=MAX_TOKENS):
    """
    Summarizes transcription text, handling long texts with chunking.
    
    Args:
        transcription (str): Input transcription text
        max_tokens (int): Maximum tokens for summary
        
    Returns:
        str: Summarized text
    """
    word_count = len(transcription.split())
    
    if word_count < 1000:
        print("Summarizing transcription...")
        return summarize_with_api(transcription, max_tokens=max_tokens)

    # For longer transcriptions, split into chunks, summarize each, then combine summaries
    print(f"Transcription is long ({word_count} words); splitting into chunks...")
    chunks = chunk_text(transcription, max_words=MAX_WORDS_PER_CHUNK)
    chunk_summaries = []
    
    for i, chunk in enumerate(chunks):
        print(f"Summarizing chunk {i+1}/{len(chunks)}...")
        try:
            summary = summarize_with_api(chunk, max_tokens=max_tokens)
            chunk_summaries.append(summary)
        except Exception as e:
            print(f"Error summarizing chunk {i+1}: {e}")
            chunk_summaries.append("")  # Optionally handle failed chunks
    
    combined_summary = " ".join(chunk_summaries)
    combined_word_count = len(combined_summary.split())
    
    print(f"Creating final summary from {len(chunk_summaries)} chunk summaries ({combined_word_count} words total)...")
    
    # If combined summary is already short enough, return it
    if combined_word_count < 300:
        print(f"Combined summary is already concise ({combined_word_count} words), using as-is.")
        return combined_summary
    
    # Create a final summary from all chunk summaries with higher token limit
    try:
        # Use 2x tokens for final summary since we're condensing multiple summaries
        final_max_tokens = min(max_tokens * 2, 1000)
        final_summary = summarize_with_api(
            combined_summary, 
            max_tokens=final_max_tokens
        )
        
        # Check if final summary is empty (edge case)
        if not final_summary or len(final_summary.strip()) == 0:
            print("Warning: Final summary was empty, using combined chunk summaries instead.")
            return combined_summary
            
        return final_summary
    except Exception as e:
        print(f"Error in final summarization: {e}")
        print("Falling back to combined chunk summaries.")
        return combined_summary
