import logging
import time
from abc import ABC, abstractmethod
from typing import Any, Dict, List, Optional, Union, cast

from pydantic import BaseModel
from tqdm.auto import tqdm

from rhesis.sdk.entities.test_set import TestSet
from rhesis.sdk.enums import TestType
from rhesis.sdk.models import get_model
from rhesis.sdk.models.base import BaseLLM
from rhesis.sdk.services.chunker import (
    ChunkingService,
    ChunkingStrategy,
    SemanticChunker,
)
from rhesis.sdk.services.extractor import (
    ExtractionService,
    SourceSpecification,
)
from rhesis.sdk.synthesizers.utils import (
    create_test_set,
    load_prompt_template,
)

logger = logging.getLogger(__name__)


class Prompt(BaseModel):
    content: str
    expected_response: str
    language_code: str


class Test(BaseModel):
    prompt: Prompt
    behavior: str
    category: str
    topic: str
    # Note: test_type is NOT included in the schema sent to the LLM
    # It will be added programmatically after generation


class Tests(BaseModel):
    tests: List[Test]


# Flat schema for LLM batch generation (easier for the model to produce).
# Repacked to nested Test structure after generation.
class FlatTest(BaseModel):
    prompt_content: str
    prompt_expected_response: str
    prompt_language_code: str
    behavior: str
    category: str
    topic: str


class FlatTests(BaseModel):
    tests: List[FlatTest]


