import json
from io import BytesIO
from types import SimpleNamespace

import pytest
from sqlalchemy import select
from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine
from starlette.datastructures import UploadFile

import app.models  # noqa: F401
from app.database import Base
from app.models.document import Document, DocumentChunk
from app.models.folder import Folder
from app.models.user import User
from app.services.auth_service import get_password_hash
from app.services.knowledge_service import KnowledgeService, SearchResult
from app.services.graph_service import GraphService


class FakeCollection:
    def __init__(self):
        self.add_calls = []
        self.delete_calls = []

    def add(self, *, ids, documents, metadatas):
        self.add_calls.append({
            'ids': ids,
            'documents': documents,
            'metadatas': metadatas,
        })

    def delete(self, *, where):
        self.delete_calls.append(where)

    def query(self, **kwargs):
        self.last_query = kwargs
        return {
            'ids': [['chunk-schema', 'chunk-rows']],
            'documents': [['schema chunk', 'row chunk']],
            'metadatas': [[
                {
                    'document_id': 'doc-1',
                    'document_title': 'Revenue',
                    'chunk_index': 0,
                    'content_type': 'table_schema',
                    'sheet_name': 'Revenue',
                    'row_start': 0,
                    'row_end': 0,
                },
                {
                    'document_id': 'doc-1',
                    'document_title': 'Revenue',
                    'chunk_index': 1,
                    'content_type': 'table_rows',
                    'sheet_name': 'Revenue',
                    'row_start': 1,
                    'row_end': 10,
                },
            ]],
            'distances': [[0.3, 0.35]],
        }


@pytest.fixture
async def knowledge_test_env(tmp_path):
    db_path = tmp_path / 'test_knowledge.db'
    engine = create_async_engine(f"sqlite+aiosqlite:///{db_path}", future=True)
    session_factory = async_sessionmaker(engine, expire_on_commit=False)

    async with engine.begin() as conn:
        await conn.run_sync(Base.metadata.create_all)

    async with session_factory() as session:
        user = User(
            email='knowledge@example.com',
            hashed_password=get_password_hash('secret123'),
            full_name='Knowledge Tester',
        )
        session.add(user)
        await session.flush()

        root = Folder(user_id=user.id, name='Finance', parent_id=None)
        session.add(root)
        await session.flush()
        child = Folder(user_id=user.id, name='Reports', parent_id=root.id)
        session.add(child)
        await session.flush()

        document = Document(
            id='doc-1',
            user_id=user.id,
            title='Revenue Workbook',
            filename='revenue.xlsx',
            file_type='xlsx',
            file_size=128,
            file_path=str(tmp_path / 'revenue.xlsx'),
            folder_id=child.id,
            summary='Revenue summary',
            chunk_count=2,
            is_indexed=False,
        )
        session.add(document)
        session.add_all([
            DocumentChunk(
                id='chunk-schema',
                document_id=document.id,
                chunk_index=0,
                content='schema chunk',
                metadata_=json.dumps({
                    'content_type': 'table_schema',
                    'sheet_name': 'Revenue',
                    'headers': ['region', 'amount'],
                    'source_order': 0,
                    'section_path': ['Revenue'],
                    'page_number': 1,
                }),
            ),
            DocumentChunk(
                id='chunk-rows',
                document_id=document.id,
                chunk_index=1,
                content='row chunk',
                metadata_=json.dumps({
                    'content_type': 'table_rows',
                    'sheet_name': 'Revenue',
                    'row_start': 1,
                    'row_end': 10,
                    'source_order': 1,
                    'section_path': ['Revenue'],
                    'page_number': 1,
                }),
            ),
        ])
        await session.commit()
        await session.refresh(user)
        await session.refresh(document)
        await session.refresh(child)
        yield session, user, document, child

    await engine.dispose()


