from __future__ import annotations

import io
from collections.abc import Sequence
from typing import Any

import pytest

from amsdal_ml.ml_ingesting import IngestionSource
from amsdal_ml.ml_ingesting import LoadedDocument
from amsdal_ml.ml_ingesting import LoadedPage
from amsdal_ml.ml_ingesting import TextChunk
from amsdal_ml.ml_ingesting.chunk_strategy import ChunkStrategy
from amsdal_ml.ml_ingesting.default_ingesting import DefaultIngesting
from amsdal_ml.ml_ingesting.embedders.embedder import Embedder
from amsdal_ml.ml_ingesting.embedders.openai_embedder import OpenAIEmbedder
from amsdal_ml.ml_ingesting.embedding_data import EmbeddingData
from amsdal_ml.ml_ingesting.loaders.loader import Loader
from amsdal_ml.ml_ingesting.loaders.pdf_loader import PdfLoader
from amsdal_ml.ml_ingesting.pipeline import DefaultIngestionPipeline
from amsdal_ml.ml_ingesting.processors.cleaner import Cleaner
from amsdal_ml.ml_ingesting.processors.text_cleaner import TextCleaner
from amsdal_ml.ml_ingesting.splitters.splitter import Splitter
from amsdal_ml.ml_ingesting.splitters.token_splitter import TokenSplitter
from amsdal_ml.ml_ingesting.stores.embedding_data import EmbeddingDataStore
from amsdal_ml.ml_ingesting.stores.store import EmbeddingStore

EXPECTED_PAGE_NUMBER = 2


class _FakePdfPage:
    def __init__(self, text: str) -> None:
        self._text = text

    def extract_text(self) -> str:
        return self._text


class _FakePdfDoc:
    def __init__(self):
        self.pages = [_FakePdfPage('Hello PDF'), _FakePdfPage('')]
        self.metadata = {'/Author': 'Tester'}


def test_pdf_loader_uses_reader(monkeypatch):
    monkeypatch.setattr('amsdal_ml.ml_ingesting.loaders.pdf_loader.PdfReader', lambda *_, **__: _FakePdfDoc())
    loader = PdfLoader()
    buf = io.BytesIO(b'%PDF-1.4 test')
    doc = loader.load(buf, filename='doc.pdf', metadata={'source': 'unit'})
    assert len(doc.pages) == 1
    assert doc.pages[0].page_number == 1
    assert doc.pages[0].text == 'Hello PDF'
    assert doc.metadata['Author'] == 'Tester'
    assert doc.metadata['source'] == 'unit'
    assert doc.metadata['filename'] == 'doc.pdf'


def test_text_cleaner_normalizes_and_drops_empty():
    doc = LoadedDocument(
        pages=[
            LoadedPage(page_number=1, text='Hello\tworld\r\n', metadata={'page_number': 1}),
            LoadedPage(page_number=2, text=' ', metadata={'page_number': 2}),
        ],
        metadata={'foo': 'bar'},
    )
    cleaner = TextCleaner()
    cleaned = cleaner.clean(doc)
    assert len(cleaned.pages) == 1
    assert cleaned.pages[0].text == 'Hello world'
    assert cleaned.metadata['foo'] == 'bar'


def test_token_splitter_preserves_metadata():
    doc = LoadedDocument(
        pages=[LoadedPage(page_number=1, text=' '.join(['w'] * 40), metadata={'page_number': 1})],
        metadata={'filename': 'file.pdf'},
    )
    splitter = TokenSplitter(max_tokens=10, overlap_tokens=5, token_len_fn=lambda _: 1)
    chunks = splitter.split(doc)
    assert chunks, 'chunks should not be empty'
    assert chunks[0].metadata['filename'] == 'file.pdf'
    assert chunks[0].metadata['page_number'] == 1
    assert chunks[0].index == 0
    assert all(ch.text for ch in chunks)


class _FakeLoader(Loader):
    def __init__(self, doc: LoadedDocument):
        self.doc = doc
        self.called = False

    def load(self, *_, **__):  # noqa: ANN001
        self.called = True
        return self.doc

    async def aload(self, *_, **__):  # noqa: ANN001
        return self.load(*_, **__)


class _FakeCleaner(Cleaner):
    def __init__(self, doc: LoadedDocument):
        self.doc = doc
        self.called = False

    def clean(self, *_):  # noqa: ANN001
        self.called = True
        return self.doc

    async def aclean(self, *_):  # noqa: ANN001
        return self.clean(*_)