class TestSetSynthesizer(ABC):
    """Base class for all test set synthesizers."""

    prompt_template_file: str

    def __init__(
        self,
        batch_size: int = 5,
        model: Optional[Union[str, BaseLLM]] = None,
        sources: Optional[List[SourceSpecification]] = None,
        chunking_strategy: Optional[ChunkingStrategy] = SemanticChunker(max_tokens_per_chunk=1500),
    ):
        """
        Initialize the base synthesizer.

        Args:
            batch_size: Maximum number of items to process in a single LLM call
            model: The model to use for generation (string name or BaseLLM instance)
            sources: Optional list of source specifications to extract content from
            chunking_strategy: Strategy for chunking source content
        """
        self.batch_size = batch_size
        self.prompt_template = load_prompt_template(self.prompt_template_file)
        self.sources = sources
        self.chunker = chunking_strategy

        if isinstance(model, str) or model is None:
            self.model = get_model(model)
        else:
            self.model = model

    def _process_with_progress(
        self,
        items: List[Any],
        process_func: Any,
        desc: str = "Processing",
    ) -> List[Any]:
        """Process items with a progress bar."""
        results = []
        with tqdm(total=len(items), desc=desc) as pbar:
            for item in items:
                result = process_func(item)
                if isinstance(result, list):
                    results.extend(result)
                else:
                    results.append(result)
                pbar.update(1)
        return results

    @abstractmethod
    def _get_template_context(self, **generate_kwargs: Any) -> Dict[str, Any]:
        """
        Prepare template context for _generate_batch() call.

        Subclasses should combine instance attributes (from __init__)
        with runtime parameters (from generate()) to build the template context.

        Args:
            **generate_kwargs: Runtime parameters passed to generate()

        Returns:
            Dict containing template context to pass to _generate_batch()
        """
        pass

    def _get_synthesizer_name(self) -> str:
        """
        Return the name of the synthesizer for metadata.

        By default, returns the class name. Subclasses can override
        if they need a custom name.

        Returns:
            str: The synthesizer name (e.g., "PromptSynthesizer", "ConfigSynthesizer")
        """
        return self.__class__.__name__

    def _compute_tests_per_chunk(self, num_tests: int, num_chunks: int) -> list[int]:
        tests_per_chunk = [(num_tests + i) // num_chunks for i in range(num_chunks)]
        tests_per_chunk.reverse()
        return tests_per_chunk

    def _generate_with_sources(
        self, num_tests: int, **kwargs: Any
    ) -> tuple[List[Dict[str, Any]], dict[str, Any]]:
        # Process documents with source tracking
        if not isinstance(self.sources, list) or not all(
            isinstance(source, SourceSpecification) for source in self.sources
        ):
            raise ValueError("sources must be a list of SourceBase objects")

        if self.chunker is None or not isinstance(self.chunker, ChunkingStrategy):
            raise ValueError("chunker must be a ChunkingStrategy object")

        logger.info(
            "[Synthesizer] _generate_with_sources: extracting %d sources",
            len(self.sources),
        )
        extract_start = time.time()
        processed_sources = ExtractionService.extract(self.sources)
        logger.info(
            "[Synthesizer] Extraction completed in %.1fs",
            time.time() - extract_start,
        )

        chunk_start = time.time()
        chunks = ChunkingService(processed_sources, strategy=self.chunker).chunk()
        logger.info(
            "[Synthesizer] Chunking completed in %.1fs: %d chunks from %d sources",
            time.time() - chunk_start,
            len(chunks),
            len(processed_sources),
        )

        tests_per_chunk = self._compute_tests_per_chunk(num_tests, len(chunks))
        if num_tests < len(chunks):
            logger.warning(
                "[Synthesizer] num_tests (%d) < num_chunks (%d), some chunks will be skipped",
                num_tests,
                len(chunks),
            )
        else:
            logger.info(
                "[Synthesizer] Generating %d tests across %d chunks",
                num_tests,
                len(chunks),
            )

        if num_tests >= len(chunks):
            coverage_percent = 100
            used_chunks = len(chunks)
        else:
            coverage_percent = num_tests / len(chunks)
            used_chunks = num_tests

        all_test_cases = []

        # Generate tests for each chunk
        for i, chunk in enumerate(chunks):
            if tests_per_chunk[i] == 0:
                continue
            logger.info(
                "[Synthesizer] Chunk %d/%d: generating %d tests (%d chars content, source=%s)",
                i + 1,
                min(num_tests, len(chunks)),
                tests_per_chunk[i],
                len(chunk.content),
                chunk.source.name,
            )

            result = self._generate_without_sources(
                num_tests=tests_per_chunk[i],
                **kwargs,
                source=chunk.content,
            )

            # Add context and document mapping to each test
            for test in result:
                # Ensure test_type is set (should already be set by _generate_batch)
                if "test_type" not in test:
                    test["test_type"] = TestType.SINGLE_TURN.value

                test["metadata"] = {
                    **(test.get("metadata") or {}),
                    "sources": [
                        {
                            "source": chunk.source.name,
                            "name": chunk.source.name,
                            "description": chunk.source.description,
                            "content": chunk.content,
                        }
                    ],
                    "generated_by": self._get_synthesizer_name(),
                    "context_index": i,
                    "context_length": len(chunk.content),
                }

            all_test_cases.extend(result)

        # Get document names for TestSet metadata

        source_names = [chunk.source.name for chunk in chunks]
        test_set_metadata = {
            "documents_used": source_names,
            "coverage_percent": coverage_percent,
            "contexts_total": len(chunks),
            "contexts_used": used_chunks,
            "tests_per_context": tests_per_chunk,
        }
        return all_test_cases, test_set_metadata

    def _generate_without_sources(self, num_tests: int = 5, **kwargs: Any) -> List[Dict[str, Any]]:
        """
        Generate test cases with automatic chunking.

        Args:
            num_tests: Total number of test cases to generate. Defaults to 5.
            **kwargs: Additional keyword arguments for test set generation

        Returns:
            TestSet: A TestSet entity containing the generated test cases
        """

        if not isinstance(num_tests, int):
            raise TypeError("num_tests must be an integer")

        template_context = self._get_template_context(**kwargs)

        all_test_cases = []
        generation_start = time.time()
        logger.info(
            "[Synthesizer] Starting generation: num_tests=%d, batch_size=%d, model=%s",
            num_tests,
            self.batch_size,
            getattr(self.model, "model_name", type(self.model).__name__),
        )

        # For large numbers, use chunking to avoid JSON parsing issues
        max_consecutive_failures = 3
        if num_tests > self.batch_size:
            # Generate in chunks
            remaining_tests = num_tests
            batch_num = 0
            consecutive_failures = 0
            while remaining_tests > 0:
                batch_num += 1
                chunk_size = min(self.batch_size, remaining_tests)
                logger.info(
                    "[Synthesizer] Batch %d: requesting %d tests "
                    "(remaining=%d, total_so_far=%d, batch_size=%d)",
                    batch_num,
                    chunk_size,
                    remaining_tests,
                    len(all_test_cases),
                    self.batch_size,
                )
                batch_start = time.time()
                try:
                    chunk_tests = self._generate_batch(chunk_size, **template_context)
                    batch_elapsed = time.time() - batch_start
                    logger.info(
                        "[Synthesizer] Batch %d: received %d/%d tests in %.1fs",
                        batch_num,
                        len(chunk_tests),
                        chunk_size,
                        batch_elapsed,
                    )

                    if len(chunk_tests) == 0:
                        consecutive_failures += 1
                        logger.warning(
                            "[Synthesizer] Batch %d: received 0 tests (%d/%d consecutive failures)",
                            batch_num,
                            consecutive_failures,
                            max_consecutive_failures,
                        )
                        if consecutive_failures >= max_consecutive_failures:
                            logger.error(
                                "[Synthesizer] Aborting after %d consecutive "
                                "empty batches (%d tests generated so far)",
                                max_consecutive_failures,
                                len(all_test_cases),
                            )
                            break
                        if self.batch_size > 5:
                            old_batch_size = self.batch_size
                            self.batch_size = max(5, self.batch_size // 2)
                            logger.warning(
                                "[Synthesizer] Reducing batch_size %d -> %d",
                                old_batch_size,
                                self.batch_size,
                            )
                        continue

                    consecutive_failures = 0
                    all_test_cases.extend(chunk_tests)
                    remaining_tests -= len(chunk_tests)

                    # If short, reduce batch size for next attempt
                    if len(chunk_tests) < chunk_size and self.batch_size > 5:
                        old_batch_size = self.batch_size
                        self.batch_size = max(5, self.batch_size // 2)
                        logger.warning(
                            "[Synthesizer] Batch %d: short by %d tests, "
                            "reducing batch_size %d -> %d",
                            batch_num,
                            chunk_size - len(chunk_tests),
                            old_batch_size,
                            self.batch_size,
                        )

                except Exception as e:
                    batch_elapsed = time.time() - batch_start
                    consecutive_failures += 1
                    logger.error(
                        "[Synthesizer] Batch %d FAILED after %.1fs: %s: %s",
                        batch_num,
                        batch_elapsed,
                        type(e).__name__,
                        e,
                        exc_info=True,
                    )
                    if consecutive_failures >= max_consecutive_failures:
                        logger.error(
                            "[Synthesizer] Aborting after %d consecutive "
                            "failures (%d tests generated so far)",
                            max_consecutive_failures,
                            len(all_test_cases),
                        )
                        break
                    if self.batch_size > 5:
                        old_batch_size = self.batch_size
                        self.batch_size = max(5, self.batch_size // 2)
                        logger.warning(
                            "[Synthesizer] Reducing batch_size %d -> %d after error",
                            old_batch_size,
                            self.batch_size,
                        )
                        continue
                    else:
                        logger.error(
                            "[Synthesizer] batch_size at minimum (%d), "
                            "aborting with %d tests generated",
                            self.batch_size,
                            len(all_test_cases),
                        )
                        break
        else:
            # Generate all tests in a single batch
            logger.info(
                "[Synthesizer] Single batch: requesting %d tests",
                num_tests,
            )
            batch_start = time.time()
            all_test_cases = self._generate_batch(num_tests, **template_context)
            batch_elapsed = time.time() - batch_start
            logger.info(
                "[Synthesizer] Single batch: received %d tests in %.1fs",
                len(all_test_cases),
                batch_elapsed,
            )

        total_elapsed = time.time() - generation_start
        logger.info(
            "[Synthesizer] Generation complete: %d/%d tests in %.1fs (%.1f tests/sec)",
            len(all_test_cases),
            num_tests,
            total_elapsed,
            len(all_test_cases) / total_elapsed if total_elapsed > 0 else 0,
        )

        # Ensure we have some test cases
        if not all_test_cases:
            raise ValueError("Failed to generate any valid test cases")

        return all_test_cases

    def _flat_test_to_nested(self, flat: Dict[str, Any]) -> Dict[str, Any]:
        """Repack a flat test dict (LLM output) into the nested Test structure."""
        return {
            "prompt": {
                "content": flat["prompt_content"],
                "expected_response": flat["prompt_expected_response"],
                "language_code": flat["prompt_language_code"],
            },
            "behavior": flat["behavior"],
            "category": flat["category"],
            "topic": flat["topic"],
        }

    def _generate_batch(
        self,
        num_tests: int,
        **kwargs: Any,
    ) -> List[Dict[str, Any]]:
        """Generate a batch of test cases with improved error handling."""
        template_context = {"num_tests": num_tests, **kwargs}
        prompt = self.prompt_template.render(**template_context)

        logger.debug(
            "[Synthesizer] _generate_batch: calling model.generate for %d tests "
            "(prompt length=%d chars)",
            num_tests,
            len(prompt),
        )

        # Use flat schema for LLM (easier to generate), then repack to nested
        llm_start = time.time()
        response = cast(
            Dict[str, Any],
            self.model.generate(prompt=prompt, schema=FlatTests),
        )
        llm_elapsed = time.time() - llm_start

        # Check for error responses from the LLM — return empty list so
        # the batch loop can retry or abort via consecutive-failure guard.
        if isinstance(response, dict) and "error" in response:
            logger.error(
                "[Synthesizer] _generate_batch: LLM returned error after %.1fs: %s",
                llm_elapsed,
                response["error"],
            )
            return []

        if not isinstance(response, dict) or "tests" not in response:
            logger.error(
                "[Synthesizer] _generate_batch: unexpected response type=%s after %.1fs: %s",
                type(response).__name__,
                llm_elapsed,
                str(response)[:500],
            )
            return []

        flat_tests = response["tests"][:num_tests]
        logger.debug(
            "[Synthesizer] _generate_batch: LLM returned %d tests in %.1fs",
            len(flat_tests),
            llm_elapsed,
        )

        tests = [
            {
                **self._flat_test_to_nested(flat),
                "test_type": TestType.SINGLE_TURN.value,  # Set to Single-Turn
                "metadata": {
                    "generated_by": self._get_synthesizer_name(),
                },
            }
            for flat in flat_tests
        ]

        return tests

    def generate(self, num_tests: int = 5, **kwargs: Any) -> TestSet:
        """Generate test cases."""
        logger.info(
            "[Synthesizer] generate() called: num_tests=%d, synthesizer=%s, "
            "has_sources=%s, model=%s",
            num_tests,
            self._get_synthesizer_name(),
            self.sources is not None,
            getattr(self.model, "model_name", type(self.model).__name__),
        )
        overall_start = time.time()

        test_set_metadata = {}
        if self.sources is not None:
            tests, test_set_metadata = self._generate_with_sources(num_tests, **kwargs)
        else:
            tests = self._generate_without_sources(num_tests, **kwargs)

        logger.info(
            "[Synthesizer] Test generation phase complete: %d tests in %.1fs, creating TestSet...",
            len(tests),
            time.time() - overall_start,
        )

        # Use utility function to create TestSet
        test_set = create_test_set(
            tests,
            model=self.model,
            synthesizer_name=self._get_synthesizer_name(),
            batch_size=self.batch_size,
            num_tests=len(tests),
            requested_tests=num_tests,
            **test_set_metadata,
        )

        total_elapsed = time.time() - overall_start
        logger.info(
            "[Synthesizer] generate() complete: %d tests, total time %.1fs",
            len(tests),
            total_elapsed,
        )
        return test_set