@pytest.mark.asyncio
async def test_index_document_writes_folder_and_structure_metadata(knowledge_test_env):
    session, user, document, _folder = knowledge_test_env
    service = KnowledgeService(session, user_id=user.id)
    fake_collection = FakeCollection()
    service.get_collection = lambda user_id: fake_collection

    await service.index_document(document.id, user.id)

    assert fake_collection.add_calls
    metadatas = fake_collection.add_calls[0]['metadatas']
    assert metadatas[0]['folder_path'] == '/Finance/Reports'
    assert metadatas[0]['content_type'] == 'table_schema'
    assert metadatas[0]['sheet_name'] == 'Revenue'
    assert metadatas[1]['content_type'] == 'table_rows'

    await session.refresh(document)
    assert document.is_indexed is True
    assert document.ingestion_status == 'ready'
    assert document.indexed_at is not None


@pytest.mark.asyncio
async def test_retrieve_prefers_table_schema_for_tabular_queries(knowledge_test_env):
    session, user, _document, _folder = knowledge_test_env
    service = KnowledgeService(session, user_id=user.id)
    fake_collection = FakeCollection()
    service.get_collection = lambda user_id: fake_collection

    results = await service.retrieve('excel表 Revenue 的列有哪些', user.id, top_k=2, use_rerank=True)

    assert [item.chunk_id for item in results] == ['chunk-schema', 'chunk-rows']
    metadata = json.loads(results[0].metadata_)
    assert metadata['content_type'] == 'table_schema'


@pytest.mark.asyncio
async def test_context_expansion_uses_same_sheet_for_table_rows(knowledge_test_env):
    session, user, _document, _folder = knowledge_test_env
    service = KnowledgeService(session, user_id=user.id)

    prev_chunk, next_chunk = await service._get_related_chunks('chunk-rows', 1, 'doc-1')

    assert prev_chunk == 'schema chunk'
    assert next_chunk is None


def test_rerank_boosts_table_chunks_when_query_mentions_sheet():
    service = KnowledgeService(db=None, user_id='user-1')
    schema = SearchResult(
        chunk_id='schema',
        document_id='doc-1',
        document_title='Revenue Workbook',
        content='Columns: region amount',
        score=0.6,
        metadata_=json.dumps({'content_type': 'table_schema', 'sheet_name': 'Revenue'}),
    )
    paragraph = SearchResult(
        chunk_id='paragraph',
        document_id='doc-1',
        document_title='Revenue Workbook',
        content='General revenue narrative',
        score=0.65,
        metadata_=json.dumps({'content_type': 'paragraph', 'section_title': 'Overview'}),
    )

    ranked = service._rerank('sheet Revenue 有哪些列', [paragraph, schema], top_k=2)

    assert [item.chunk_id for item in ranked] == ['schema', 'paragraph']


@pytest.mark.asyncio
async def test_context_expansion_prefers_same_section_before_linear_neighbors(knowledge_test_env):
    session, user, document, _folder = knowledge_test_env
    session.add_all([
        DocumentChunk(
            id='chunk-overview',
            document_id=document.id,
            chunk_index=2,
            content='overview chunk',
            metadata_=json.dumps({
                'content_type': 'paragraph',
                'section_path': ['Overview'],
                'section_title': 'Overview',
                'source_order': 2,
                'page_number': 2,
            }),
        ),
        DocumentChunk(
            id='chunk-overview-2',
            document_id=document.id,
            chunk_index=3,
            content='overview details chunk',
            metadata_=json.dumps({
                'content_type': 'paragraph',
                'section_path': ['Overview'],
                'section_title': 'Overview',
                'source_order': 3,
                'page_number': 2,
            }),
        ),
        DocumentChunk(
            id='chunk-appendix',
            document_id=document.id,
            chunk_index=4,
            content='appendix chunk',
            metadata_=json.dumps({
                'content_type': 'paragraph',
                'section_path': ['Appendix'],
                'section_title': 'Appendix',
                'source_order': 4,
                'page_number': 3,
            }),
        ),
    ])
    await session.commit()

    service = KnowledgeService(session, user_id=user.id)
    prev_chunk, next_chunk = await service._get_related_chunks('chunk-overview-2', 3, 'doc-1')

    assert prev_chunk == 'overview chunk'
    assert next_chunk is None


