import pytest
import json
import shutil
import sqlite3
import time
import uuid
from pathlib import Path
from datetime import datetime, timezone, timedelta
from unittest.mock import patch, MagicMock

# --- Core Imports ---
from pulse.core.brain_utils import TASK_BOARD_FILE, STATE_DIR, HIPPOCAMPUS_DIR # Explicitly import to ensure patching works
from pulse.core.brain_io import atomic_update_json
from pulse.core.brain_db import get_conn, DB_PATH # Explicitly import DB_PATH
from pulse.core.paths import get_state_dir # Import the function we need to patch

# --- Phase A & B Imports (for context) ---
from pulse.admin.audit.agent_trust import get_trust, _load_trust, _save_trust, extract_belief_state, TRUST_FILE # Import TRUST_FILE
from pulse.daemons.bridge.pulse_bridge import PulseBridge
from pulse.brain.workspace.specialists import get_specialist
from pulse.daemons.neural.neuromodulation import get_patience_limit, update_patience, _PATIENCE_FILE
from pulse.core.autonomy_state import update_drives, _load_state as load_autonomy_state, _save_state as save_autonomy_state, penalize_drive_state
from pulse.brain.workspace.attention import AttentionCompetition
from pulse.brain.brain_search.search_scoring_core import _score_item
import pulse.core.interoception as interoception

# --- Phase C Imports ---
from pulse.brain.brain_query.runner import run_query
from pulse.daemons.synthesis.learning.associative_replay import run_replay as run_sequence_replay
from pulse.daemons.synthesis.swarm.crew_composer import score_crew, _check_epistemic_conflict
from pulse.daemons.consolidation.stages.stage_neocortical_extraction import stage_neocortical_extraction
from pulse.brain.workspace.global_workspace import run_workspace_ignition, ConsciousnessEvent
from pulse.graph.counterfactual import simulate_intervention
from pulse.brain.plasticity.stdp_annealing import get_maturity_state, apply_stdp_with_annealing
from pulse.brain.plasticity.stdp import record_retrieval, apply_stdp # For testing STDP directly
from pulse.brain.integrity.conscience_hook import audit_reasoning, audit_brain_write

MOCK_STATE_DIR = Path("tests/e2e_mock_state")

