"""LLM data sources from 2024-2025 models (GPT-4, Claude 3, Llama 3, etc.)."""

from __future__ import annotations

import hashlib
import logging
from typing import TYPE_CHECKING

from llm_detector.types import TextSample
from .base import BaseDataSource, BatchConfig

if TYPE_CHECKING:
    from collections.abc import Iterator

logger = logging.getLogger(__name__)


class GPT4AlpacaSource(BaseDataSource):
    """GPT-4 generated instruction-following dataset.

    52K English instruction-following examples generated by GPT-4
    using the same prompts as Stanford Alpaca.
    """

    def __init__(
        self,
        *,
        batch_config: BatchConfig | None = None,
        min_length: int = 10,
        max_length: int | None = None,
    ) -> None:
        super().__init__(config=batch_config)
        self.min_length = min_length
        self.max_length = max_length
        self._dataset_name = "vicgalle/alpaca-gpt4"

    @property
    def name(self) -> str:
        return "gpt4_alpaca"

    @property
    def category(self) -> str:
        return "llm"

    def _generate_samples(self, limit: int | None = None) -> Iterator[TextSample]:
        """Stream GPT-4 generated instruction responses."""
        try:
            from datasets import load_dataset
        except ImportError as e:
            raise RuntimeError(f"{self.name}: datasets library required") from e

        try:
            dataset = load_dataset(
                self._dataset_name,
                split="train",
                streaming=True,
            )
        except Exception as e:
            logger.error(f"Failed to load {self._dataset_name}: {e}")
            raise RuntimeError(f"Could not load dataset {self._dataset_name}") from e

        count = 0
        for idx, item in enumerate(dataset):
            if limit is not None and count >= limit:
                break

            # Extract the GPT-4 generated output
            text = item.get("output", "").strip()
            if not text:
                continue

            # Apply length filters
            if len(text) < self.min_length:
                continue
            if self.max_length and len(text) > self.max_length:
                continue

            # Clean the text
            cleaned = text.strip()
            if not cleaned or len(cleaned) < self.min_length:
                continue

            # Create sample
            sample_id = f"{self.name}_{idx:08d}"
            text_hash = hashlib.md5(cleaned.encode("utf-8")).hexdigest()[:8]

            yield TextSample(
                text=cleaned,
                source=self.name,
                is_llm=True,
                metadata={
                    "dataset": self._dataset_name,
                    "sample_id": sample_id,
                    "text_hash": text_hash,
                    "instruction": item.get("instruction", "")[:100],  # First 100 chars
                    "model": "gpt-4",
                },
            )
            count += 1


class Claude3OpusSource(BaseDataSource):
    """Claude 3 Opus instruction dataset.

    15K high-quality instruction-following examples from Claude 3 Opus.
    """

    def __init__(
        self,
        *,
        batch_config: BatchConfig | None = None,
        min_length: int = 10,
        max_length: int | None = None,
    ) -> None:
        super().__init__(config=batch_config)
        self.min_length = min_length
        self.max_length = max_length
        self._dataset_name = "nothingiisreal/Claude-3-Opus-Instruct-15K"

    @property
    def name(self) -> str:
        return "claude3_opus"

    @property
    def category(self) -> str:
        return "llm"

    def _generate_samples(self, limit: int | None = None) -> Iterator[TextSample]:
        """Stream Claude 3 Opus generated responses."""
        try:
            from datasets import load_dataset
        except ImportError as e:
            raise RuntimeError(f"{self.name}: datasets library required") from e

        try:
            dataset = load_dataset(
                self._dataset_name,
                "Instruct Data v1 - Merged",  # Need to specify config for this dataset
                split="train",
                streaming=True,
            )
        except Exception as e:
            logger.error(f"Failed to load {self._dataset_name}: {e}")
            raise RuntimeError(f"Could not load dataset {self._dataset_name}") from e

        count = 0
        for idx, item in enumerate(dataset):
            if limit is not None and count >= limit:
                break

            # Extract Claude's response - field names may vary
            text = item.get("response", item.get("output", "")).strip()
            if not text:
                continue

            # Apply filters
            if len(text) < self.min_length:
                continue
            if self.max_length and len(text) > self.max_length:
                continue

            cleaned = text.strip()
            if not cleaned or len(cleaned) < self.min_length:
                continue

            sample_id = f"{self.name}_{idx:08d}"
            text_hash = hashlib.md5(cleaned.encode("utf-8")).hexdigest()[:8]

            yield TextSample(
                text=cleaned,
                source=self.name,
                is_llm=True,
                metadata={
                    "dataset": self._dataset_name,
                    "sample_id": sample_id,
                    "text_hash": text_hash,
                    "model": "claude-3-opus",
                },
            )
            count += 1


