import unittest
from unittest.mock import Mock

from src.step3_question_answer.knowledge_answers import synthesize_answers
from src.step3_question_answer.recommendation import Recommender
from src.step3_question_answer.tests.test_knowledge_answers import fixture, context, RECORD


def entry(key, answer='提交申请表，再在线提交。'):
    return {'source_id': key, 'answer': answer, 'exact_qa': False,
            'citations': [{'path': '/data/answer', 'quote': answer}]}


class AnswerSynthesisTests(unittest.TestCase):
    def test_final_sources_repair_regenerates_from_original_evidence(self):
        model = Mock(side_effect=[{'text': '错误草稿', 'source_ids': ['unknown']},
                                  {'text': '重新核验后的回答', 'source_ids': ['a']}])
        result = synthesize_answers(model, '如何办理', [entry('a'), entry('b')], [], 'specific_need')
        self.assertEqual(result['status'], 'generated')
        self.assertEqual(result['text'], '重新核验后的回答')
        first, second = [c.args[1] for c in model.call_args_list]
        self.assertEqual(first['evidence'], second['evidence'])
        self.assertEqual(second['source_repair']['allowed_source_ids'], ['a', 'b'])
        self.assertNotIn('错误草稿', str(second))
        self.assertEqual(result['diagnostics']['source_repair_attempts'], 1)

    def test_reduction_sources_repair_is_batch_local_and_bounded(self):
        from src.step3_question_answer.knowledge_answers import REDUCE_PROMPT, input_size
        for recover in (True, False):
            calls = []
            def model(prompt, data):
                calls.append((prompt, data))
                self.assertLessEqual(input_size(prompt, data), 32768 - 16384 - 904)
                refs = list(dict.fromkeys(r for e in data['evidence']
                                         for r in e.get('source_ids', [e.get('source_id')])))
                if prompt.startswith(REDUCE_PROMPT) and len(calls) <= 2:
                    if len(calls) == 1 or not recover:
                        return {'text': '笔记', 'source_ids': ['other-batch']}
                    self.assertEqual(data['source_repair']['allowed_source_ids'], sorted(refs))
                    self.assertEqual(data['evidence'], calls[0][1]['evidence'])
                return {'text': '条件仍待核验。', 'source_ids': refs}
            result = synthesize_answers(model, '如何办理',
                                        [entry('a', '合成材料。' * 1000), entry('b')], [], 'specific_need')
            self.assertEqual(result['status'], 'generated' if recover else 'fallback')
            self.assertEqual(result['diagnostics']['source_repair_attempts'], 1)
            if not recover:
                self.assertEqual(len(calls), 2)
                self.assertEqual(result['error'], 'summary_invalid_sources')

    def test_short_guide_uses_synthesis_and_keeps_sources_in_both_modes(self):
        from src.step3_question_answer.report import markdown
        record = {**entry('guide', '合成事项可在合成窗口提交申请表，受理后审查并告知结果。'),
                  'knowledge_type': 'service_guide', 'title': '合成办事指南'}
        text = ('可在合成窗口申请该事项。\n\n### 办理渠道\n\n- 合成窗口。'
                '\n\n### 申请材料\n\n- 申请表。\n\n### 办理流程\n\n1. 提交申请。\n2. 受理后审查并告知结果。')
        for kind in ('policy_overview', 'specific_need'):
            model = Mock(return_value={'text': text, 'source_ids': ['guide']})
            result = synthesize_answers(model, '介绍办理指南', [record], [], kind)
            model.assert_called_once()
            self.assertEqual(result['status'], 'generated')
            self.assertEqual(model.call_args.args[1]['evidence'][0]['knowledge_type'], 'service_guide')
            self.assertEqual(result['citations'][0]['quote'], record['answer'])
            self.assertEqual(result['text'], text.replace('### 办理渠道', '### 一、办理渠道').replace('### 申请材料', '### 二、申请材料').replace('### 办理流程', '### 三、办理流程'))
            rendered = markdown(result['text'])
            self.assertLess(rendered.index('<h3>一、办理渠道'), rendered.index('<h3>二、申请材料'))
            self.assertLess(rendered.index('<h3>二、申请材料'), rendered.index('<h3>三、办理流程'))
            self.assertIn('<ol>', rendered)

    def test_short_finance_answer_is_formatted_without_losing_evidence(self):
        record = {**entry('finance', '合成周转贷最高500万元，须提供抵押。'),
                  'knowledge_type': 'finance', 'title': '合成周转贷'}
        text = '### 合成周转贷\n\n- **额度**：最高500万元。\n- **条件**：须提供抵押。'
        model = Mock(return_value={'text': text, 'source_ids': ['finance']})
        result = synthesize_answers(model, '融资300万元有哪些产品', [record], [], 'specific_need')
        model.assert_called_once()
        evidence = model.call_args.args[1]['evidence'][0]
        self.assertEqual(evidence['knowledge_type'], 'finance')
        self.assertEqual(evidence['title'], record['title'])
        self.assertEqual(result['text'], text)
        self.assertEqual(result['citations'][0]['quote'], record['citations'][0]['quote'])

    def test_short_nonfinance_answer_keeps_fast_path(self):
        model = Mock(side_effect=AssertionError('unnecessary generation'))
        record = entry('qa')
        result = synthesize_answers(model, '提交方式', [record], [], 'policy_overview')
        model.assert_not_called()
        self.assertEqual(result['text'], record['answer'])

    def test_large_input_is_fully_batched_with_wrapped_model(self):
        from types import SimpleNamespace
        from threading import Semaphore
        from src.step4_api.runtime import limited
        from src.step3_question_answer.knowledge_answers import input_size, REDUCE_PROMPT
        seen = []
        def call(prompt, data):
            self.assertLessEqual(input_size(prompt, data) + 16384 + 904, 32768)
            seen.append((prompt, data))
            refs = list(dict.fromkeys(r for e in data['evidence']
                                     for r in e.get('source_ids', [e.get('source_id')])))
            return {'text': '保留地区限制，先申请后办理。', 'source_ids': refs}
        call.config = SimpleNamespace(context_tokens=32768, margin=512)
        call.MAX_OUTPUT_TOKENS = 16384
        model = limited(call, Semaphore(1))
        records = [entry(str(i), ('地区限制不可忽略。' * 650) + str(i)) for i in range(4)]
        result = synthesize_answers(model, '如何办理', records, [], 'policy_overview')
        self.assertEqual(result['status'], 'generated', result.get('diagnostics'))
        self.assertEqual(set(result['source_ids']), {'0', '1', '2', '3'})
        self.assertEqual(len(result['citations']), 4)
        self.assertGreater(result['diagnostics']['model_calls'], 1)
        self.assertEqual(result['diagnostics']['context_limit'], 32768)
        fragments = [e for p, d in seen if p.startswith(REDUCE_PROMPT) for e in d['evidence']]
        import json
        for record in records:
            pieces = [e['fragment'] for e in fragments if e.get('source_ids') == [record['source_id']]]
            reconstructed = json.loads(''.join(pieces))
            self.assertEqual(reconstructed['answer'], record['answer'])
            self.assertEqual(reconstructed['citations'], record['citations'])

    def test_batch_failure_does_not_publish_partial_success(self):
        model = Mock(side_effect=RuntimeError('model_request_timeout'))
        result = synthesize_answers(model, '如何办理', [entry('a', '材料。' * 10000)], [], 'policy_overview')
        self.assertEqual(result['status'], 'fallback')
        self.assertEqual(result['error'], 'model_request_timeout')
        self.assertEqual(len(result['citations'][0]['quote']), 30000)

    def test_sources_and_cards_share_one_synthesis_without_appending(self):
        model = Mock(return_value={'text': '提交申请表后在线办理。', 'source_ids': ['a', 'b']})
        result = synthesize_answers(model, '如何办理', [entry('a')], [entry('b')], 'policy_overview')
        model.assert_called_once()
        self.assertEqual(result['text'], '提交申请表后在线办理。')
        self.assertEqual(len(model.call_args.args[1]['evidence']), 2)
        self.assertEqual(len(result['citations']), 2)

    def test_long_or_chunked_single_answer_is_synthesized(self):
        for answer in ('提交申请。\n提交申请。', '材料说明。' * 100):
            model = Mock(return_value={'text': '提交申请。', 'source_ids': ['a']})
            result = synthesize_answers(model, '如何办理', [entry('a', answer)], [], 'policy_overview')
            model.assert_called_once()
            self.assertEqual(result['text'], '提交申请。')

    def test_invalid_answers_are_still_rejected(self):
        for value, code in (({'text': '编造内容', 'source_ids': ['unknown']}, 'summary_invalid_sources'),
                            ({}, 'summary_empty_text'), (None, 'summary_invalid_schema')):
            result = synthesize_answers(Mock(return_value=value), '如何办理',
                                        [entry('a'), entry('b')], [], 'policy_overview')
            self.assertEqual(result['error'], code)
            self.assertEqual(result['diagnostics']['stage'], 'validation')
            self.assertNotIn('提交申请表', result['text'])
            self.assertEqual(len(result['citations']), 2)

    def test_long_valid_answer_is_returned_without_length_retry_or_truncation(self):
        for length in (601, 1046, 1801, 3000):
            text = '字' * length
            model = Mock(return_value={'text': text, 'source_ids': ['a', 'b']})
            result = synthesize_answers(model, '如何办理', [entry('a'), entry('b')], [], 'policy_overview')
            self.assertEqual(result['status'], 'generated')
            self.assertEqual(result['text'], text)
            self.assertNotIn('error', result)
            self.assertNotIn('text_limit', result['diagnostics'])
            self.assertEqual(result['diagnostics']['output_chars'], length)
            model.assert_called_once()
            self.assertEqual(len(result['citations']), 2)

    def test_model_failures_keep_only_safe_codes_and_numeric_budget(self):
        from types import SimpleNamespace
        for code in ('analysis_context_limit', 'model_request_timeout', 'model_invalid_content_json',
                     'model_http_429', 'SECRET_EXCEPTION_BODY'):
            model = Mock(side_effect=RuntimeError(code))
            model.config = SimpleNamespace(context_tokens=32768, margin=100)
            model.MAX_OUTPUT_TOKENS = 16384
            result = synthesize_answers(model, '如何办理', [entry('a'), entry('b')], [], 'policy_overview')
            self.assertEqual(result['error'], 'summary_generation_failed' if code.startswith('SECRET') else code)
            d = result['diagnostics']
            self.assertEqual(d['stage'], 'model')
            self.assertEqual(d['reserved_total'], d['input_bytes'] + 16384 + 904)
            self.assertEqual(d['context_limit'], 32768)
            self.assertNotIn('SECRET', str(result))

    def test_explicit_detail_can_exceed_default_without_truncation(self):
        text = '资料' * 400
        result = synthesize_answers(Mock(return_value={'text': text, 'source_ids': ['a']}),
                                    '请详细介绍', [entry('a'), entry('b')], [], 'policy_overview')
        self.assertEqual(result['text'], text)
        self.assertEqual([c['source_id'] for c in result['citations']], ['a'])

    def test_recommender_merges_multiple_sources_and_preserves_full_records(self):
        records = [{**RECORD, 'id': key} for key in ('a', 'b')]
        def model(prompt, data):
            if prompt.startswith('将已核验'):
                return {'text': '接受电子材料，无需复印件。', 'source_ids': ['a', 'b']}
            return {'relevant': True, 'reason': '相关', 'answer': RECORD['data']['answer'],
                    'citations': [{'path': '/data/answer', 'quote': RECORD['data']['answer']}],
                    'conditions': [], 'card': None, 'match_basis': []}
        mock = Mock(side_effect=model)
        result = Recommender(mock).answer(context('该事项怎么交材料'), fixture(records))
        self.assertEqual(result['summary']['text'], '接受电子材料，无需复印件。')
        self.assertEqual(mock.call_count, 3)
        self.assertEqual([s['full_record'] for s in result['sources']], records)
        self.assertEqual(len(result['knowledge_answers']), 2)