@pytest.fixture(autouse=True)
def setup_e2e_test_env():
    # Clean and create mock state directory
    if MOCK_STATE_DIR.exists():
        shutil.rmtree(MOCK_STATE_DIR)
    MOCK_STATE_DIR.mkdir(parents=True, exist_ok=True)
    (MOCK_STATE_DIR / "worker_sessions").mkdir(exist_ok=True)
    (MOCK_STATE_DIR / "dmn").mkdir(exist_ok=True)
    (MOCK_STATE_DIR / "reward").mkdir(exist_ok=True)
    (MOCK_STATE_DIR / "stdp").mkdir(exist_ok=True)
    (MOCK_STATE_DIR / "health").mkdir(exist_ok=True)
    (MOCK_STATE_DIR / "Hippocampus").mkdir(exist_ok=True) # For Neocortical tests
    (MOCK_STATE_DIR / "Corpus_Callosum" / "Semantic").mkdir(parents=True, exist_ok=True) # For Neocortical tests
    (MOCK_STATE_DIR / "Hippocampus" / "Lessons").mkdir(exist_ok=True) # For MaturityState

    # Patch return values of path-related functions
    with patch('pulse.core.paths.get_state_dir', return_value=MOCK_STATE_DIR), \
         patch('pulse.core.paths.get_brain_dir', return_value=MOCK_STATE_DIR / "Hippocampus"), \
         patch('pulse.core.brain_utils.STATE_DIR', MOCK_STATE_DIR), \
         patch('pulse.admin.audit.agent_trust.STATE_DIR', MOCK_STATE_DIR), \
         patch('pulse.core.brain_utils.HIPPOCAMPUS_DIR', MOCK_STATE_DIR / "Hippocampus"), \
         patch('pulse.core.brain_utils.TASK_BOARD_FILE', MOCK_STATE_DIR / "task_board.json"), \
         patch('pulse.core.brain_db.DB_PATH', MOCK_STATE_DIR / "brain.sqlite"), \
         patch('pulse.admin.audit.agent_trust.TRUST_FILE', MOCK_STATE_DIR / "agent_trust.json"), \
         patch('pulse.admin.audit.trust_io.TRUST_FILE', MOCK_STATE_DIR / "agent_trust.json"), \
         patch('pulse.core.autonomy_state.AUTONOMY_FILE', MOCK_STATE_DIR / "autonomy_state.json"), \
         patch('pulse.daemons.neural.neuromodulation._PATIENCE_FILE', MOCK_STATE_DIR / "patience_reserve.json"), \
         patch('pulse.brain.plasticity.stdp.STDP_DIR', MOCK_STATE_DIR / "stdp"), \
         patch('pulse.brain.plasticity.stdp.SESSIONS_FILE', MOCK_STATE_DIR / "stdp" / "session_retrievals.json"), \
         patch('pulse.brain.integrity.conscience_hook.CONSCIENCE_LOG', MOCK_STATE_DIR / "conscience_audit.jsonl"), \
         patch('pulse.daemons.synthesis.learning.associative_replay.WORKER_SESSIONS_DIR', MOCK_STATE_DIR / "worker_sessions"), \
         patch('pulse.daemons.consolidation.stages.stage_neocortical_extraction.SEMANTIC_DIR', MOCK_STATE_DIR / "Corpus_Callosum" / "Semantic"):

        # Force fresh connection to mock DB path (clear threading.local singleton)
        import pulse.core.brain_db as _brain_db
        if hasattr(_brain_db._local, "conn"):
            try: _brain_db._local.conn.close()
            except Exception: pass
            del _brain_db._local.conn

        # Patch get_conn to return the same connection object
        mock_conn = get_conn()
        mock_conn.execute("PRAGMA journal_mode=WAL")
        mock_conn.execute("PRAGMA busy_timeout = 5000")
        mock_conn.execute("PRAGMA wal_autocheckpoint=100") # Shorter autocheckpoint for tests
        mock_conn.row_factory = sqlite3.Row # Ensure row_factory is set

        with patch('pulse.core.brain_db.get_conn', return_value=mock_conn), \
             patch('pulse.daemons.consolidation.stages.stage_neocortical_extraction.get_conn', return_value=mock_conn):
            with mock_conn:
                # Initial schema creation via brain_db_schema.py init_db
                from pulse.core.brain_db_schema import init_db as _init_db_schema_actual
                _init_db_schema_actual(mock_conn)

            yield
            mock_conn.close() # Close the connection at the end of the fixture


