import json
from io import BytesIO
from datetime import date, datetime, timedelta, timezone
from types import SimpleNamespace
from unittest.mock import AsyncMock, Mock
from uuid import uuid4
import pytest
import pytest_asyncio
from fastapi import UploadFile
from sqlalchemy import select, func
from sqlalchemy.ext.asyncio import create_async_engine, async_sessionmaker
from src.models import (RecognitionJob, StoredFile, FinancialStatementMeta,
                        FinancialStatementLine, FinancialStatementValue)
from src.models.base import Base
from src.schemas.enums import RecognitionStatus, RecognitionPipeline, FileStatus, DocumentType, FileType, StatementType, FslRowKind
from src.schemas.pdfworker import PDFWorkerResult
from src.services import pdfworker_results
from src.services.file_service import FileService
from src.repositories.file_repository import FileRepository
from src.repositories.recognition_job_repo import RecognitionJobRepository
from src.core import clients


@pytest_asyncio.fixture
async def db():
    engine = create_async_engine('sqlite+aiosqlite:///:memory:')
    tables = [StoredFile.__table__, RecognitionJob.__table__, FinancialStatementMeta.__table__,
              FinancialStatementLine.__table__, FinancialStatementValue.__table__]
    async with engine.begin() as connection:
        await connection.run_sync(lambda conn: Base.metadata.create_all(conn, tables=tables))
    factory = async_sessionmaker(engine, expire_on_commit=False)
    ids = SimpleNamespace(task=uuid4(), file=uuid4(), legal=uuid4())
    async with factory() as session:
        session.add(StoredFile(id=ids.file, legal_entity_id=ids.legal, file_type=FileType.PDF,
                              document_type=DocumentType.OFR,
                              original_name='report.pdf', s3_key='folder/nested/report.pdf',
                              size_bytes=123, period_year=2026, current_recognition_job_id=ids.task))
        session.add(RecognitionJob(id=ids.task, file_id=ids.file, legal_entity_id=ids.legal,
                                   pipeline=RecognitionPipeline.IMAGE, status=RecognitionStatus.QUEUED, attempt=1))
        await session.commit()
    yield factory, ids
    await engine.dispose()


def ready_event(ids, status='ready'):
    return PDFWorkerResult(task_id=ids.task, file_id=ids.file, status=status,
        result_s3_key='results/full/path.json' if status == 'ready' else None,
        error_code=None if status == 'ready' else 'storage_unavailable', finished_at=datetime.now(timezone.utc))


def setup_storage(ids, monkeypatch):
    payload = dict(task_id=str(ids.task), job_id=str(ids.task), file_id=str(ids.file), files=[{}])
    response = SimpleNamespace(read=AsyncMock(return_value=json.dumps(payload).encode()), release=Mock())
    storage = SimpleNamespace(get_object=AsyncMock(return_value=response))
    meta_id, line_id = uuid4(), uuid4()
    converter = Mock(return_value=dict(
        meta=[dict(id=meta_id, stored_file_id=ids.file, statement_type=StatementType.BALANCE,
                   inn='123', inn_verified=True, units='руб', units_verified=True,
                   report_date=date(2026,1,1), report_date_verified=True)],
        rows=[dict(id=line_id, recognition_job_id=ids.task, meta_id=meta_id,
                   legal_entity_id=ids.legal, line_code='1600', line_code_recognized=True,
                   line_name='Итого', line_name_recognized=True, row_kind=FslRowKind.ROW, ordinal=1)],
        values=[dict(id=uuid4(), line_id=line_id, amount=100, string_value='100', column_index=1, is_recognized=True)],
        skipped_files=[]))
    monkeypatch.setattr(pdfworker_results, 'json_to_db_rows', converter)
    return storage, response, converter


@pytest.mark.asyncio
async def test_ready_commits_once(db, monkeypatch):
    factory, ids = db
    storage, response, converter = setup_storage(ids, monkeypatch)
    event = ready_event(ids)
    await pdfworker_results.apply_pdfworker_result(event, factory, storage)
    await pdfworker_results.apply_pdfworker_result(event, factory, storage)
    storage.get_object.assert_awaited_once_with('results/full/path.json')
    response.release.assert_called_once()
    assert converter.call_args.args[0]['legal_entity_id'] == str(ids.legal)
    async with factory() as session:
        assert (await session.get(RecognitionJob, ids.task)).status == RecognitionStatus.DONE
        stored_file = await session.get(StoredFile, ids.file)
        assert stored_file.status == FileStatus.CHECK
        assert stored_file.document_type == DocumentType.OFR
        for model in (FinancialStatementMeta, FinancialStatementLine, FinancialStatementValue):
            assert await session.scalar(select(func.count()).select_from(model)) == 1


@pytest.mark.asyncio
async def test_upload_accepts_new_extension_and_saves_ofr_type(db):
    factory, ids = db
    storage = Mock(put_object=AsyncMock())
    async with factory() as session:
        service = FileService(storage, session, FileRepository(session),
                              RecognitionJobRepository(session), None)
        uploaded = UploadFile(filename='report.docx', file=BytesIO(b'data'), size=4)
        created = await service.save_file(uploaded, ids.legal, 'Q1', 2026)
        stored = await session.get(StoredFile, created['file']['id'])
        job = await session.get(RecognitionJob, created['rec_job']['id'])
        assert stored.file_type == FileType.FILE
        assert stored.document_type == DocumentType.OFR
        assert job.file_id == stored.id
        assert service._get_transformed_file_job_info(stored)['file_type'] == 'DOCX'
        storage.put_object.assert_awaited_once_with(f'{stored.id}.docx', b'data')


