"""Regression for work listings losing titles/types and merging same-year work."""
import copy
import json
from pathlib import Path

import pytest

from tools.answer_evaluation.answer_format import parse_answer
from tools.answer_evaluation.checked_generation import generate_checked_answer
from tools.answer_evaluation.context_policy import answer_context_issues
from tools.answer_evaluation.graph_context import graph_context
from tools.answer_evaluation.render_answer import render_graph_answer
from tools.graphrag import graph_search as gs
from tools.graphrag.tests.test_graph_search import FakeSession, sample_catalog


QUESTION = 'Aké práce má Ján Holp a v ktorých rokoch?'
URL = 'https://example.test/jan'


@pytest.mark.parametrize('question', [
    QUESTION, 'Aké práce má Ján Holp?', 'Ktoré práce má Ján Holp?',
    'Vypíš práce Jána Holpa.', 'Uveď všetky práce Jána Holpa.',
    'Zoznam prác Jána Holpa', 'Prehľad projektov Jána Holpa',
    'Aký diplomový projekt mal Ján Holp v roku 2020?',
    'Aké má Ján Holp práce a v ktorých rokoch?',
])
def test_work_listing_intent_without_explicit_title_word(question):
    assert gs.extract_query_constraints(question)['mentions_work_listing']


@pytest.mark.parametrize('question', [
    'V ktorých rokoch má Ján Holp práce?', 'Aký je rok práce Jána Holpa?',
    'Aký je názov diplomovej práce Jána Holpa?',
    'Koľko prác má Ján Holp?', 'Ktoré osoby majú práce z roku 2020?',
    'Aké kľúčové slová má Ján Holp?',
])
def test_other_projections_are_not_changed_into_work_listings(question):
    assert not gs.extract_query_constraints(question)['mentions_work_listing']


@pytest.fixture
def retrieved():
    common = {'person': 'Ján Holp', 'document_path': 'p.md', 'source_url': URL}
    theses = [
        {'thesis_id': 'thesis:bp', 'thesis_type': 'bachelor_thesis', 'year': 2019,
         'heading': 'Bakalárska práca 2019', 'title': 'Vyhľadávač na slovenskom internete'},
        {'thesis_id': 'thesis:dp', 'thesis_type': 'diploma_thesis', 'year': 2021,
         'heading': 'Diplomová práca 2021', 'title': 'Systém získavania informácií v slovenskom jazyku.'}]
    projects = [
        {'work_id': 'work:tp', 'work_type': 'team_project', 'year': 2019,
         'heading': 'Tímový projekt 2019', 'title': None},
        {'work_id': 'work:dp', 'work_type': 'diploma_project', 'year': 2020,
         'heading': 'Diplomový projekt 2020', 'title': None},
        {'work_id': 'work:dp2', 'work_type': 'diploma_project', 'year': 2020,
         'heading': 'Diplomový projekt 2  2020', 'title': None}]
    return {'status': 'ok', 'reason': None, 'complete': True, 'truncated': False,
            'truncation_reasons': [], 'index_version': {'version': 'fixture'}, 'aggregation': None,
            'constraints': gs.extract_query_constraints(QUESTION),
            'sources': [{'document_path': 'p.md', 'source_url': URL}],
            'facts': [{**common, **row, 'fact_type': 'person_thesis'} for row in theses] +
                     [{**common, **row, 'fact_type': 'person_work'} for row in projects]}


def listing_answer(context):
    keys = ['person', 'record_id', 'title', 'year', 'work_type', 'work_label']
    return {'status': 'ok', 'answer': '', 'source_urls': [URL], 'count': None,
            'claims': [{key: row[key] for key in keys if row.get(key) is not None}
                       for row in context['required_academic_records']]}


def completion(answer):
    return {'choices': [{'message': {'content': json.dumps(answer, ensure_ascii=False)}}]}


