import copy
import json
from pathlib import Path

import pytest

from test_bounded_contract_v2 import context, loop_for
from test_case_builder import case_setup
from test_investigation_tools import tool_case
from test_investigation_startup import Adapter
from test_investigator import CASE, run, state
from viano_bug_analysis.analysis import Investigator, InvestigationError
from viano_bug_analysis.analysis.investigation_lean_v20 import validate_reply, normalize
from viano_bug_analysis.analysis.investigation_prompt import build_prompt
from viano_bug_analysis.budget import BudgetManager
from viano_bug_analysis.lifecycle import CheckpointStore, Workspace
from viano_bug_analysis.llm import ProviderResponse
from viano_bug_analysis.package import PackageReader, PackageValidator
from test_case_interpreter import gateway


CORPUS = Path(__file__).parent/'fixtures/llm_rejections/20261020/bug17-checkpoint-recall'


def checkpoint_context():
    c = context()
    c['lean_session'] = {'checkpoint':True,'intent':'review','turns':[{
        'evidence_id':None,'observation':'Read the next page.','llm_ref':'fixture',
        'read_ids':list(c['evidence'])}]}
    return c


@pytest.mark.parametrize('name',['6849c974','dfceb150','a641cb5c'])
def test_actual_checkpoint_answers_pass_both_validation_layers(name):
    c = checkpoint_context()
    answer = json.loads((CORPUS/name/'parsed-answer.json').read_text(encoding='utf-8'))
    before = copy.deepcopy(answer)
    validate_reply(c,answer)
    normalized = normalize(c,answer)
    assert normalized['next'] == answer['next']
    assert answer == before
    assert normalized['update']['working_brief']['findings']


def test_invalid_action_feedback_names_allowed_checkpoint_actions():
    c = checkpoint_context()
    answer = json.loads((CORPUS/'6849c974/parsed-answer.json').read_text(encoding='utf-8'))
    answer['next'] = {'kind':'invented'}
    with pytest.raises(InvestigationError) as error:
        validate_reply(c,answer)
    info = error.value.feedback
    assert info['validation_path'] == '/next/kind'
    assert 'recall' in info['requirement'] and 'finish' in info['requirement']
    assert len(json.dumps(info)) < 500


class RecallCheckpoint(Adapter):
    async def generate(self,config,request,**kwargs):
        body = json.loads(request['messages'][1]['content'])
        if body['task'] == 'investigate_one_step':
            self.calls.append(body)
            return ProviderResponse({'assessment':None,'observation':'Consolidate before reading another page.',
                'next':{'kind':'checkpoint','intent':'review'}},'fixture',100,50)
        if body['task'] == 'consolidate_investigation_session':
            self.calls.append(body)
            return ProviderResponse({'assessment':None,'confidence':'plausible','findings':[],
                'uncertainties':['More pages needed.'],'next_question':'Read the next evidence page.',
                'retire':[],'next':{'kind':'recall','evidence_refs':['E000001'],'page':0}},'fixture',100,50)
        return await super().generate(config,request,**kwargs)


def test_checkpoint_recall_persists_brief_then_returns_to_short_loop(tool_case,tmp_path):
    adapter = RecallCheckpoint()
    loop = loop_for(tool_case,adapter)
    run(loop,steps=5)
    before = BudgetManager(tool_case[0]).snapshot(*CASE).state
    result = run(loop,steps=1)
    saved = state(result)
    assert result.status != 'BLOCKED'
    assert saved['hypotheses']['H1']['lean_session'] == {'turns':[],'checkpoint':False}
    assert saved['recalled_pages']['H1'] == {'E000001':0}
    assert saved['hypotheses']['H1']['working_brief_ref']
    assert BudgetManager(tool_case[0]).snapshot(*CASE).state == before
    assert PackageValidator().validate(result.package).valid
    recipient = CheckpointStore(Workspace(tmp_path/'recipient'))
    recipient.intake(result.package)
    restored = Investigator(recipient,gateway(adapter),router=loop.router)
    run(restored,steps=1)
    assert adapter.calls[-1]['task'] == 'investigate_one_step'
    assert 'E000001' in adapter.calls[-1]['input']['evidence']


@pytest.mark.parametrize('change,code',[('unknown','INVESTIGATION_EVIDENCE_UNKNOWN'),('page','INVESTIGATION_EVIDENCE_PAGE')])
def test_checkpoint_recall_cannot_read_unknown_evidence_or_page(tool_case,change,code):
    class Invalid(RecallCheckpoint):
        async def generate(self,config,request,**kwargs):
            response = await super().generate(config,request,**kwargs)
            if json.loads(request['messages'][1]['content'])['task'] == 'consolidate_investigation_session':
                response.document['next']['evidence_refs'] = ['E999999'] if change == 'unknown' else ['E000001']
                response.document['next']['page'] = 99999 if change == 'page' else 0
            return response
    loop = loop_for(tool_case,Invalid())
    run(loop,steps=5)
    before = BudgetManager(tool_case[0]).snapshot(*CASE).state
    result = run(loop,steps=1)
    assert state(result)['hypotheses']['H1']['notes'][-1]['code'] == code
    assert BudgetManager(tool_case[0]).snapshot(*CASE).state == before
    assert 'H1' not in state(result).get('recalled_pages',{})


def test_early_diagnosis_review_still_forbids_recall():
    c = checkpoint_context()
    c['early_review'] = {'request':{'status':'preparing'}}
    answer = json.loads((CORPUS/'6849c974/parsed-answer.json').read_text(encoding='utf-8'))
    with pytest.raises(InvestigationError):
        validate_reply(c,answer)