import os
from dataclasses import dataclass
from openai import AzureOpenAI

from .chunking import chunk_text
from .prompts import SYSTEM_PROMPT, build_user_prompt
from .utils import validate_input, length_mapper, max_tokens_mapper, format_mapper, logger
from .token_counter import count_tokens


@dataclass
class SummaryResult:
    """
    Result returned by :meth:`AzureSummarizer.summarize`.

    Attributes:
        summary (str):
            The generated summary text.
        input_tokens (int):
            Estimated number of tokens in the original input text.
        output_tokens (int):
            Estimated number of tokens in the generated summary.

    Example::

        result = summarizer.summarize("Long article...")
        print(result.summary)
        print(f"Input : {result.input_tokens} tokens")
        print(f"Output: {result.output_tokens} tokens")
    """

    summary: str
    input_tokens: int
    output_tokens: int

    def __str__(self) -> str:
        return (
            f"{self.summary}\n\n"
            f"[Tokens — Input: {self.input_tokens} | Summary: {self.output_tokens}]"
        )


class AzureSummarizer:
    """
    Azure OpenAI-backed text summarizer.

    Args:
        azure_endpoint (str, optional):
            Azure OpenAI endpoint URL. Falls back to the
            ``AZURE_OPENAI_ENDPOINT`` environment variable.
        api_key (str, optional):
            Azure OpenAI API key. Falls back to the
            ``AZURE_OPENAI_API_KEY`` environment variable.
        api_version (str):
            Azure OpenAI API version string.
        deployment_name (str):
            Name of the deployed model to use.

    Example::

        from llm_summarizer import AzureSummarizer

        summarizer = AzureSummarizer(
            api_key="<key>",
            azure_endpoint="https://<resource>.openai.azure.com",
        )
        result = summarizer.summarize(
            text="Your long document here...",
            summary_type="medium",   # "short" | "medium" | "detailed"
            tone="neutral",
            focus_area="general",
            output_format="text",    # "text" | "bullets" | "json"
        )
        print(result)                    # summary + token counts
        print(result.summary)            # summary text only
        print(result.input_tokens)       # token count of original text
        print(result.output_tokens)      # token count of summary
    """

    def __init__(
        self,
        azure_endpoint: str = None,
        api_key: str = None,
        api_version: str = "2024-02-15-preview",
        deployment_name: str = "gi-local-gpt-5-mini",
    ):
        self.azure_endpoint = azure_endpoint or os.getenv("AZURE_OPENAI_ENDPOINT")
        self.api_key = api_key or os.getenv("AZURE_OPENAI_API_KEY")
        self.deployment_name = deployment_name

        if not self.azure_endpoint or not self.api_key:
            raise ValueError("Azure endpoint and API key must be provided")

        self.client = AzureOpenAI(
            api_key=self.api_key,
            api_version=api_version,
            azure_endpoint=self.azure_endpoint,
        )

    def summarize(
        self,
        text: str,
        summary_type: str = "medium",
        tone: str = "neutral",
        focus_area: str = "general",
        output_format: str = "text",
    ) -> SummaryResult:
        """
        Summarize *text* and return a :class:`SummaryResult` containing
        the summary and token counts for both the input and the output.

        Args:
            text (str):
                The passage to summarise.
            summary_type (str):
                Desired length -- ``"short"``, ``"medium"``, or
                ``"detailed"``. Defaults to ``"medium"``.
            tone (str):
                Writing tone, e.g. ``"neutral"``, ``"formal"``,
                ``"casual"``. Defaults to ``"neutral"``.
            focus_area (str):
                Topic lens, e.g. ``"general"``, ``"technical insights"``,
                ``"financial"``. Defaults to ``"general"``.
            output_format (str):
                Output structure -- ``"text"``, ``"bullets"``, or
                ``"json"``. Defaults to ``"text"``.

        Returns:
            :class:`SummaryResult`:
                Object with ``summary``, ``input_tokens``, and
                ``output_tokens`` attributes.
        """
        should_continue = validate_input(text)
        input_tokens = count_tokens(text)

        if not should_continue:
            return SummaryResult(
                summary=text,
                input_tokens=input_tokens,
                output_tokens=count_tokens(text),
            )

        summary_length = length_mapper(summary_type)
        max_tokens = max_tokens_mapper(summary_type)
        format_instruction = format_mapper(output_format)
        chunks = chunk_text(text)

        logger.info(f"Total chunks: {len(chunks)}")

        if len(chunks) == 1:
            summary_text = self._summarize_chunk(
                chunks[0], summary_length, tone, focus_area, format_instruction, max_tokens
            )
        else:
            partial_summaries = [
                self._summarize_chunk(
                    chunk, summary_length, tone, focus_area, format_instruction, max_tokens
                )
                for chunk in chunks
            ]
            combined = " ".join(partial_summaries)
            logger.info("Running reduce summarization")
            summary_text = self._summarize_chunk(
                combined, summary_length, tone, focus_area, format_instruction, max_tokens
            )

        output_tokens = count_tokens(summary_text)
        logger.info(f"Token counts -- input: {input_tokens}, summary: {output_tokens}")

        return SummaryResult(
            summary=summary_text,
            input_tokens=input_tokens,
            output_tokens=output_tokens,
        )

    def _summarize_chunk(
        self,
        text: str,
        summary_length: str,
        tone: str,
        focus_area: str,
        format_instruction: str,
        max_tokens: int = 300,
    ) -> str:
        user_prompt = build_user_prompt(
            text, summary_length, tone, focus_area, format_instruction
        )
        response = self.client.chat.completions.create(
            model=self.deployment_name,
            max_completion_tokens=max_tokens,
            messages=[
                {"role": "system", "content": SYSTEM_PROMPT},
                {"role": "user", "content": user_prompt},
            ],
        )
        return response.choices[0].message.content.strip()
