from .config import *
from .imports import (
    os,
    re,
    quote,
    Counter,
    List,
    Optional,
    get_torch,
    get_transformers
    )

SUMMARIZER_DIR: str = DEFAULT_PATHS["summarizer"]
FLAN_MODEL_NAME: str = DEFAULT_PATHS["flan"]


# ------------------------------------------------------------------------------
# 3.1. FLAN-BASED (google/flan-t5-xl) SUMMARIZER
# ------------------------------------------------------------------------------

class FlanManager:
    def __init__(self, model_name: Optional[str] = None):
        # FIX: was `model or FLAN_MODEL_NAME` — `model` is undefined; param is `model_name`
        self.model_name = model_name or FLAN_MODEL_NAME
        self.flan_tokenizer = get_transformers('AutoTokenizer').from_pretrained(self.model_name)
        self.flan_model = get_transformers('AutoModelForSeq2SeqLM').from_pretrained(self.model_name)
        self.flan_device = 0 if get_torch().cuda.is_available() else -1
        self.flan_summarizer = get_transformers('pipeline')(
            "text2text-generation",
            model=self.flan_model,
            tokenizer=self.flan_tokenizer,
            device=self.flan_device,
        )


def get_flan_summarizer(model_name: Optional[str] = None):
    # FIX: consistent spelling — was get_flan_sumarizer (one 'm') in some places
    flan_mgr = FlanManager(model_name=model_name)
    return flan_mgr.flan_summarizer


def get_flan_summary(
    text: str,
    max_chunk: int = 512,
    max_length: Optional[int] = None,
    min_length: Optional[int] = None,
    do_sample: bool = False,
    model_name: Optional[str] = None,
) -> str:
    """
    Use google/flan-t5-xl to generate a human-like summary of an input text.

    Args:
        text (str): Input text to summarize.
        max_chunk (int): Maximum tokens for each chunk. Defaults to 512.
        max_length (int | None): Max tokens for generated summary. Defaults to 512.
        min_length (int | None): Min tokens for generated summary. Defaults to 100.
        do_sample (bool): Whether to sample. Defaults to False.
        model_name (str | None): Override model path/id.

    Returns:
        str: The generated summary.
    """
    max_length = max_length or 512
    min_length = min_length or 100
    prompt = f"Summarize the following text in a coherent, concise paragraph:\n\n{text}"

    # FIX: was get_flan_sumarizer / flan_summarizer — now consistent
    flan_summarizer = get_flan_summarizer(model_name=model_name)
    output = flan_summarizer(
        prompt,
        max_length=max_length,
        min_length=min_length,
        do_sample=do_sample,
    )[0]["generated_text"]
    return output.strip()


# ------------------------------------------------------------------------------
# 3.2. T5 MANAGER
# ------------------------------------------------------------------------------

class T5Manager:
    def __init__(self, model_directory: Optional[str] = None):
        self.model_directory = model_directory or SUMMARIZER_DIR
        self.t5_tokenizer = get_transformers('AutoTokenizer').from_pretrained(self.model_directory)
        self.t5_model = get_transformers('AutoModelForSeq2SeqLM').from_pretrained(self.model_directory)


def get_t5_manager(model_directory: Optional[str] = None) -> T5Manager:
    return T5Manager(model_directory=model_directory)


def get_t5_tokenizer(model_directory: Optional[str] = None):
    # FIX: was get_t5__tokenizer (double underscore) — removed extra underscore
    return get_t5_manager(model_directory=model_directory).t5_tokenizer


def get_t5_model(model_directory: Optional[str] = None):
    return get_t5_manager(model_directory=model_directory).t5_model


# ------------------------------------------------------------------------------
# 3.3. CHUNK-BASED T5 SUMMARIZER (for arbitrarily long text)
# ------------------------------------------------------------------------------

def split_to_chunk(full_text: str, max_words: int = 300) -> List[str]:
    """
    Break a long text into smaller chunks (approximately max_words per chunk),
    splitting on sentence-like boundaries (". ").
    """
    sentences = full_text.split(". ")
    chunks = []
    buffer = ""
    for sent in sentences:
        candidate = (buffer + sent).strip()
        if len(candidate.split()) <= max_words:
            buffer = candidate + ". "
        else:
            if buffer:
                chunks.append(buffer.strip())
            buffer = sent + ". "
    if buffer:
        chunks.append(buffer.strip())
    return chunks


def chunk_summaries(
    chunks: List[str],
    max_length: int = 160,
    min_length: int = 40,
    truncation: bool = False,
    model_directory: Optional[str] = None,
) -> List[str]:

    summaries = []

    t5_tokenizer = get_t5_tokenizer(model_directory=model_directory)
    t5_model = get_t5_model(model_directory=model_directory)

    for chunk in chunks:

        inputs = t5_tokenizer(
            chunk,
            return_tensors="pt",
            truncation=truncation,
            max_length=512,
        )

        token_ids = inputs["input_ids"]

        with get_torch().no_grad():

            outputs = t5_model.generate(
                token_ids,
                max_length=max_length,
                min_length=min_length,
                num_beams=4,
                early_stopping=True,
            )

        summary_text = t5_tokenizer.decode(
            outputs[0],
            skip_special_tokens=True
        )

        summaries.append(summary_text.strip())

    return summaries

def get_summary(
    full_text: str,
    max_words: int = 300,
    max_length: int = 160,
    min_length: int = 40,
    truncation: bool = False,
) -> str:
    """
    Summarize arbitrarily long text by chunking, summarizing each chunk, then stitching.
    """
    if not full_text:
        return ""

    chunks = split_to_chunk(full_text, max_words=max_words)
    summaries = chunk_summaries(
        chunks,
        max_length=max_length,
        min_length=min_length,
        truncation=truncation,
    )
    return " ".join(summaries).strip()


# ------------------------------------------------------------------------------
# 3.4. FALCONSAI SUMMARIZER
# ------------------------------------------------------------------------------

class FalconsManager:
    def __init__(self):
        self.falcons_summarizer = get_transformers('pipeline')(
            "summarization",
            model="Falconsai/text_summarization",
            device=0 if get_torch().cuda.is_available() else -1,
        )


def get_falcons_summarizer():
    # FIX: was get_falcons_sumarizer (one 'm') — consistent spelling throughout
    falcons_mgr = FalconsManager()
    return falcons_mgr.falcons_summarizer


def chunk_falcons_summaries(
    chunks: List[str],
    max_length: int = 160,
    min_length: int = 40,
    truncation: bool = True,
) -> List[str]:
    """
    Summarize each chunk using Falconsai/text_summarization.
    """
    summaries = []
    falcons_summarizer = get_falcons_summarizer()
    for chunk in chunks:
        out = falcons_summarizer(
            chunk,
            max_length=max_length,
            min_length=min_length,
            truncation=truncation,
        )
        summaries.append(out[0]["summary_text"].strip())
    return summaries