class _FakeSplitter(Splitter):
    def __init__(self, chunks: list[TextChunk]):
        self.chunks = chunks
        self.called = False

    def split(self, *_):  # noqa: ANN001
        self.called = True
        return self.chunks

    async def asplit(self, *_):  # noqa: ANN001
        return self.split(*_)


class _FakeEmbedder(Embedder):
    def __init__(self):
        self.calls: list[str] = []

    def embed(self, text: str) -> list[float]:
        self.calls.append(text)
        return [1.0, 0.0, 0.0]

    async def aembed(self, text: str) -> list[float]:
        return self.embed(text)


class _FakeStore(EmbeddingStore):
    def __init__(self):
        self.saved = None

    def save(self, embeddings: Sequence[EmbeddingData], *, source: IngestionSource | None = None) -> list[Any]:
        self.saved = embeddings
        return list(embeddings)

    async def asave(self, embeddings: Sequence[EmbeddingData], *, source: IngestionSource | None = None) -> list[Any]:
        return self.save(embeddings, source=source)


def test_pipeline_assembles_and_saves():
    doc = LoadedDocument(pages=[LoadedPage(page_number=1, text='hello world', metadata={})], metadata={})
    chunks = [TextChunk(index=0, text='hello world', tags=['c'], metadata={'page_number': 1})]

    loader = _FakeLoader(doc)
    cleaner = _FakeCleaner(doc)
    splitter = _FakeSplitter(chunks)
    embedder = _FakeEmbedder()
    store = _FakeStore()

    pipeline = DefaultIngestionPipeline(
        loader=loader,
        cleaner=cleaner,
        splitter=splitter,
        embedder=embedder,
        store=store,
    )

    source = IngestionSource(object_class='Doc', object_id='1', tags=['a'], metadata={'filename': 'f.pdf'})
    result = pipeline.run(io.BytesIO(b'data'), filename='f.pdf', tags=['b'], source=source)

    assert result == store.saved
    assert loader.called and cleaner.called and splitter.called
    assert embedder.calls == ['hello world']
    assert store.saved[0].tags == ['a', 'b', 'c']
    assert store.saved[0].metadata['filename'] == 'f.pdf'
    assert store.saved[0].metadata['page_number'] == 1


class _DummyManager:
    def __init__(self):
        self.saved = None

    def bulk_create(self, objs):  # noqa: ANN001
        self.saved = objs
        return objs


class _DummyModel:
    objects = _DummyManager()

    def __init__(self, **kwargs):
        for k, v in kwargs.items():
            setattr(self, k, v)


def test_embedding_store_merges_tags_and_metadata():
    store = EmbeddingDataStore(model_cls=_DummyModel)
    embeddings = [
        EmbeddingData(
            chunk_index=0,
            raw_text='t',
            embedding=[0.1],
            tags=['x'],
            metadata={'page_number': 2},
        )
    ]
    source = IngestionSource(object_class='Doc', object_id='42', tags=['a'], metadata={'filename': 'f.pdf'})

    saved = store.save(embeddings, source=source)
    assert len(saved) == 1
    obj = saved[0]
    assert obj.data_object_class == 'Doc'
    assert obj.data_object_id == '42'
    assert obj.tags == ['a', 'x']
    assert obj.ml_metadata['filename'] == 'f.pdf'
    assert obj.ml_metadata['page_number'] == EXPECTED_PAGE_NUMBER


@pytest.mark.skip(reason='integration: requires real OpenAI key')
def test_openai_embedder_smoke():
    embedder = OpenAIEmbedder()
    vec = embedder.embed('ping')
    assert isinstance(vec, list)
    assert vec, 'embedding should not be empty'