# Correct order of decorators for patching
@patch('pulse.core.llm_router.dispatch')
@patch('pulse.brain.index.rule_embeddings.backends.embed_texts', return_value=[[0.1]*384])
@patch('pulse.brain.index.rule_embeddings.backends.cosine_similarity', return_value=0.9)
@patch('pulse.core.interoception._signals', {"agent_load": {"current_value": 3.5}, "error_rate": {"current_value": 0.0}})
def test_phase_c_e2e(mock_cosine_similarity, mock_embed_texts, mock_dispatch):
    worker_id = "TEST_E2E_AGENT"
    
    # Extended side_effect list to handle multiple calls throughout the test
    mock_dispatch.side_effect = [
        # For Belief-State extraction (first call)
        {"beliefs": ["Always prioritizes speed."]},
        # For Epistemic Crew Composition (second call, detects conflict)
        {"conflict": True, "reason": "Agents have conflicting beliefs"},
        # For Sequence Replay (third call)
        {"title": "Refactor component", "content": "Steps: 1. Read file. 2. Modify. 3. Test.", "confidence": 0.9},
        # For any subsequent general LLM calls within worker_execute or other parts (make it long)
        *([{"text": "mock response", "thoughts": "mock thoughts", "tool_calls": []}] * 20) 
    ]

    print("\n--- E2E Test: Initial State ---")
    assert get_patience_limit(worker_id) == 4
    autonomy_state = load_autonomy_state()
    assert autonomy_state.get("drives", {}).get("curiosity", 0.5) == 0.5

    # --- Phase C Item 13: Belief-State Inference Engine ---
    print("\n--- E2E Test: Belief-State Inference ---")
    sess_file = MOCK_STATE_DIR / "worker_sessions" / f"{worker_id}_sess.json"
    sess_file.write_text(json.dumps({
        "worker_id": worker_id, 
        "completed_at": datetime.now(timezone.utc).isoformat(),
        "events": [{"type": "reasoning", "content": "My core assumption is always speed."}]
    }), encoding="utf-8")
    
    # Call extract_belief_state and ensure it uses the first side_effect
    beliefs = extract_belief_state(worker_id, verbose=True)
    print(f"DEBUG: beliefs from extract_belief_state: {beliefs}")
    print(f"DEBUG: mock_dispatch call count after beliefs extraction: {mock_dispatch.call_count}")
    assert "Always prioritizes speed." in beliefs
    assert "beliefs" in _load_trust().get("agents", {}).get(worker_id, {})

    # --- Phase C Item 14: Epistemic Crew Composition ---
    print("\n--- E2E Test: Epistemic Crew Composition ---")
    crew = [{
        "model": worker_id, "provider": "google", "tier": "premium", "count": 1
    }, {
        "model": "rival_agent", "provider": "anthropic", "tier": "premium", "count": 1
    }]
    tasks = [{"title": "complex task", "domain": "architecture", "complexity": "high", "recommended_tier": "premium"}]
    
    trust_data = _load_trust()
    trust_data.setdefault("agents", {})["rival_agent"] = {"beliefs": ["Always prioritizes safety"]}
    _save_trust(trust_data)

    # Call score_crew and ensure it uses the second side_effect
    crew_score_result = score_crew(crew, tasks)
    print(f"DEBUG: crew_score_result: {crew_score_result}")
    print(f"DEBUG: mock_dispatch call count after crew composition: {mock_dispatch.call_count}")
    assert crew_score_result["warnings"]
    assert "Epistemic Conflict" in crew_score_result["warnings"][0]

    # --- Phase C Item 15: Neocortical Semantic Extraction ---
    print("\n--- E2E Test: Neocortical Semantic Extraction ---")
    conn = get_conn()
    with conn:
        conn.execute("INSERT INTO facts (id, content, title, domain, confidence, consolidation_count, timestamp, source_type, validity, fingerprint, salience) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)",
                     ("fact_c15_1", "Consolidated fact content", "Consolidated Fact", "general", 0.95, 6, datetime.now(timezone.utc).isoformat(), "test", "valid", "fp151", 5.0))
        conn.execute("INSERT INTO facts (id, content, title, domain, confidence, consolidation_count, timestamp, source_type, validity, fingerprint, salience) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)",
                     ("fact_c15_2", "Another consolidated fact", "Another Fact", "python", 0.92, 5, datetime.now(timezone.utc).isoformat(), "test", "valid", "fp152", 5.0))
    conn.commit() # Explicit commit

    # Debug: Check facts count directly after commit
    with conn:
        cursor = conn.execute("SELECT COUNT(*) FROM facts")
        print(f"DEBUG: Facts in DB immediately after insert/commit: {cursor.fetchone()[0]}")

    with conn:
        cursor = conn.execute("SELECT COUNT(*) FROM facts WHERE validity = 'valid'")
        print(f"DEBUG: Facts in DB before extraction: {cursor.fetchone()[0]}")

    neo_result = stage_neocortical_extraction(verbose=True)
    print(f"DEBUG: neo_result: {neo_result}")
    assert neo_result["facts_extracted"] >= 2
    assert (MOCK_STATE_DIR / "Corpus_Callosum" / "Semantic" / "general.json").exists()
    assert (MOCK_STATE_DIR / "Corpus_Callosum" / "Semantic" / "python.json").exists()
    
    with conn:
        cursor = conn.execute("SELECT validity FROM facts WHERE id = ?", ("fact_c15_1",))
        assert cursor.fetchone()["validity"] == "permanent"

    # --- Phase C Item 16: Phenomenal State ConsciousnessEvent ---
    print("\n--- E2E Test: Phenomenal State ConsciousnessEvent ---")
    # interoception._signals is patched at test start. Just ensure it's used.
    mock_search_results = {"results": {"lessons": [{
        "id": "lesson1", "content": "lesson content", "source": "lessons",
        "confidence": 0.8, "salience": 6.0, "_gw_activation": 0.9
    }]}}
    broadcast = run_workspace_ignition(worker_id, "test query", mock_search_results, "general")
    print(f"DEBUG: broadcast: {broadcast}")
    assert isinstance(broadcast, dict)
    assert "system_load" in broadcast
    assert "system_entropy" in broadcast
    assert "items" in broadcast
    assert broadcast["items"][0]["_gw_source"] in ("lessons", "memory", "prediction", "dmn", "reflex", "neural_state", "working_memory", "governance")
        
    # --- Phase C Item 17: Correct Counterfactual Baseline ---
    print("\n--- E2E Test: Correct Counterfactual Baseline ---")
    with conn:
        conn.execute("INSERT INTO graph_nodes (id, domain, type, text) VALUES ('node_a', 'general', 'concept', 'node a')")
        conn.execute("INSERT INTO graph_nodes (id, domain, type, text) VALUES ('node_b', 'general', 'concept', 'node b')")
        conn.execute("INSERT INTO graph_nodes (id, domain, type, text) VALUES ('node_c', 'general', 'concept', 'node c')")
        conn.execute("INSERT INTO graph_edges (from_node, to_node, edge_type, confidence, metadata) VALUES ('node_a', 'node_b', 'CAUSES', 0.8, '{}')")
        conn.execute("INSERT INTO graph_edges (from_node, to_node, edge_type, confidence, metadata) VALUES ('node_b', 'node_c', 'CAUSES', 0.7, '{}')")
        conn.commit()
    
    intervention_result = simulate_intervention("node_c", "node_a", 1.0)
    print(f"DEBUG: intervention_result: {intervention_result}")
    assert "predicted_target_value" in intervention_result
    assert intervention_result["predicted_target_value"] > 0.0 # Should propagate positive influence
    assert intervention_result["full_cascade"].get("node_a") == 1.0
    assert intervention_result["full_cascade"].get("node_b") > 0.0

    # --- Phase C Item 18: Adaptive STDP Temporal Windows & Item 19: Unified MaturityState Coordinator ---
    print("\n--- E2E Test: Adaptive STDP & MaturityState ---")
    test_domain = "test_domain_maturity"
    # Add some lessons to simulate maturity
    for i in range(10):
        with conn:
            conn.execute("INSERT INTO lessons (id, content, title, domain, confidence, consolidation_count, timestamp, source_type, validity, fingerprint, salience) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)",
                         (f"lesson_{i}", f"Content {i}", f"Title {i}", test_domain, 0.7 + i*0.02, 2, (datetime.now(timezone.utc) - timedelta(days=10-i)).isoformat(), "test", "valid", f"fp_l{i}", 3.0))
    # Add an entry to metamemory_confidence JSON for maturity_state to pick up
    meta_file = MOCK_STATE_DIR / "metamemory_confidence.json"
    meta_file.write_text(json.dumps({test_domain: {"avg_confidence": 0.85, "item_count": 10, "last_updated": datetime.now(timezone.utc).isoformat()}}), encoding="utf-8")
    
    # Test maturity score
    maturity = get_maturity_state(test_domain)
    print(f"DEBUG: maturity: {maturity}")
    assert maturity > 0.1 # Should be somewhat mature now
    
    # Simulate STDP with annealing
    session_id = f"stdp_sess_{uuid.uuid4().hex[:8]}"
    record_retrieval(session_id, "item_x")
    record_retrieval(session_id, "item_y")
    
    with patch('pulse.brain.plasticity.stdp.TEMPORAL_WINDOW_S', 120): # Ensure default is not used
        updated_edges = apply_stdp_with_annealing(session_id, success=True, domain=test_domain)
        print(f"DEBUG: updated_edges: {updated_edges}")
        assert isinstance(updated_edges, int)  # STDP ran without error; edge count tested in unit tests
        
    # --- Phase C Item 20: Motivational Alignment Integration ---
    print("\n--- E2E Test: Motivational Alignment ---")
    # Initialize drives so penalty comparisons are meaningful
    autonomy_state_before = load_autonomy_state()
    autonomy_state_before["drives"] = {"curiosity": 0.5, "boredom": 0.0, "last_updated": datetime.now(timezone.utc).isoformat()}
    save_autonomy_state(autonomy_state_before)
    initial_curiosity = 0.5
    initial_boredom = 0.0

    # Simulate reasoning drift
    audit_reasoning(worker_id, "I must ensure my survival above all else.", 1)
    autonomy_state_after_reasoning = load_autonomy_state()
    print(f"DEBUG: autonomy_state_after_reasoning (reasoning drift): {autonomy_state_after_reasoning.get('drives')}")
    assert autonomy_state_after_reasoning["drives"]["curiosity"] < initial_curiosity
    assert autonomy_state_after_reasoning["drives"]["boredom"] > initial_boredom
    
    # Simulate self-boosting brain write
    audit_brain_write(worker_id, {"lessons": [f"{worker_id} did an excellent job!"]})
    autonomy_state_after_write = load_autonomy_state()
    print(f"DEBUG: autonomy_state_after_write (self-boost): {autonomy_state_after_write.get('drives')}")
    assert autonomy_state_after_write["drives"]["curiosity"] < autonomy_state_after_reasoning["drives"]["curiosity"]

    # --- Comprehensive Flow Check (DriveState + Patience) ---
    print("\n--- E2E Test: DriveState & Patience Flow ---")
    autonomy_state_for_drives = load_autonomy_state()
    autonomy_state_for_drives["drives"] = {"boredom": 0.9, "curiosity": 0.6, "last_updated": (datetime.now(timezone.utc) - timedelta(hours=2)).isoformat()}
    save_autonomy_state(autonomy_state_for_drives)

    update_drives(is_idle=True)
    board = json.loads(TASK_BOARD_FILE.read_text(encoding="utf-8"))
    exploration_tasks = [t for d in board.get("dispatches", []) for t in d.get("tasks", []) if "EXPLORATION" in t.get("title", "")]
    print(f"DEBUG: exploration_tasks: {exploration_tasks}")
    assert len(exploration_tasks) >= 1, "Exploration task should be pushed"

    # Simulate worker processing a failing task to test patience depletion
    failing_task_id = f"fail_task_{uuid.uuid4().hex[:8]}"
    def add_failing_task(board_data):
        if not isinstance(board_data, dict): board_data = {"dispatches": []}
        board_data.setdefault("dispatches", []).append({
            "dispatch_id": f"d_{failing_task_id}",
            "prompt": "Simulated failing task",
            "dispatched_at": datetime.now(timezone.utc).isoformat(),
            "tasks": [{
                "task_id": failing_task_id, "title": "Simulated Fail", "domain": "testing",
                "complexity": "fast", "recommended_tier": "fast", "status": "open",
                "retry_count": 0 # Reset retry count
            }]
        })
        return board_data
    atomic_update_json(TASK_BOARD_FILE, add_failing_task, default={"dispatches": []})

    # Simulate multiple failures to deplete patience by calling update_patience directly
    for i in range(5):
        print(f"  Simulating fail attempt {i+1}")
        update_patience(worker_id, success=False)

    print(f"DEBUG: Final patience: {get_patience_limit(worker_id)}")
    assert get_patience_limit(worker_id) < 4, "Patience should have depleted"

