"""Main TextSummarizer class"""
import os
from dataclasses import dataclass
from typing import Optional

try:
    from openai import AzureOpenAI, OpenAI
except ImportError:
    raise ImportError("Install openai: pip install openai>=1.0.0")

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

@dataclass
class SummaryResult:
    """Result from summarize()"""
    summary: str
    input_tokens: int
    output_tokens: int
    num_chunks: int = 1

    @property
    def compression(self) -> str:
        if self.input_tokens == 0:
            return "N/A"
        saved = self.input_tokens - self.output_tokens
        pct = (saved / self.input_tokens) * 100
        return f"{pct:.0f}% shorter" if saved > 0 else "summary longer!"

    def __str__(self) -> str:
        return (f"{self.summary}\n\n"
                f"[Compression: {self.compression} | "
                f"Input: {self.input_tokens} tokens | "
                f"Output: {self.output_tokens} tokens | "
                f"Chunks: {self.num_chunks}]")

class TextSummarizer:
    """Production-ready text summarizer with chunking and token budgets."""

    def __init__(self, api_key=None, azure_endpoint=None, deployment_name=None,
                 api_version="2024-02-15-preview", base_url=None, model="gpt-4o-mini", provider="auto"):
        self.api_key = api_key or os.getenv("OPENAI_API_KEY") or os.getenv("AZURE_OPENAI_API_KEY")
        self.azure_endpoint = azure_endpoint or os.getenv("AZURE_OPENAI_ENDPOINT")
        self.deployment_name = deployment_name or os.getenv("AZURE_OPENAI_DEPLOYMENT")
        self.api_version = api_version
        self.base_url = base_url
        self.model = model

        if provider == "auto":
            provider = "azure" if (azure_endpoint or self.azure_endpoint) else ("custom" if base_url else "openai")

        self.provider = provider

        if self.provider == "azure":
            if not self.azure_endpoint or not self.api_key:
                raise ValueError("Azure requires azure_endpoint and api_key")
            self.client = AzureOpenAI(api_key=self.api_key, api_version=self.api_version,
                                     azure_endpoint=self.azure_endpoint)
            self.model = self.deployment_name or self.model
        elif self.provider == "custom":
            if not self.base_url or not self.api_key:
                raise ValueError("Custom requires base_url and api_key")
            self.client = OpenAI(api_key=self.api_key, base_url=self.base_url)
        else:
            if not self.api_key:
                raise ValueError("OpenAI requires api_key")
            self.client = OpenAI(api_key=self.api_key)

        logger.info(f"TextSummarizer initialized: {self.provider}")

    def summarize(self, text: str, summary_type="medium", tone="neutral", focus_area="general",
                  output_format="text", chunk_strategy="character", chunk_size=3000) -> SummaryResult:
        """Summarize text with token budget guardrails."""
        should_continue = validate_input(text)
        input_tokens = count_tokens(text)

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

        summary_length = length_mapper(summary_type)
        format_instruction = format_mapper(output_format)
        max_tokens = token_budget(summary_type, input_tokens)

        logger.info(f"Summarizing: {input_tokens} input tokens, budget={max_tokens}")

        chunks = (chunk_text_by_sentences(text, max_chunk_size=chunk_size)
                 if chunk_strategy == "sentence"
                 else chunk_text(text, chunk_size=chunk_size))
        logger.info(f"Split into {len(chunks)} 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(c, summary_length, tone,
                                focus_area, format_instruction, max_tokens)
                               for c in chunks]
            combined = "\n\n".join(partial_summaries)
            logger.info("Re-summarizing combined chunks...")
            summary_text = self._summarize_chunk(combined, summary_length, tone,
                                                focus_area, format_instruction, max_tokens)

        output_tokens = count_tokens(summary_text)
        logger.info(f"Done: {input_tokens} → {output_tokens} tokens "
                   f"({((input_tokens-output_tokens)/input_tokens*100):.1f}%)")

        return SummaryResult(summary=summary_text, input_tokens=input_tokens,
                           output_tokens=output_tokens, num_chunks=len(chunks))

    def _summarize_chunk(self, text, summary_length, tone, focus_area, format_instruction, max_tokens):
    # For Azure: use SIMPLE, direct prompt
        if self.provider == "azure":
            simple_prompt = f"Summarize this in a {tone} tone:\n\n{text}"
        else:
            # For OpenAI: use complex prompt
            simple_prompt = build_user_prompt(text, summary_length, tone, focus_area,
                                            format_instruction, max_tokens)
        
        # Messages - NO system prompt for Azure
        if self.provider == "azure":
            messages = [{"role": "user", "content": simple_prompt}]
        else:
            messages = [
                {"role": "system", "content": SYSTEM_PROMPT},
                {"role": "user", "content": simple_prompt}
            ]
        
        # API kwargs
        api_kwargs = {
            "model": self.model,
            "messages": messages,
        }
        
        # Token parameters
        if self.provider == "azure":
            api_kwargs["max_completion_tokens"] = max_tokens
        else:
            api_kwargs["max_tokens"] = max_tokens
            api_kwargs["temperature"] = 0.1
        
        response = self.client.chat.completions.create(**api_kwargs)
        content = response.choices[0].message.content.strip()
        return content
