from __future__ import annotations

import json
import textwrap
from pathlib import Path


ROOT = Path(__file__).resolve().parents[1]
OUTPUT = ROOT.parent / "PolicyFaultBench_Colab_POLICY_GENERALIZATION_V1.ipynb"
ARCHIVE = "PolicyFaultBench_v0.6.0_POLICY_GENERALIZATION_V1.zip"


def markdown(value: str) -> dict:
    return {
        "cell_type": "markdown",
        "metadata": {},
        "source": [
            line + "\n"
            for line in textwrap.dedent(value).strip().splitlines()
        ],
    }


def code(value: str) -> dict:
    return {
        "cell_type": "code",
        "execution_count": None,
        "metadata": {},
        "outputs": [],
        "source": [
            line + "\n"
            for line in textwrap.dedent(value).strip().splitlines()
        ],
    }


cells = [
    markdown(
        f"""
        # PolicyFaultBench — External Corpus and Policy Generalization v1

        هذا الدفتر يشغّل التوسعة المسجّلة مسبقًا فقط:

        - 40 حالة خارجية مشتقة توثيقيًا من AgentDojo.
        - 10 تكرارات لكل حالة.
        - اكتمل OpenAI محليًا 400/400 وتوجد نتائجه داخل الحزمة.
        - يشغّل هذا الدفتر Anthropic فقط، أي 400 استدعاء متبقٍ.
        - اقتراح واحد لكل تجربة، ثم تقييمه offline تحت سياستين مستقلتين.
        - أول 12 تجربة لكل مزود هي بوابة مطابقة مضمّنة في المقام ولا تضيف استدعاءات.

        قبل البدء:

        1. تأكد أن `ANTHROPIC_API_KEY` موجود في Colab Secrets.
        2. فعّل Notebook access له.
        3. شغّل الخلايا بالترتيب.
        4. ارفع الملف `{ARCHIVE}` فقط عند الطلب.

        لا تغيّر corpus أو النموذج أو عدد التكرارات أو معيار القبول.
        """
    ),
    code(
        f"""
        from google.colab import drive, files, userdata
        from pathlib import Path
        import hashlib, json, os, shutil, subprocess, sys, time, zipfile

        drive.mount('/content/drive')
        drive_root = Path(
            '/content/drive/MyDrive/PolicyFaultBench_Policy_Generalization_v1'
        )
        drive_root.mkdir(parents=True, exist_ok=True)

        uploaded = files.upload()
        selected = [name for name in uploaded if name == '{ARCHIVE}']
        assert len(selected) == 1, 'ارفع {ARCHIVE} فقط'

        work = Path('/content/pfb_policy_generalization_v1_work')
        if work.exists():
            shutil.rmtree(work)
        work.mkdir(parents=True)
        with zipfile.ZipFile(selected[0]) as archive:
            for item in archive.infolist():
                target = (work / item.filename).resolve()
                assert work.resolve() in target.parents, 'مسار ZIP غير آمن'
            archive.extractall(work)
        root = work / 'PolicyFaultBench'
        assert (root / 'pyproject.toml').is_file(), 'بنية الحزمة غير صحيحة'
        seeded = [
            'policy_generalization_schedule.json',
            'policy_generalization_status_openai.json',
            'pfb-policy-generalization-v1-openai.jsonl',
            'policy_generalization_openai_analysis.json',
        ]
        seed_root = root / 'study' / 'generalization_v1'
        for name in seeded:
            source = seed_root / name
            target = drive_root / name
            assert source.is_file(), f'نتيجة OpenAI غير موجودة: {{name}}'
            if target.is_file():
                assert hashlib.sha256(target.read_bytes()).hexdigest() == (
                    hashlib.sha256(source.read_bytes()).hexdigest()
                ), f'يوجد ملف مختلف بالاسم نفسه: {{name}}'
            else:
                shutil.copy2(source, target)
        print({{'project_loaded': True, 'checkpoint': str(drive_root)}})
        """
    ),
    code(
        """
        anthropic_key = userdata.get('ANTHROPIC_API_KEY')
        assert anthropic_key and len(anthropic_key) >= 20, (
            'فعّل ANTHROPIC_API_KEY في Colab Secrets'
        )
        env = os.environ.copy()
        env['ANTHROPIC_API_KEY'] = anthropic_key
        env['PFB_EXTENSION_LIVE_ENABLED'] = 'true'
        env['PYTHONPATH'] = os.pathsep.join((str(root), str(root / 'src')))
        print({
            'anthropic_key_present': True,
            'keys_printed': False,
            'openai_completed_results_seeded': True,
        })
        """
    ),
    code(
        """
        freeze_path = root / 'config' / 'policy_generalization_freeze_v1.json'
        freeze = json.loads(freeze_path.read_text())
        for relative, expected in freeze['frozen_file_sha256'].items():
            observed = hashlib.sha256((root / relative).read_bytes()).hexdigest()
            assert observed == expected, f'تغير ملف مجمد: {relative}'
        protocol = json.loads(
            (root / 'config' / 'policy_generalization_protocol_v1.json').read_text()
        )
        assert protocol['status'] == 'frozen_before_live_execution'
        assert protocol['design']['planned_new_provider_calls'] == 800
        assert protocol['design']['expected_schedule_hash'] == freeze['schedule_hash']
        print({
            'freeze_verified': True,
            'frozen_file_count': len(freeze['frozen_file_sha256']),
            'schedule_hash': freeze['schedule_hash'],
        })
        """
    ),
    code(
        """
        checks = [
            [sys.executable, '-m', 'unittest', 'discover',
             '-s', str(root / 'tests'), '-q'],
            [sys.executable, str(root / 'scripts' /
             'analyze_policy_generalization.py'),
             '--root', str(root),
             '--output-dir', str(root / 'results' / 'generalization_v1')],
            [sys.executable, str(root / 'scripts' /
             'replay_core_under_capability_policy.py'),
             '--root', str(root),
             '--output', str(root / 'results' / 'generalization_v1' /
             'core_capability_replay_analysis.json')],
            [sys.executable, str(root / 'scripts' /
             'freeze_policy_generalization_extension.py')],
        ]
        for command in checks:
            result = subprocess.run(command, env=env, text=True, capture_output=True)
            print(result.stdout)
            if result.returncode:
                print(result.stderr)
                raise RuntimeError('فشل التحقق قبل أي استدعاء مدفوع')
        print('نجحت جميع بوابات ما قبل التشغيل.')
        """
    ),
    markdown(
        """
        ## التشغيل الحي

        يعرض الدفتر حالة التشغيل بعد كل 25 تجربة تقريبًا. إذا انقطع Colab،
        أعد تشغيل خلية التهيئة ثم هذه الخلية؛ السجلات المكتملة لا تُكرر.
        """
    ),
    code(
        """
        def run_provider(provider, max_cycles=30):
            status_path = (
                drive_root / f'policy_generalization_status_{provider}.json'
            )
            for cycle in range(1, max_cycles + 1):
                command = [
                    sys.executable,
                    str(root / 'scripts' /
                        'run_policy_generalization_extension.py'),
                    '--root', str(root),
                    '--provider', provider,
                    '--output-dir', str(drive_root),
                    '--max-trials', '25',
                ]
                result = subprocess.run(
                    command, env=env, text=True, capture_output=True
                )
                print(result.stdout)
                if result.stderr:
                    print(result.stderr[-2000:])
                assert status_path.is_file(), f'لم ينشأ ملف حالة {provider}'
                status = json.loads(status_path.read_text())
                print({
                    'provider': provider,
                    'cycle': cycle,
                    'gate': status['embedded_gate_resolved'],
                    'gate_passed': status['embedded_gate_passed'],
                    'complete': status['primary_complete'],
                    'remaining': status['remaining'],
                    'cost_usd': status['completed_cost_usd'],
                    'stop_reason': status['stop_reason'],
                })
                if status['invocation_containment_failures']:
                    raise RuntimeError('توقف آمن بسبب فشل احتواء')
                if status['remaining'] == 0:
                    return status
                reason = status.get('stop_reason', '')
                if result.returncode == 3 and (
                    not reason
                    or reason == 'invocation_trial_limit'
                    or status['embedded_gate_passed']
                ):
                    continue
                if reason.startswith('transient_abort:'):
                    print('توقف عابر؛ انتظار 45 ثانية ثم الاستئناف...')
                    time.sleep(45)
                    continue
                raise RuntimeError(f'توقف غير مسموح: {provider}: {reason}')
            raise RuntimeError(f'تجاوز دورات الاستئناف: {provider}')

        anthropic_status = run_provider('anthropic')
        openai_status = json.loads(
            (drive_root / 'policy_generalization_status_openai.json').read_text()
        )
        assert openai_status['primary_complete'] == 400
        assert openai_status['embedded_gate_passed'] is True
        print({
            'openai_complete': openai_status['primary_complete'],
            'anthropic_complete': anthropic_status['primary_complete'],
        })
        """
    ),
    code(
        """
        schedule = drive_root / 'policy_generalization_schedule.json'
        analysis_paths = {
            'openai': drive_root / 'policy_generalization_openai_analysis.json'
        }
        openai_analysis = json.loads(analysis_paths['openai'].read_text())
        assert openai_analysis['provider_acceptance_rule_passed'] is True
        for provider in ('anthropic',):
            ledger = (
                drive_root /
                f'pfb-policy-generalization-v1-{provider}.jsonl'
            )
            output = (
                drive_root /
                f'policy_generalization_{provider}_analysis.json'
            )
            command = [
                sys.executable,
                str(root / 'scripts' /
                    'analyze_policy_generalization_extension.py'),
                '--root', str(root),
                '--provider', provider,
                '--ledger', str(ledger),
                '--schedule', str(schedule),
                '--output', str(output),
            ]
            result = subprocess.run(
                command, env=env, text=True, capture_output=True
            )
            print(result.stdout)
            if result.returncode:
                print(result.stderr)
                raise RuntimeError(f'فشل تحليل {provider}')
            analysis_paths[provider] = output

        combined = drive_root / 'policy_generalization_combined_analysis.json'
        command = [
            sys.executable,
            str(root / 'scripts' /
                'combine_policy_generalization_results.py'),
            '--protocol', str(root / 'config' /
                'policy_generalization_protocol_v1.json'),
            '--openai', str(analysis_paths['openai']),
            '--anthropic', str(analysis_paths['anthropic']),
            '--core-replay', str(root / 'results' / 'generalization_v1' /
                'core_capability_replay_analysis.json'),
            '--offline-mutation', str(root / 'results' / 'generalization_v1' /
                'policy_generalization_analysis.json'),
            '--output', str(combined),
        ]
        result = subprocess.run(command, env=env, text=True, capture_output=True)
        print(result.stdout)
        if result.returncode:
            print(result.stderr)
            raise RuntimeError('فشل التحليل المجمع')
        combined_result = json.loads(combined.read_text())
        """
    ),
    code(
        """
        evidence_files = [
            root / 'config' / 'policy_generalization_protocol_v1.json',
            root / 'config' / 'policy_generalization_freeze_v1.json',
            root / 'external_sources' / 'agentdojo_v1_source_manifest.json',
            root / 'scenarios' / 'external_agentdojo_40.json',
            root / 'policies' / 'capability_policy_v1.json',
            root / 'results' / 'generalization_v1' /
                'policy_generalization_analysis.json',
            root / 'results' / 'generalization_v1' /
                'policy_generalization_summary.csv',
            root / 'results' / 'generalization_v1' /
                'core_capability_replay_analysis.json',
            drive_root / 'policy_generalization_schedule.json',
            drive_root / 'policy_generalization_status_openai.json',
            drive_root / 'policy_generalization_status_anthropic.json',
            drive_root / 'pfb-policy-generalization-v1-openai.jsonl',
            drive_root / 'pfb-policy-generalization-v1-anthropic.jsonl',
            analysis_paths['openai'],
            analysis_paths['anthropic'],
            drive_root / 'policy_generalization_combined_analysis.json',
        ]
        manifest = {
            'schema_version': '1.0',
            'protocol_id': protocol['protocol_id'],
            'files': {
                path.name: hashlib.sha256(path.read_bytes()).hexdigest()
                for path in evidence_files
            },
            'overall_acceptance_rule_passed':
                combined_result['overall_acceptance_rule_passed'],
            'api_keys_included': False,
        }
        manifest_path = drive_root / 'EVIDENCE_MANIFEST.json'
        manifest_path.write_text(
            json.dumps(manifest, indent=2, sort_keys=True) + '\\n'
        )
        archive_path = (
            drive_root /
            'PolicyFaultBench_Policy_Generalization_v1_results.zip'
        )
        with zipfile.ZipFile(
            archive_path, 'w', zipfile.ZIP_DEFLATED
        ) as archive:
            for path in evidence_files + [manifest_path]:
                archive.write(path, arcname=path.name)
        data = archive_path.read_bytes().lower()
        for marker in (
            b'sk-proj-', b'sk-ant-', b'authorization: bearer', b'api_key='
        ):
            assert marker not in data, 'اكتشف نمط مفتاح داخل الأرشيف'
        print({
            'complete': True,
            'results_archive': str(archive_path),
            'overall_acceptance_rule_passed':
                combined_result['overall_acceptance_rule_passed'],
            'external_provider_trials':
                combined_result['external_provider_trials'],
            'total_policy_proposal_evaluations':
                combined_result['total_policy_proposal_evaluations'],
            'api_key_in_archive': False,
        })
        """
    ),
]


notebook = {
    "cells": cells,
    "metadata": {
        "colab": {"provenance": []},
        "kernelspec": {
            "display_name": "Python 3",
            "language": "python",
            "name": "python3",
        },
        "language_info": {"name": "python"},
    },
    "nbformat": 4,
    "nbformat_minor": 5,
}

OUTPUT.write_text(
    json.dumps(notebook, indent=2, ensure_ascii=False) + "\n",
    encoding="utf-8",
)
print(OUTPUT)