class Claude35SonnetSource(BaseDataSource):
    """Claude 3.5 Sonnet reflection dataset.

    Detailed reflection examples from Claude 3.5 Sonnet,
    inspired by OpenAI's o1 reasoning approach.
    """

    def __init__(
        self,
        *,
        batch_config: BatchConfig | None = None,
        min_length: int = 10,
        max_length: int | None = None,
    ) -> None:
        super().__init__(config=batch_config)
        self.min_length = min_length
        self.max_length = max_length
        self._dataset_name = "QuietImpostor/Claude-3-Opus-Claude-3.5-Sonnnet-9k"

    @property
    def name(self) -> str:
        return "claude35_sonnet"

    @property
    def category(self) -> str:
        return "llm"

    def _generate_samples(self, limit: int | None = None) -> Iterator[TextSample]:
        """Stream Claude 3.5 Sonnet generated text."""
        try:
            from datasets import load_dataset
        except ImportError as e:
            raise RuntimeError(f"{self.name}: datasets library required") from e

        try:
            dataset = load_dataset(
                self._dataset_name,
                split="train",
                streaming=True,
            )
        except Exception as e:
            logger.error(f"Failed to load {self._dataset_name}: {e}")
            # Try alternative dataset
            alt_dataset = "leafspark/DetailedReflection-Claude-v3_5-Sonnet"
            logger.info(f"Trying alternative: {alt_dataset}")
            try:
                dataset = load_dataset(alt_dataset, split="train", streaming=True)
                self._dataset_name = alt_dataset
            except Exception as e2:
                raise RuntimeError(f"Could not load any Claude 3.5 dataset") from e2

        count = 0
        for idx, item in enumerate(dataset):
            if limit is not None and count >= limit:
                break

            # Try multiple possible field names
            text = (
                item.get("response", "") or
                item.get("output", "") or
                item.get("completion", "") or
                item.get("text", "")
            ).strip()

            if not text:
                continue

            if len(text) < self.min_length:
                continue
            if self.max_length and len(text) > self.max_length:
                continue

            cleaned = text.strip()
            if not cleaned or len(cleaned) < self.min_length:
                continue

            sample_id = f"{self.name}_{idx:08d}"
            text_hash = hashlib.md5(cleaned.encode("utf-8")).hexdigest()[:8]

            yield TextSample(
                text=cleaned,
                source=self.name,
                is_llm=True,
                metadata={
                    "dataset": self._dataset_name,
                    "sample_id": sample_id,
                    "text_hash": text_hash,
                    "model": "claude-3.5-sonnet",
                },
            )
            count += 1


class MixtralInstructSource(BaseDataSource):
    """Mixtral-8x7B instruction dataset.

    High-quality synthetic data from Mixtral models,
    which are widely used for open-source applications.
    """

    def __init__(
        self,
        *,
        batch_config: BatchConfig | None = None,
        min_length: int = 10,
        max_length: int | None = None,
    ) -> None:
        super().__init__(config=batch_config)
        self.min_length = min_length
        self.max_length = max_length
        # Cosmopedia subset or similar Mixtral-generated content
        self._dataset_name = "HuggingFaceTB/cosmopedia"

    @property
    def name(self) -> str:
        return "mixtral_instruct"

    @property
    def category(self) -> str:
        return "llm"

    def _generate_samples(self, limit: int | None = None) -> Iterator[TextSample]:
        """Stream Mixtral-generated synthetic content."""
        try:
            from datasets import load_dataset
        except ImportError as e:
            raise RuntimeError(f"{self.name}: datasets library required") from e

        try:
            # Load a subset of Cosmopedia (it's huge)
            dataset = load_dataset(
                self._dataset_name,
                "web_samples_v1",  # One of the smaller subsets
                split="train",
                streaming=True,
            )
        except Exception as e:
            logger.error(f"Failed to load {self._dataset_name}: {e}")
            raise RuntimeError(f"Could not load dataset {self._dataset_name}") from e

        count = 0
        for idx, item in enumerate(dataset):
            if limit is not None and count >= limit:
                break

            text = item.get("text", "").strip()
            if not text:
                continue

            # Take a reasonable chunk (Cosmopedia texts can be very long)
            text = text[:10000]  # Cap at 10k chars

            if len(text) < self.min_length:
                continue
            if self.max_length and len(text) > self.max_length:
                continue

            cleaned = text.strip()
            if not cleaned or len(cleaned) < self.min_length:
                continue

            sample_id = f"{self.name}_{idx:08d}"
            text_hash = hashlib.md5(cleaned.encode("utf-8")).hexdigest()[:8]

            yield TextSample(
                text=cleaned,
                source=self.name,
                is_llm=True,
                metadata={
                    "dataset": self._dataset_name,
                    "sample_id": sample_id,
                    "text_hash": text_hash,
                    "model": "mixtral-8x7b",
                    "prompt": item.get("prompt", "")[:100] if "prompt" in item else "",
                },
            )
            count += 1


# Registry helper
def register_2024_llm_sources(registry):
    """Register all 2024-2025 LLM sources."""
    from .registry import SourceDefinition

    sources = [
        SourceDefinition(
            name="gpt4_alpaca",
            factory=GPT4AlpacaSource,
            category="llm",
            description="GPT-4 instruction-following dataset (52K examples)",
            enabled=True,
        ),
        SourceDefinition(
            name="claude3_opus",
            factory=Claude3OpusSource,
            category="llm",
            description="Claude 3 Opus instruction dataset (15K examples)",
            enabled=True,
        ),
        SourceDefinition(
            name="claude35_sonnet",
            factory=Claude35SonnetSource,
            category="llm",
            description="Claude 3.5 Sonnet reflection dataset",
            enabled=True,
        ),
        SourceDefinition(
            name="mixtral_instruct",
            factory=MixtralInstructSource,
            category="llm",
            description="Mixtral-8x7B synthetic content",
            enabled=True,
        ),
    ]

    for source_def in sources:
        registry.register(source_def.name, source_def)


__all__ = [
    "GPT4AlpacaSource",
    "Claude3OpusSource",
    "Claude35SonnetSource",
    "MixtralInstructSource",
    "register_2024_llm_sources",
]