@pytest.mark.asyncio
@pytest.mark.parametrize('status', ['error', 'timeout'])
async def test_error_without_minio(db, status):
    factory, ids = db
    data = ready_event(ids, 'error').model_dump()
    if status == 'timeout':
        data.update(status='timeout', error_code='deadline_exceeded')
    event = PDFWorkerResult(**data)
    storage = Mock(get_object=AsyncMock())
    await pdfworker_results.apply_pdfworker_result(event, factory, storage)
    await pdfworker_results.apply_pdfworker_result(event, factory, storage)
    storage.get_object.assert_not_awaited()
    async with factory() as session:
        job = await session.get(RecognitionJob, ids.task)
        assert job.status == RecognitionStatus.ERROR
        assert event.error_code in job.error_message
        assert (await session.get(StoredFile, ids.file)).status == FileStatus.ERROR


@pytest.mark.asyncio
async def test_invalid_json_rolls_back(db, monkeypatch):
    factory, ids = db
    storage, response, _ = setup_storage(ids, monkeypatch)
    response.read.return_value = b'{"files": []}'
    with pytest.raises(ValueError):
        await pdfworker_results.apply_pdfworker_result(ready_event(ids), factory, storage)
    response.release.assert_called_once()
    async with factory() as session:
        assert (await session.get(RecognitionJob, ids.task)).status == RecognitionStatus.QUEUED
        assert await session.scalar(select(func.count()).select_from(FinancialStatementMeta)) == 0


@pytest.mark.asyncio
async def test_pollworker_recovers_same_command(db, monkeypatch):
    factory, ids = db
    monkeypatch.setattr(clients, 'async_session_maker', factory)
    worker = clients.PollWorker(SimpleNamespace(jetstream=lambda: SimpleNamespace()))
    first, = await worker.claim_pending_rows()
    assert first.source_s3_key == 'folder/nested/report.pdf'
    assert first.task_id == ids.task
    assert not await worker.claim_pending_rows()
    async with factory() as session:
        job = await session.get(RecognitionJob, ids.task)
        job.started_at = datetime.now(timezone.utc) - timedelta(minutes=2)
        await session.commit()
    retry, = await worker.claim_pending_rows()
    assert first.model_dump_json() == retry.model_dump_json()
    await worker.mark_published(first.task_id)
    assert not await worker.claim_pending_rows()


@pytest.mark.asyncio
async def test_cannot_requeue_finished_job(db, monkeypatch):
    factory, ids = db
    monkeypatch.setattr(clients, 'async_session_maker', factory)
    await pdfworker_results.apply_pdfworker_result(ready_event(ids, 'error'), factory, Mock())
    worker = clients.PollWorker(SimpleNamespace(jetstream=lambda: None))
    await worker.mark_queued(ids.task)
    async with factory() as session:
        assert (await session.get(RecognitionJob, ids.task)).status == RecognitionStatus.ERROR
        assert (await session.get(StoredFile, ids.file)).status == FileStatus.ERROR


@pytest.mark.asyncio
async def test_invalid_stored_key_does_not_block_queue(db, monkeypatch):
    factory, ids = db
    monkeypatch.setattr(clients, 'async_session_maker', factory)
    async with factory() as session:
        file = await session.get(StoredFile, ids.file)
        file.s3_key = None
        await session.commit()
    worker = clients.PollWorker(SimpleNamespace(jetstream=lambda: None))
    assert await worker.claim_pending_rows() == []
    async with factory() as session:
        job = await session.get(RecognitionJob, ids.task)
        assert job.status == RecognitionStatus.ERROR
        assert 'source_s3_key' in job.error_message


@pytest.mark.asyncio
async def test_wrong_file_id_does_not_change_task(db):
    factory, ids = db
    data = ready_event(ids, 'error').model_dump()
    data['file_id'] = uuid4()
    with pytest.raises(ValueError, match='file_id'):
        await pdfworker_results.apply_pdfworker_result(PDFWorkerResult(**data), factory, Mock())
    async with factory() as session:
        assert (await session.get(RecognitionJob, ids.task)).status == RecognitionStatus.QUEUED


@pytest.mark.asyncio
async def test_dispatches_flagged_document_type(db, monkeypatch):
    factory, ids = db
    storage, _, _ = setup_storage(ids, monkeypatch)
    async with factory() as session:
        stored = await session.get(StoredFile, ids.file)
        stored.document_type = DocumentType.OSV
        await session.commit()
    handler = AsyncMock()
    monkeypatch.setitem(pdfworker_results.RESULT_HANDLERS, DocumentType.OFR, handler)
    await pdfworker_results.apply_pdfworker_result(ready_event(ids), factory, storage)
    handler.assert_awaited_once()
    async with factory() as session:
        job = await session.get(RecognitionJob, ids.task)
        assert job.status == RecognitionStatus.DONE
        assert (await session.get(StoredFile, ids.file)).document_type == DocumentType.OFR


@pytest.mark.asyncio
async def test_missing_handler_does_not_complete_job(db, monkeypatch):
    factory, ids = db
    monkeypatch.delitem(pdfworker_results.RESULT_HANDLERS, DocumentType.OFR)
    storage = Mock(get_object=AsyncMock())
    with pytest.raises(ValueError, match='handler'):
        await pdfworker_results.apply_pdfworker_result(ready_event(ids), factory, storage)
    storage.get_object.assert_not_awaited()
    async with factory() as session:
        assert (await session.get(RecognitionJob, ids.task)).status == RecognitionStatus.QUEUED