# ==============================================================================
# ☁️ r2.py (V38.0 - GEM-Architect: Diagnostic Matrix & Guardian Persistence)
# ==============================================================================

import os
import traceback
import json
import time
import io
from datetime import datetime, timedelta
import asyncio
import boto3
from botocore.exceptions import NoCredentialsError, ClientError
from typing import List, Dict, Any, Optional, Union

# ==============================================================================
# ⚙️ التكوين والإعدادات
# ==============================================================================
R2_ACCOUNT_ID = os.getenv("R2_ACCOUNT_ID")
R2_ACCESS_KEY_ID = os.getenv("R2_ACCESS_KEY_ID")
R2_SECRET_ACCESS_KEY = os.getenv("R2_SECRET_ACCESS_KEY")
BUCKET_NAME = "trading"

# 🔴 الرصيد الافتراضي
INITIAL_CAPITAL = 10.0 

# 📁 مفاتيح الملفات الأساسية في R2
WHALE_LEARNING_PENDING_KEY = "learning_whale_pending_records.json"
WHALE_LEARNING_COMPLETED_KEY = "learning_whale_completed_records.json"
WHALE_LEARNING_CONFIG_KEY = "learning_whale_optimal_config.json"

PORTFOLIO_STATE_KEY = "portfolio_state.json"
SMART_PORTFOLIO_STATE_KEY = "smart_portfolio_state.json"

OPEN_TRADES_KEY = "open_trades.json"
CLOSED_TRADES_KEY = "closed_trades_history.json"

CANDIDATES_KEY = "Candidates.json"
CONTRACTS_DB_KEY = "contracts.json"
SYSTEM_LOGS_KEY = "system_logs.json"
LEARNING_DATA_KEY = "learning_data.json"
ANALYSIS_AUDIT_KEY = "analysis_audit_log.json"
DEEP_STEWARD_AUDIT_KEY = "DeepSteward_Audit_Log.json"

# ✅ مفتاح جديد خاص ببيانات تدريب الحوكمة (156 مؤشر + النتيجة)
GOVERNANCE_TRAINING_KEY = "datasets/governance_training_data.json"

# 📊 مفاتيح التحليلات والتشخيص (جديد)
DIAGNOSTIC_STATS_KEY = "analytics/model_diagnostic_matrix.json"
GUARDIAN_STATS_KEY = "analytics/guardian_performance_stats.json"

# 🧬 مفاتيح النظام السيبراني
STRATEGIC_DNA_KEY = "learning/strategic_dna_v2.json"

# 🆕 مفاتيح التدريب (Legacy Support)
TRAINING_PENDING_BATCH_KEY = "datasets/pending_training_batch.json"
MODEL_TITAN_KEY = "ml_models/layer2/Titan_XGB_V1.json"