class _CustomChunkStrategy(ChunkStrategy):
    """Custom chunk strategy that varies params based on model type."""

    AUTHOR_MAX_DEPTH = 1
    AUTHOR_MAX_CHUNKS = 3
    AUTHOR_MAX_TOKENS = 500
    VEHICLE_MAX_DEPTH = 2
    VEHICLE_MAX_CHUNKS = 8
    VEHICLE_MAX_TOKENS = 700
    BOOK_MAX_DEPTH = 4
    BOOK_MAX_CHUNKS = 20
    BOOK_MAX_TOKENS = 1000
    DEFAULT_MAX_DEPTH = 2
    DEFAULT_MAX_CHUNKS = 10
    DEFAULT_MAX_TOKENS = 800
    EMBEDDING_DIM = 10

    def get_params(self, obj):
        from amsdal_ml.ml_ingesting.chunk_strategy import ChunkParams
        from tests.fixtures.models.author import Author
        from tests.fixtures.models.book import Book
        from tests.fixtures.models.vehicle import Vehicle

        if isinstance(obj, Author):
            return ChunkParams(
                max_depth=self.AUTHOR_MAX_DEPTH,
                max_chunks=self.AUTHOR_MAX_CHUNKS,
                max_tokens_per_chunk=self.AUTHOR_MAX_TOKENS,
            )
        elif isinstance(obj, Vehicle):
            return ChunkParams(
                max_depth=self.VEHICLE_MAX_DEPTH,
                max_chunks=self.VEHICLE_MAX_CHUNKS,
                max_tokens_per_chunk=self.VEHICLE_MAX_TOKENS,
            )
        elif isinstance(obj, Book):
            return ChunkParams(
                max_depth=self.BOOK_MAX_DEPTH,
                max_chunks=self.BOOK_MAX_CHUNKS,
                max_tokens_per_chunk=self.BOOK_MAX_TOKENS,
            )
        return ChunkParams(
            max_depth=self.DEFAULT_MAX_DEPTH,
            max_chunks=self.DEFAULT_MAX_CHUNKS,
            max_tokens_per_chunk=self.DEFAULT_MAX_TOKENS,
        )


def test_default_ingesting_with_custom_chunk_strategy():
    """Test that DefaultIngesting uses custom chunk strategy params based on model types."""
    from tests.fixtures.models.author import Author
    from tests.fixtures.models.book import Book
    from tests.fixtures.models.vehicle import Vehicle

    custom_strategy = _CustomChunkStrategy()
    custom_ingester = DefaultIngesting(chunk_strategy=custom_strategy)

    author = Author(id=1, name='Test Author', email='test@example.com')
    vehicle = Vehicle(
        id=1,
        make='Toyota',
        model='Camry',
        year=2020,
        engine_type='gasoline',
        fuel_efficiency=25.0,
        specs={'engine': '2.5L'},
        price=25000.0,
        warranty_years=5,
        maintenance_schedule={'10000': {'oil_change': 'required'}},
        safety_ratings={'nhtsa': 5},
    )
    book = Book(id=1, title='Test Book', author_id=1, published_year=2020)

    author_params = custom_ingester.chunk_strategy.get_params(author)
    assert author_params.max_depth == _CustomChunkStrategy.AUTHOR_MAX_DEPTH
    assert author_params.max_chunks == _CustomChunkStrategy.AUTHOR_MAX_CHUNKS
    assert author_params.max_tokens_per_chunk == _CustomChunkStrategy.AUTHOR_MAX_TOKENS

    vehicle_params = custom_ingester.chunk_strategy.get_params(vehicle)
    assert vehicle_params.max_depth == _CustomChunkStrategy.VEHICLE_MAX_DEPTH
    assert vehicle_params.max_chunks == _CustomChunkStrategy.VEHICLE_MAX_CHUNKS
    assert vehicle_params.max_tokens_per_chunk == _CustomChunkStrategy.VEHICLE_MAX_TOKENS

    book_params = custom_ingester.chunk_strategy.get_params(book)
    assert book_params.max_depth == _CustomChunkStrategy.BOOK_MAX_DEPTH
    assert book_params.max_chunks == _CustomChunkStrategy.BOOK_MAX_CHUNKS
    assert book_params.max_tokens_per_chunk == _CustomChunkStrategy.BOOK_MAX_TOKENS

    def mock_embed(text: str) -> list[float]:  # noqa: ARG001
        return [0.1] * _CustomChunkStrategy.EMBEDDING_DIM

    author_records = custom_ingester.generate_embeddings(author, embed_func=mock_embed)
    assert len(author_records) <= _CustomChunkStrategy.AUTHOR_MAX_CHUNKS

    vehicle_records = custom_ingester.generate_embeddings(vehicle, embed_func=mock_embed)
    assert len(vehicle_records) <= _CustomChunkStrategy.VEHICLE_MAX_CHUNKS

    book_records = custom_ingester.generate_embeddings(book, embed_func=mock_embed)
    assert len(book_records) <= _CustomChunkStrategy.BOOK_MAX_CHUNKS

    for records in [author_records, vehicle_records, book_records]:
        assert all(len(r.embedding) == _CustomChunkStrategy.EMBEDDING_DIM for r in records)