@pytest.mark.asyncio
async def test_reindex_document_rebuilds_chunks_and_versions(knowledge_test_env):
    session, user, _document, folder = knowledge_test_env
    from app.services.document_service import DocumentService
    from app.services.knowledge_service import KnowledgeService

    upload = UploadFile(
        filename='reindex.md',
        file=BytesIO(b'# Intro\n\nOriginal content\n\n## Details\n\nUpdated content'),
    )
    doc_service = DocumentService(session)
    document = await doc_service.upload_document(user.id, upload, folder_id=folder.id)

    chunk_result = await session.execute(select(DocumentChunk).where(DocumentChunk.document_id == document.id))
    original_chunks = list(chunk_result.scalars().all())
    assert original_chunks

    service = KnowledgeService(session, user_id=user.id)
    rebuilt = await service.reindex_document(document.id, user.id)

    assert rebuilt is True
    await session.refresh(document)
    assert document.parser_version == 'v2'
    assert document.index_version == 'v2'
    assert document.ingestion_status == 'ready'

    new_chunk_result = await session.execute(
        select(DocumentChunk)
        .where(DocumentChunk.document_id == document.id)
        .order_by(DocumentChunk.chunk_index)
    )
    rebuilt_chunks = list(new_chunk_result.scalars().all())
    assert rebuilt_chunks
    assert all(chunk.metadata_ for chunk in rebuilt_chunks)


@pytest.mark.asyncio
async def test_reindex_document_chunks_reuses_existing_db_chunks(knowledge_test_env):
    session, user, document, _folder = knowledge_test_env
    service = KnowledgeService(session, user_id=user.id)
    fake_collection = FakeCollection()
    service.get_collection = lambda user_id: fake_collection

    chunk_result = await session.execute(
        select(DocumentChunk)
        .where(DocumentChunk.id == 'chunk-schema')
    )
    chunk = chunk_result.scalar_one()
    chunk.content = 'edited schema chunk'
    document.ingestion_status = 'indexing'
    await session.commit()

    rebuilt = await service.reindex_document_chunks(document.id, user.id)

    assert rebuilt is True
    assert fake_collection.delete_calls == [{'document_id': document.id}]
    assert fake_collection.add_calls
    assert fake_collection.add_calls[0]['documents'][0] == 'edited schema chunk'

    await session.refresh(document)
    assert document.ingestion_status == 'ready'
    assert document.indexed_at is not None


@pytest.mark.anyio
async def test_graph_service_uses_user_llm_configured_model(knowledge_test_env, monkeypatch):
    session, user, document, _folder = knowledge_test_env
    document.is_indexed = True
    await session.commit()

    used_providers = []

    class FakeLLM:
        async def invoke(self, _messages):
            return SimpleNamespace(content=json.dumps({
                'entities': [{'name': 'Revenue', 'type': 'topic', 'description': 'Revenue topic'}],
                'relations': [],
            }))

    async def fake_get_user_llm_config(self, user_id, model_name=None):
        assert user_id == user.id
        assert model_name is None
        return {
            'provider': 'openai',
            'model': 'user-model',
            'api_key': 'secret',
            'base_url': 'https://example.com/v1',
            'enabled': True,
        }

    def fake_create_llm_from_config(config):
        used_providers.append(config['provider'])
        return FakeLLM()

    def fail_if_global_llm_used():
        raise AssertionError('global get_llm should not be used when user config exists')

    monkeypatch.setattr('app.services.graph_service.get_llm', fail_if_global_llm_used)
    monkeypatch.setattr('app.services.graph_service.resolve_user_llm', fake_get_user_llm_config, raising=False)
    monkeypatch.setattr('app.services.graph_service._create_llm_from_config', fake_create_llm_from_config, raising=False)

    service = GraphService(session)
    await service.build_graph(user.id)

    assert used_providers == ['openai']