def test_search_without_year_or_title_fetches_both_academic_kinds():
    fact = {'person': 'Jakub Maruniak', 'person_element_id': 'p2', 'work_id': 'work:dp2',
            'work_element_id': 'w2', 'heading': 'Diplomový projekt 2 2020', 'year': 2020,
            'work_type': 'diploma_project', 'title': None}
    session = FakeSession({gs.CATALOG_QUERY: sample_catalog(), gs.WORK_FACTS_QUERY: [fact]})
    result = gs.search_graph_with_session(session, 'Aké práce má Jakub Maruniak?')
    calls = {query for query, _ in session.calls}
    assert {gs.THESIS_FACTS_QUERY, gs.WORK_FACTS_QUERY} <= calls
    assert result['facts'][0]['work_id'] == 'work:dp2'
    assert result['facts'][0]['title'] is None


@pytest.mark.parametrize('question,expected,excluded', [
    ('Aké projekty má Jakub Maruniak?', gs.WORK_FACTS_QUERY, gs.THESIS_FACTS_QUERY),
    ('Aké záverečné práce má Jakub Maruniak?', gs.THESIS_FACTS_QUERY, gs.WORK_FACTS_QUERY),
])
def test_listing_preserves_requested_domain(question, expected, excluded):
    session = FakeSession({gs.CATALOG_QUERY: sample_catalog()})
    gs.search_graph_with_session(session, question)
    calls = {query for query, _ in session.calls}
    assert expected in calls and excluded not in calls


def test_context_preserves_five_records_instead_of_three_years(retrieved):
    context = graph_context(retrieved)
    assert len(context['required_academic_records']) == 5
    assert len(context['facts']) == 5
    assert {r['record_id'] for r in context['facts']} == {'thesis:bp', 'thesis:dp', 'work:tp', 'work:dp', 'work:dp2'}
    assert [r['title'] for r in context['facts']].count(None) == 3
    assert context['complete'] is True


@pytest.mark.parametrize('change', ['years_only', 'no_types', 'missing_project', 'no_record_ids'])
def test_guard_rejects_incomplete_response_even_when_retrieval_is_complete(retrieved, change):
    context = graph_context(retrieved)
    answer = listing_answer(context)
    if change == 'years_only':
        answer['claims'] = [{'person': 'Ján Holp', 'year': year} for year in (2019, 2020, 2021)]
    elif change == 'no_types':
        for claim in answer['claims']: claim.pop('work_type')
    elif change == 'missing_project':
        answer['claims'].pop()
    else:
        for claim in answer['claims']: claim.pop('record_id')
    assert 'missing_academic_record_details' in answer_context_issues(answer, context, check_text=False)


@pytest.mark.parametrize('field,value', [('year', 2020), ('work_type', 'team_project'),
                                       ('record_id', 'work:dp2'), ('record_id', 'thesis/dp'),
                                       ('work_label', 'Diplomový projekt 2020')])
def test_record_fields_cannot_be_mixed_between_works(retrieved, field, value):
    context = graph_context(retrieved)
    answer = listing_answer(context)
    answer['claims'][1][field] = value
    assert 'unsupported_academic_binding' in answer_context_issues(answer, context, check_text=False)


def test_citation_must_support_the_whole_record_not_separate_field_occurrences(retrieved):
    context = graph_context(retrieved)
    claim = {'person': 'Ján Holp', 'title': 'A', 'year': 2021}
    context.pop('required_academic_records')
    other = 'https://example.test/other'
    context['sources'].append({'source_url': other})
    context['facts'] = [
        {'person': 'Ján Holp', 'title': 'A', 'year': 2021, 'source_url': URL},
        {'person': 'Ján Holp', 'title': 'A', 'year': 2020, 'source_url': other},
        {'person': 'Ján Holp', 'title': 'B', 'year': 2021, 'source_url': other}]
    answer = {'status': 'ok', 'answer': '', 'claims': [claim], 'source_urls': [other]}
    assert 'missing_academic_record_source' in answer_context_issues(answer, context, check_text=False)
    answer['source_urls'].append(URL)
    assert answer_context_issues(answer, context, check_text=False) == []