class R2Service:
    def __init__(self):
        try:
            endpoint_url = f"https://{R2_ACCOUNT_ID}.r2.cloudflarestorage.com"
            self.s3_client = boto3.client(
                's3',
                endpoint_url=endpoint_url,
                aws_access_key_id=R2_ACCESS_KEY_ID,
                aws_secret_access_key=R2_SECRET_ACCESS_KEY,
            )
            self.lock_acquired = False
            self.BUCKET_NAME = BUCKET_NAME
            print(f"✅ [R2 V38.0] Service Loaded (Diagnostics Ready).")
            
        except Exception as e:
            raise RuntimeError(f"Failed to initialize S3 client: {e}")

    # ==============================================================================
    # 🔴 دالة التصفير الشامل (Updated & Non-Blocking)
    # ==============================================================================
    async def reset_all_stats_async(self):
        """تصفير المحفظة والسجلات التاريخية والعودة لنقطة الصفر"""
        try:
            print("🔄 [R2 Reset] بدء عملية التصفير...")
            
            initial_state = {
                "current_capital_usd": INITIAL_CAPITAL,
                "invested_capital_usd": 0.0,
                "allocated_capital_usd": 0.0,
                "initial_capital_usd": INITIAL_CAPITAL,
                "total_trades": 0,
                "winning_trades": 0,
                "losing_trades": 0,
                "total_profit_usd": 0.0,
                "total_loss_usd": 0.0,
                "win_rate": 0.0,
                "daily_net_pnl": 0.0,
                "is_trading_halted": False,
                "last_update": datetime.now().isoformat()
            }
            
            await self.save_portfolio_state_async(initial_state)
            
            smart_state = {
                "current_capital": INITIAL_CAPITAL,
                "allocated_capital_usd": 0.0,
                "session_start_balance": INITIAL_CAPITAL,
                "daily_net_pnl": 0.0,
                "is_trading_halted": False,
                "halt_reason": None,
                "last_session_reset": datetime.now().isoformat()
            }
            await self.upload_json_async(smart_state, SMART_PORTFOLIO_STATE_KEY)
            
            empty_list_json = json.dumps([], indent=2).encode('utf-8')
            empty_logs_json = json.dumps({"logs": []}, indent=2).encode('utf-8')
            
            def _reset_sync():
                self.s3_client.put_object(
                    Bucket=BUCKET_NAME, Key=CLOSED_TRADES_KEY, Body=empty_list_json, ContentType="application/json"
                )
                self.s3_client.put_object(
                    Bucket=BUCKET_NAME, Key=OPEN_TRADES_KEY, Body=empty_list_json, ContentType="application/json"
                )
                self.s3_client.put_object(
                    Bucket=BUCKET_NAME, Key=SYSTEM_LOGS_KEY, Body=empty_logs_json, ContentType="application/json"
                )
            
            await asyncio.to_thread(_reset_sync)
            
            # تصفير المصفوفات أيضاً عند الطلب الشامل
            await self.reset_diagnostic_stats_async()
            await self.reset_guardian_stats_async()
            
            print("✅ [R2 Reset] تم تصفير جميع البيانات والمحفظة بنجاح.")
            return True
            
        except Exception as e:
            print(f"❌ [R2 Reset Error] فشل التصفير: {e}")
            traceback.print_exc()
            return False

    # ==============================================================================
    # 🆕 دوال دعم الملفات العامة
    # ==============================================================================
    async def upload_json_async(self, data: Any, key: str):
        try:
            json_bytes = json.dumps(data, indent=2, ensure_ascii=False).encode('utf-8')
            await asyncio.to_thread(
                self.s3_client.put_object,
                Bucket=self.BUCKET_NAME, 
                Key=key, 
                Body=json_bytes, 
                ContentType="application/json"
            )
        except Exception as e:
            print(f"❌ [R2] Upload JSON failed for {key}: {e}")
            raise

    async def get_file_async(self, key: str) -> Optional[bytes]:
        try:
            response = await asyncio.to_thread(
                self.s3_client.get_object,
                Bucket=self.BUCKET_NAME, 
                Key=key
            )
            return response['Body'].read()
        except ClientError as e:
            if e.response['Error']['Code'] == 'NoSuchKey':
                return None
            raise

    async def upload_file_async(self, file_obj, key: str):
        try:
            await asyncio.to_thread(
                self.s3_client.upload_fileobj,
                file_obj,
                self.BUCKET_NAME,
                key
            )
        except Exception as e:
            print(f"❌ [R2] File upload failed for {key}: {e}")
            raise

    # ==============================================================================
    # 📊 Diagnostic & Guardian Stats (New Feature)
    # ==============================================================================
    async def get_diagnostic_stats_async(self) -> Dict[str, Any]:
        """استرجاع مصفوفة التشخيص"""
        try:
            data = await self.get_file_async(DIAGNOSTIC_STATS_KEY)
            defaults = {
                "Titan": {"wins": 0, "losses": 0, "pnl": 0.0},
                "Patterns": {"wins": 0, "losses": 0, "pnl": 0.0},
                "Oracle": {"wins": 0, "losses": 0, "pnl": 0.0},
                "Sniper": {"wins": 0, "losses": 0, "pnl": 0.0},
                "MonteCarlo_L": {"wins": 0, "losses": 0, "pnl": 0.0},
                "MonteCarlo_A": {"wins": 0, "losses": 0, "pnl": 0.0},
                "Governance": {"wins": 0, "losses": 0, "pnl": 0.0}
            }
            loaded = json.loads(data) if data else {}
            # دمج القيم الافتراضية لضمان وجود المفاتيح
            for k, v in defaults.items():
                if k not in loaded: loaded[k] = v
            return loaded
        except: return {}

    async def update_diagnostic_stats_async(self, updates: Dict[str, Dict[str, float]]):
        """تحديث إحصائيات النماذج بشكل آمن"""
        try:
            current = await self.get_diagnostic_stats_async()
            for model, metrics in updates.items():
                if model not in current: current[model] = {"wins": 0, "losses": 0, "pnl": 0.0}
                current[model]['wins'] += metrics.get('wins', 0)
                current[model]['losses'] += metrics.get('losses', 0)
                current[model]['pnl'] += metrics.get('pnl', 0.0)
            
            await self.upload_json_async(current, DIAGNOSTIC_STATS_KEY)
        except Exception as e:
            print(f"❌ [R2] Failed to update diagnostics: {e}")

    async def reset_diagnostic_stats_async(self):
        """تصفير مصفوفة التشخيص"""
        empty = {k: {"wins": 0, "losses": 0, "pnl": 0.0} for k in 
                 ["Titan", "Patterns", "Oracle", "Sniper", "MonteCarlo_L", "MonteCarlo_A", "Governance"]}
        await self.upload_json_async(empty, DIAGNOSTIC_STATS_KEY)
        print("📊 [R2] Diagnostic Matrix Reset.")

    # --- Guardian Stats ---
    async def get_guardian_stats_async(self) -> Dict[str, Any]:
        data = await self.get_file_async(GUARDIAN_STATS_KEY)
        defaults = {
            "hybrid": {"total": 0, "good": 0, "saved": 0.0, "missed": 0.0},
            "crash":  {"total": 0, "good": 0, "saved": 0.0, "missed": 0.0},
            "giveback": {"total": 0, "good": 0, "saved": 0.0, "missed": 0.0},
            "stagnation": {"total": 0, "good": 0, "saved": 0.0, "missed": 0.0}
        }
        loaded = json.loads(data) if data else {}
        for k, v in defaults.items():
            if k not in loaded: loaded[k] = v
        return loaded

    async def save_guardian_stats_async(self, stats: Dict[str, Any]):
        await self.upload_json_async(stats, GUARDIAN_STATS_KEY)
        
    async def reset_guardian_stats_async(self):
        """تصفير إحصائيات الحراس"""
        defaults = {
            "hybrid": {"total": 0, "good": 0, "saved": 0.0, "missed": 0.0},
            "crash":  {"total": 0, "good": 0, "saved": 0.0, "missed": 0.0},
            "giveback": {"total": 0, "good": 0, "saved": 0.0, "missed": 0.0},
            "stagnation": {"total": 0, "good": 0, "saved": 0.0, "missed": 0.0}
        }
        await self.upload_json_async(defaults, GUARDIAN_STATS_KEY)
        print("🛡️ [R2] Guardian Stats Reset.")

    # ==============================================================================
    # 🏛️ بيانات تدريب الحوكمة (Governance Training Data)
    # ==============================================================================
    async def append_governance_training_data(self, record: Dict[str, Any]):
        """
        حفظ سجل كامل لبيانات الحوكمة (المؤشرات + القرار + النتيجة)
        يستخدم هذا الملف لاحقاً لتدريب Governance Meta-Model.
        """
        try:
            # 1. جلب البيانات الحالية
            data = await self.get_file_async(GOVERNANCE_TRAINING_KEY)
            dataset = json.loads(data) if data else []
            
            # 2. إضافة السجل الجديد
            dataset.append(record)
            
            # 3. الحفاظ على حجم معقول (آخر 5000 صفقة)
            if len(dataset) > 5000: 
                dataset = dataset[-5000:]
            
            # 4. الحفظ
            await self.upload_json_async(dataset, GOVERNANCE_TRAINING_KEY)
            print(f"🏛️ [R2 Governance] تم حفظ بيانات التدريب لـ {record.get('symbol', 'UNKNOWN')}")
            
        except Exception as e:
            print(f"❌ [R2 Error] فشل حفظ بيانات تدريب الحوكمة: {e}")

    # ==============================================================================
    # 🧬 إدارة الاستراتيجية (Strategic DNA)
    # ==============================================================================
    async def load_strategy_dna_async(self) -> Dict[str, Any]:
        try:
            data = await self.get_file_async(STRATEGIC_DNA_KEY)
            return json.loads(data) if data else {}
        except: return {}

    async def save_strategy_dna_async(self, dna_data: Dict[str, Any]):
        await self.upload_json_async(dna_data, STRATEGIC_DNA_KEY)

    # ==============================================================================
    # 📊 إدارة المرشحين
    # ==============================================================================
    async def save_candidates_async(self, candidates):
        try:
            data = {"timestamp": datetime.now().isoformat(), "total_candidates": len(candidates), "candidates": candidates}
            await self.upload_json_async(data, CANDIDATES_KEY)
        except Exception: pass

    async def load_candidates_async(self):
        try:
            data_bytes = await self.get_file_async(CANDIDATES_KEY)
            if data_bytes: return json.loads(data_bytes).get('candidates', [])
            return []
        except Exception: return []

    # ==============================================================================
    # 📝 السجلات والتعلم
    # ==============================================================================
    async def save_system_logs_async(self, log_data):
        pass 

    async def append_deep_steward_audit(self, audit_record: Dict[str, Any]):
        try:
            data = await self.get_file_async(DEEP_STEWARD_AUDIT_KEY)
            history = json.loads(data) if data else []
            history.append(audit_record)
            if len(history) > 2000: history = history[-2000:]
            await self.upload_json_async(history, DEEP_STEWARD_AUDIT_KEY)
        except Exception as e:
            print(f"❌ [R2 Error] فشل حفظ التدقيق: {e}")

    # ==============================================================================
    # 💰 إدارة المحفظة
    # ==============================================================================
    async def get_portfolio_state_async(self):
        try:
            data_bytes = await self.get_file_async(PORTFOLIO_STATE_KEY)
            if data_bytes:
                state = json.loads(data_bytes)
                if 'losing_trades' not in state: state['losing_trades'] = 0
                if 'initial_capital_usd' not in state: state['initial_capital_usd'] = INITIAL_CAPITAL
                if 'allocated_capital_usd' not in state: state['allocated_capital_usd'] = 0.0
                if 'daily_net_pnl' not in state: state['daily_net_pnl'] = 0.0
                return state
            
            initial_state = {
                "current_capital_usd": INITIAL_CAPITAL,
                "allocated_capital_usd": 0.0,
                "invested_capital_usd": 0.0,
                "initial_capital_usd": INITIAL_CAPITAL,
                "total_trades": 0, "winning_trades": 0, "losing_trades": 0,
                "total_profit_usd": 0.0, "total_loss_usd": 0.0, "win_rate": 0.0,
                "daily_net_pnl": 0.0,
                "is_trading_halted": False,
                "last_update": datetime.now().isoformat()
            }
            await self.save_portfolio_state_async(initial_state)
            return initial_state
        except Exception: raise

    async def save_portfolio_state_async(self, state):
        state['last_update'] = datetime.now().isoformat()
        await self.upload_json_async(state, PORTFOLIO_STATE_KEY)

    # ==============================================================================
    # 📈 إدارة الصفقات
    # ==============================================================================
    async def get_open_trades_async(self):
        try:
            data = await self.get_file_async(OPEN_TRADES_KEY)
            return json.loads(data) if data else []
        except: return []
    
    async def save_open_trades_async(self, trades):
        await self.upload_json_async(trades, OPEN_TRADES_KEY)

    async def append_to_closed_trades_history(self, trade_data: Dict[str, Any]):
        try:
            data = await self.get_file_async(CLOSED_TRADES_KEY)
            history = json.loads(data) if data else []
            history.append(trade_data)
            if len(history) > 1000: history = history[-1000:]
            await self.upload_json_async(history, CLOSED_TRADES_KEY)
        except Exception: pass

    # ==============================================================================
    # 📜 قاعدة العقود
    # ==============================================================================
    async def load_contracts_db_async(self):
        try:
            data = await self.get_file_async(CONTRACTS_DB_KEY)
            return json.loads(data) if data else {}
        except: return {}
    
    async def save_contracts_db_async(self, data):
        await self.upload_json_async(data, CONTRACTS_DB_KEY)

    # ==============================================================================
    # 🐋 تعلم الحيتان (Legacy Support)
    # ==============================================================================
    async def save_whale_learning_record_async(self, record: Dict[str, Any]):
        try:
            data = await self.get_file_async(WHALE_LEARNING_PENDING_KEY)
            pending = json.loads(data) if data else []
            pending.append(record)
            await self.upload_json_async(pending, WHALE_LEARNING_PENDING_KEY)
        except: pass

    async def get_pending_whale_learning_records_async(self) -> List[Dict[str, Any]]:
        data = await self.get_file_async(WHALE_LEARNING_PENDING_KEY)
        return json.loads(data) if data else []

    async def update_completed_whale_learning_record_async(self, completed_record: Dict[str, Any]):
        try:
            rec_id = completed_record.get("record_id")
            if not rec_id: return

            comp_data = await self.get_file_async(WHALE_LEARNING_COMPLETED_KEY)
            completed = json.loads(comp_data) if comp_data else []
            completed.append(completed_record)
            if len(completed) > 5000: completed = completed[-5000:]
            await self.upload_json_async(completed, WHALE_LEARNING_COMPLETED_KEY)

            pend_data = await self.get_file_async(WHALE_LEARNING_PENDING_KEY)
            pending = json.loads(pend_data) if pend_data else []
            updated_pending = [r for r in pending if r.get("record_id") != rec_id]
            await self.upload_json_async(updated_pending, WHALE_LEARNING_PENDING_KEY)
        except: pass

    async def get_all_completed_whale_records_async(self) -> List[Dict[str, Any]]:
        data = await self.get_file_async(WHALE_LEARNING_COMPLETED_KEY)
        return json.loads(data) if data else []
            
    async def save_whale_learning_config_async(self, config: Dict[str, Any]):
        await self.upload_json_async(config, WHALE_LEARNING_CONFIG_KEY)
        
    async def load_whale_learning_config_async(self) -> Dict[str, Any]:
        data = await self.get_file_async(WHALE_LEARNING_CONFIG_KEY)
        return json.loads(data) if data else {}