from __future__ import annotations

import json
from copy import deepcopy

from app.services.memory import apply_extracted_local, rollback_auto_continuity_for_messages


class FakeDB:
    def __init__(self):
        self.memories = []
        self.state = {
            'world_bible': {
                'schema_version': 4,
                'history_backstory': [], 'background_facts': [], 'characters': [],
                'relationship_history': [], 'core_canon': [], 'locations': [],
                'secrets_knowledge': [], 'groups_factions': [], 'important_objects': [],
                'world_rules': [], 'current_state': [], 'physical_continuity': [],
                'story_timeline': [], 'custom_entries': [], 'auto_review_conflicts': [],
            },
            'hard_canon': {'manual_fact': 'User canon survives', 'auto_facts': []},
            'relationship': {'manual_status': 'user-authored'},
            'knowledge': {'facts': []},
            'scene': {'auto_facts': []},
            'physical': {'continuity_notes': []},
        }
        self.next_id = 1

    def row(self, sql, params=()):
        if 'FROM memories' in sql and 'stable_key=?' in sql:
            cid, stable = params
            xs = [m for m in self.memories if m['conversation_id'] == cid and m.get('stable_key') == stable]
            return deepcopy(xs[-1]) if xs else None
        return None

    def rows(self, sql, params=()):
        if 'FROM memories' in sql:
            cid = params[0]
            return [deepcopy(m) for m in self.memories if m['conversation_id'] == cid and m.get('source_kind') == 'Auto-created from RP']
        return []

    def execute(self, sql, params=()):
        if sql.startswith('INSERT INTO memories'):
            keys = [
                'conversation_id','memory_type','title','content','importance','source_message_id','confidence','active',
                'what_happened','who_remembers','who_involved','story_when','story_where','emotional_weight',
                'current_relevance','source_kind','original_context','stable_key','revision_history'
            ]
            row = dict(zip(keys, params))
            row['id'] = self.next_id; self.next_id += 1
            self.memories.append(row)
            return row['id']
        if sql.startswith('UPDATE memories SET memory_type='):
            # Handles both normal auto-update and rollback restore; both end with id.
            row_id = params[-1]
            row = next(x for x in self.memories if x['id'] == row_id)
            if len(params) == 17:  # normal apply_extracted_local update
                keys = ['memory_type','title','content','what_happened','importance','who_remembers','who_involved','story_when','story_where','emotional_weight','current_relevance','original_context','revision_history','confidence','source_message_id','active']
                for k, v in zip(keys, params[:-1]): row[k] = v
                row['source_kind'] = 'Auto-created from RP'
            elif len(params) == 18:  # rollback restore
                keys = ['memory_type','title','content','what_happened','importance','who_remembers','who_involved','story_when','story_where','emotional_weight','current_relevance','source_kind','original_context','revision_history','confidence','source_message_id','active']
                for k, v in zip(keys, params[:-1]): row[k] = v
            else:
                raise AssertionError(f'unexpected memory UPDATE params: {len(params)}')
            return row_id
        if sql.startswith('DELETE FROM memories'):
            row_id, cid = params
            self.memories = [m for m in self.memories if not (m['id'] == row_id and m['conversation_id'] == cid)]
            return 1
        raise AssertionError(f'unhandled SQL: {sql[:120]}')

    def get_state(self, conversation_id):
        return deepcopy(self.state)

    def update_state(self, conversation_id, doc_type, payload):
        self.state[doc_type] = deepcopy(payload)


def extraction(label: str, relationship: str, message_fact: str):
    return {
        'memories': [{
            'stable_key': 'relationship:xaden:mara:trust', 'action': 'update', 'title': 'Trust state',
            'what_happened': label, 'who_remembers': ['Mara','Xaden'], 'who_involved': ['Mara','Xaden'],
            'memory_type': 'Relationship Moment', 'importance': 'Critical', 'emotional_weight': 'mixed',
            'current_relevance': 'active', 'confidence': .99,
        }],
        'world_bible_updates': [
            {'section':'story_timeline','stable_key':'event:reroll-test','action':'update','item':{'event':message_fact},'confidence':.99,'conflicts_with_user_canon':False},
            {'section':'current_state','stable_key':'state:door','action':'update','item':{'status':message_fact},'confidence':.99,'conflicts_with_user_canon':False},
        ],
        'state_updates': {'relationship': {'status': relationship}},
        'conflicts': [],
    }


db = FakeDB(); cid = 1
# Earlier canonical turn creates the baseline memory/state.
apply_extracted_local(db, cid, extraction('They cautiously began trusting each other.', 'cautious trust', 'The west door is locked.'), source_message_id=100, user_text='baseline', assistant_text='baseline')
# First reroll variant becomes canonical and updates the same stable facts.
apply_extracted_local(db, cid, extraction('Xaden said he trusted Mara completely.', 'full trust', 'The west door is open.'), source_message_id=201, user_text='reroll A', assistant_text='reroll A')

# Also add an auto-only fact created solely by abandoned reroll A and a protected manual fact.
apply_extracted_local(db, cid, {
    'memories': [{'stable_key':'promise:abandoned','title':'Abandoned promise','what_happened':'Xaden promised to leave tonight.','who_remembers':['Xaden','Mara'],'who_involved':['Xaden','Mara'],'memory_type':'Promise','importance':'Important','current_relevance':'active','confidence':.99}],
    'world_bible_updates': [{'section':'custom_entries','stable_key':'branch:abandoned','item':{'details':'They planned to flee tonight.'},'confidence':.99}],
    'state_updates': {}, 'conflicts': []
}, source_message_id=201, user_text='reroll A', assistant_text='reroll A')
db.state['world_bible']['custom_entries'].append({'stable_key':'manual:keep','details':'Manual Bible fact must survive.','source':'user-created'})

rolled = rollback_auto_continuity_for_messages(db, cid, [201])
assert rolled['memories'] >= 2, rolled

# Evolving baseline must be restored, not deleted.
trust = next(m for m in db.memories if m.get('stable_key') == 'relationship:xaden:mara:trust')
assert 'cautiously began trusting' in trust['what_happened'], trust
assert str(trust.get('source_message_id')) == '100', trust
assert not any(m.get('stable_key') == 'promise:abandoned' for m in db.memories), db.memories

world = db.state['world_bible']
assert any(x.get('stable_key') == 'manual:keep' for x in world['custom_entries']), world
assert not any(x.get('stable_key') == 'branch:abandoned' for x in world['custom_entries']), world
assert any('locked' in str(x).lower() for x in world['story_timeline']), world['story_timeline']
assert db.state['relationship']['manual_status'] == 'user-authored'
assert db.state['relationship']['auto_current']['status'] == 'cautious trust', db.state['relationship']
assert str(db.state['relationship']['auto_current'].get('source_message_id')) == '100'

# New canonical reroll B can now establish different continuity cleanly.
apply_extracted_local(db, cid, extraction('Xaden said trust would have to be earned.', 'guarded', 'The west door remains locked.'), source_message_id=202, user_text='reroll B', assistant_text='reroll B')
trust = next(m for m in db.memories if m.get('stable_key') == 'relationship:xaden:mara:trust')
assert 'earned' in trust['what_happened']
assert str(trust.get('source_message_id')) == '202'
assert db.state['relationship']['auto_current']['status'] == 'guarded'
assert db.state['hard_canon']['manual_fact'] == 'User canon survives'

print('PASS — reroll canon switching rolls back only abandoned auto continuity')