def test_missing_title_is_reported_without_invention_and_ids_stay_private(retrieved):
    context = graph_context(retrieved)
    answer, _ = parse_answer(json.dumps(listing_answer(context)))
    assert answer_context_issues(answer, context, check_text=False) == []
    rendered = render_graph_answer(answer, context)
    assert answer_context_issues(rendered, context) == []
    assert rendered['answer'].count('názov v zdroji nie je uvedený') == 3
    assert 'Diplomový projekt 2  2020' in rendered['answer']
    assert 'Systém získavania informácií' in rendered['answer']
    assert '**Bakalárska práca 2019**' in rendered['answer']
    assert 'typ práce:' not in rendered['answer']
    assert all(claim['record_id'] not in rendered['answer'] for claim in answer['claims'])
    answer['claims'][2]['title'] = 'Vymyslený názov projektu'
    assert 'unsupported_academic_binding' in answer_context_issues(answer, context, check_text=False)


def test_unknown_year_is_not_filled_from_heading_or_another_record(retrieved):
    retrieved['facts'][2]['year'] = None
    retrieved['facts'][2]['heading'] = 'Tímový projekt'
    context = graph_context(retrieved)
    answer = listing_answer(context)
    rendered = render_graph_answer(answer, context)
    assert 'rok v zdroji nie je uvedený' in rendered['answer']
    assert answer_context_issues(rendered, context) == []
    answer['claims'][2]['year'] = 2019
    assert 'unsupported_academic_binding' in answer_context_issues(answer, context, check_text=False)


def test_known_record_cannot_be_refused(retrieved):
    context = graph_context(retrieved)
    answer = {'status': 'no_answer', 'answer': '', 'claims': [], 'source_urls': []}
    assert 'rejected_supported_academic_listing' in answer_context_issues(answer, context, check_text=False)


def test_listing_does_not_require_records_outside_truncated_context(retrieved):
    retrieved.update(complete=False, truncated=True, truncation_reasons=['result_limit'])
    retrieved['facts'] = retrieved['facts'][:2]
    context = graph_context(retrieved)
    rendered = render_graph_answer(listing_answer(context), context)
    assert answer_context_issues(rendered, context) == []
    assert 'Výpis je neúplný.' in rendered['answer']


def test_year_only_question_still_returns_only_years(retrieved):
    retrieved['constraints'] = gs.extract_query_constraints('V ktorých rokoch má Ján Holp práce?')
    context = graph_context(retrieved)
    assert 'required_academic_records' not in context
    assert len(context['facts']) == 3
    assert all(set(row) == {'person', 'year', 'document_path', 'source_url'} for row in context['facts'])


def test_production_delivery_repairs_year_only_answer_and_renders_all_records(monkeypatch, retrieved):
    from tools.answers import answer_delivery as delivery
    monkeypatch.setattr(delivery, 'search_graph', lambda query, **kwargs: retrieved)
    good = listing_answer(graph_context(retrieved))
    bad = {**good, 'claims': [{'person': 'Ján Holp', 'year': year} for year in (2019, 2020, 2021)]}
    replies = iter([bad, good]); messages = []
    def complete(rows):
        messages.append(copy.deepcopy(rows))
        return completion(next(replies))
    result = delivery.deliver_answer(QUESTION, engine='graphrag', complete=complete)
    assert result['accepted_for_delivery']
    assert len(result['claims']) == 5 and len(messages) == 2
    assert 'missing_academic_record_details' in result['attempts'][0]['context_issues']
    assert 'Diplomový projekt 2  2020' in result['answer']
    assert 'expected_bindings' not in json.dumps(messages)


def test_two_year_only_attempts_are_never_delivered(retrieved):
    context = graph_context(retrieved)
    answer = {'status': 'ok', 'answer': '', 'source_urls': [URL],
              'claims': [{'person': 'Ján Holp', 'year': 2021}]}
    prompt = Path('tools/answer_evaluation/graph_answer_prompt.txt').read_text()
    result = generate_checked_answer(QUESTION, context, prompt, lambda _: completion(answer))
    assert not result['accepted_for_delivery'] and len(result['attempts']) == 2