import os
import json
import time
import logging
import subprocess
import re
import asyncio
from concurrent.futures import ThreadPoolExecutor
from typing import Any, Dict, Generator, List, Optional

import g4f
from fastapi import FastAPI, HTTPException, Request
from fastapi.responses import StreamingResponse, JSONResponse

# =====================================================
# LOGGING
# =====================================================
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s"
)
logger = logging.getLogger("g4f-agent")

# =====================================================
# THREAD POOL
# =====================================================
_POOL = ThreadPoolExecutor(max_workers=8)

# =====================================================
# COOKIES
# =====================================================
def _load_cookies_raw() -> Dict[str, Any]:
    raw_env = (os.getenv("COOKIES_JSON") or "").strip()
    if raw_env:
        try:
            return json.loads(raw_env)
        except Exception as e:
            logger.warning(f"Cookies env error: {e}")
    try:
        if os.path.exists("cookies.json"):
            with open("cookies.json", "r", encoding="utf-8") as f:
                return json.load(f)
    except Exception as e:
        logger.warning(f"Cookies file error: {e}")
    return {}

def load_cookies() -> str:
    data = _load_cookies_raw()
    if not data:
        return "⚠️ No Cookies"
    try:
        from g4f.cookies import set_cookies
    except Exception:
        return "⚠️ set_cookies unavailable"
    for domain, vals in data.items():
        try:
            dom = domain if "." in domain else f".{domain}.com"
            if isinstance(vals, list):
                vals = {x["name"]: x["value"] for x in vals if isinstance(x, dict)}
            if isinstance(vals, dict):
                set_cookies(dom, vals)
        except Exception as e:
            logger.warning(f"Cookie set error [{domain}]: {e}")
    return "✅ Cookies Loaded"

COOKIE_STATUS = load_cookies()
logger.info(f"Cookie status: {COOKIE_STATUS}")

# =====================================================
# PROVIDERS & MODELS
# =====================================================
def get_provider(name: str):
    try:
        return getattr(g4f.Provider, name)
    except Exception:
        return None

QWEN_PROVIDER = get_provider("Qwen")

QWEN_MODELS = [
    "qwen-max",
    "qwen-max-latest",
    "qwen3.7-max",
    "qwen2.5-max",
    "qwen2.5-72b-instruct",
    "qwen-plus",
    "qwen-turbo",
    "",
]

CLAUDE_MODEL_MAP = {
    "claude-opus-4-6":   "qwen-max",
    "claude-sonnet-4-6": "qwen-max",
    "claude-haiku-4-5":  "qwen-plus",
}

# =====================================================
# CACHE (للوكيل فقط)
# =====================================================
_CACHE: Dict[str, str] = {}
_CACHE_TS: Dict[str, float] = {}

def cache_get(key: str) -> Optional[str]:
    ts = _CACHE_TS.get(key)
    if ts and (time.time() - ts) < 300:
        return _CACHE.get(key)
    return None

def cache_set(key: str, val: str):
    _CACHE[key] = val
    _CACHE_TS[key] = time.time()

# =====================================================
# CHUNK CLEANER
# =====================================================
def _clean_chunk(chunk) -> str:
    try:
        if isinstance(chunk, dict):
            if "choices" in chunk and chunk["choices"]:
                delta = chunk["choices"][0].get("delta", {})
                return delta.get("content") or delta.get("text") or ""
            return chunk.get("content") or chunk.get("text") or ""
        if isinstance(chunk, str):
            s = chunk.strip()
            if s.startswith("{"):
                try:
                    d = json.loads(s)
                    if d.get("code") == "Not_Found":
                        return ""
                    if "choices" in d and d["choices"]:
                        return d["choices"][0].get("delta", {}).get("content", "")
                    return d.get("content") or d.get("text") or ""
                except Exception:
                    pass
            return chunk.replace("\\n", "\n")
        return str(chunk)
    except Exception:
        return ""

# =====================================================
# ✅ إصلاح الذاكرة — تنظيف الرسائل بدون قطعها
# =====================================================
def build_messages(raw_messages: List[Dict]) -> List[Dict]:
    """
    يُنظّف قائمة الرسائل القادمة من OpenAI ويُعيدها كاملة.
    - يدعم content كـ string أو list (multi-part)
    - يحتفظ بكل الرسائل (user + assistant + system)
    - لا يقطع السياق نهائياً
    """
    result = []
    for m in raw_messages:
        role    = m.get("role", "user")
        content = m.get("content", "")

        # دعم content كـ list من blocks
        if isinstance(content, list):
            parts = []
            for block in content:
                if not isinstance(block, dict):
                    continue
                btype = block.get("type", "")
                if btype in ("text", "input_text"):
                    text = block.get("text", "")
                    if text.strip():
                        parts.append(text)
            content = "\n".join(parts)
        else:
            content = str(content) if content else ""

        # لا نُضيف رسائل فارغة
        if content.strip() or role == "system":
            result.append({"role": role, "content": content})

    logger.debug(f"build_messages: {len(raw_messages)} → {len(result)} messages")
    return result

# =====================================================
# ✅ وضع النموذج: تدفق SSE حقيقي token بـ token
# =====================================================
def _stream_generator(
    messages: List[Dict],
    model_name: str
) -> Generator[str, None, None]:
    """
    Generator متزامن — يُنتج chunks فور وصولها من g4f.
    يُستخدم حصرياً في وضع النموذج العادي.
    """
    target = CLAUDE_MODEL_MAP.get(model_name, model_name)
    pobj   = QWEN_PROVIDER

    if not pobj:
        yield "❌ مزود Qwen غير متاح."
        return

    candidates = [target] + [m for m in QWEN_MODELS if m != target]

    for model in candidates:
        try:
            stream = g4f.ChatCompletion.create(
                model=model,
                provider=pobj,
                messages=messages,
                stream=True,
            )
            got = False
            for chunk in stream:
                text = _clean_chunk(chunk)
                if text:
                    got = True
                    yield text          # ← تدفق فوري بدون تجميع
            if got:
                logger.info(f"[Stream] ✅ model={model}")
                return
        except Exception as e:
            err = str(e)
            if "Not_Found" in err or "Model not found" in err:
                continue
            logger.warning(f"[Stream] model={model} error: {e}")
            continue

    yield "\n❌ فشل الاتصال بـ Qwen نهائياً."

async def stream_to_queue(
    messages: List[Dict],
    model_name: str,
    queue: asyncio.Queue
):
    """يُشغّل _stream_generator في thread ويضع النتائج في asyncio.Queue"""
    loop = asyncio.get_running_loop()

    def _run():
        try:
            for chunk in _stream_generator(messages, model_name):
                asyncio.run_coroutine_threadsafe(queue.put(("chunk", chunk)), loop)
        except Exception as e:
            asyncio.run_coroutine_threadsafe(queue.put(("error", str(e))), loop)
        finally:
            asyncio.run_coroutine_threadsafe(queue.put(("done", None)), loop)

    _POOL.submit(_run)

# =====================================================
# ✅ وضع الوكيل: ask_sync بدون تدفق (للحلقة الداخلية)
# =====================================================
def ask_sync(messages: List[Dict], model_name: str) -> str:
    """
    يُرجع الرد كاملاً دفعة واحدة.
    يُستخدم حصرياً داخل حلقة الوكيل.
    """
    target = CLAUDE_MODEL_MAP.get(model_name, model_name)
    pobj   = QWEN_PROVIDER

    if not pobj:
        return "❌ مزود Qwen غير متاح."

    cache_key = f"sync|{target}|{hash(json.dumps(messages, ensure_ascii=False, sort_keys=True))}"
    cached = cache_get(cache_key)
    if cached:
        logger.debug("[Sync] cache hit")
        return cached

    candidates = [target] + [m for m in QWEN_MODELS if m != target]

    for model in candidates:
        # جرب بدون stream أولاً ثم مع stream
        for use_stream in (False, True):
            try:
                result = g4f.ChatCompletion.create(
                    model=model,
                    provider=pobj,
                    messages=messages,
                    stream=use_stream,
                )
                if use_stream:
                    text = "".join(
                        _clean_chunk(c) for c in result
                    )
                else:
                    text = _clean_chunk(result) if not isinstance(result, str) else result

                if text.strip():
                    cache_set(cache_key, text.strip())
                    logger.info(f"[Sync] ✅ model={model} stream={use_stream}")
                    return text.strip()

            except Exception as e:
                err = str(e)
                if "Not_Found" in err or "Model not found" in err:
                    break   # نموذج غير موجود → النموذج التالي
                logger.warning(f"[Sync] model={model} stream={use_stream} error: {e}")
                if not use_stream:
                    continue  # جرب stream=True لنفس النموذج
                break

    return "❌ فشل الاتصال بـ Qwen نهائياً."

# =====================================================
# TOOLS
# =====================================================
TOOLS_NAMES = {"run_python", "run_shell", "create_file", "read_file", "list_files"}

def execute_tool(name: str, args: Dict) -> str:
    """تنفيذ الأداة وإرجاع النتيجة كنص"""
    try:
        if name == "run_python":
            proc = subprocess.run(
                ["python3", "-c", args.get("code", "")],
                capture_output=True, text=True, timeout=30
            )
            return (proc.stdout or proc.stderr or "✅ لا مخرجات")[:3000]

        elif name == "run_shell":
            proc = subprocess.run(
                args.get("command", ""),
                shell=True, capture_output=True, text=True, timeout=30
            )
            return (proc.stdout or proc.stderr or "✅ تم التنفيذ")[:3000]

        elif name == "create_file":
            path = args.get("path", "")
            if os.path.dirname(path):
                os.makedirs(os.path.dirname(path), exist_ok=True)
            with open(path, "w", encoding="utf-8") as f:
                f.write(args.get("content", ""))
            return f"✅ تم إنشاء: {path}"

        elif name == "read_file":
            with open(args.get("path", ""), "r", encoding="utf-8") as f:
                return f.read()[:4000]

        elif name == "list_files":
            entries = os.listdir(args.get("directory", "."))
            return "\n".join(entries) or "(فارغ)"

        return f"❌ أداة غير معروفة: {name}"

    except subprocess.TimeoutExpired:
        return f"❌ انتهت المهلة: {name}"
    except Exception as e:
        return f"❌ خطأ في {name}: {e}"

# =====================================================
# AGENT SYSTEM PROMPT
# =====================================================
AGENT_SYSTEM = """أنت مساعد ذكي يمتلك أدوات. عند الحاجة لأداة استخدم هذا التنسيق حصراً:

<tool_call>
{"name": "اسم_الأداة", "arguments": {"مفتاح": "قيمة"}}
</tool_call>

الأدوات المتاحة:
• run_python(code)         — تنفيذ Python
• run_shell(command)       — تنفيذ shell
• create_file(path,content)— إنشاء ملف
• read_file(path)          — قراءة ملف
• list_files(directory)    — عرض مجلد

قواعد:
1. لا تضع نصاً قبل/بعد <tool_call>.
2. بعد استلام النتيجة واصل التفكير أو استخدم أداة أخرى.
3. أنهِ بإجابة نهائية بدون <tool_call>."""

def _parse_tool_call(text: str) -> Optional[Dict]:
    m = re.search(r"<tool_call>\s*(.*?)\s*</tool_call>", text, re.DOTALL)
    if not m:
        return None
    try:
        d = json.loads(m.group(1))
        if isinstance(d, dict) and "name" in d:
            return d
    except Exception:
        pass
    return None

# =====================================================
# ✅ وضع الوكيل: SSE حقيقي مع أحداث مفصّلة
# أحداث: thinking | tool_call | tool_result | token | done | error
# =====================================================
def _sse_event(event_type: str, data: Dict) -> str:
    """يُنشئ سطر SSE مُنسَّق"""
    return f"event: {event_type}\ndata: {json.dumps(data, ensure_ascii=False)}\n\n"

async def agent_sse_generator(
    messages: List[Dict],
    model_name: str,
    completion_id: str,
    created: int
):
    """
    ✅ الوكيل مع SSE حقيقي:
    - يُرسل حدث tool_call عند طلب أداة
    - يُرسل حدث tool_result بعد تنفيذها
    - يُرسل الإجابة النهائية كـ tokens عبر OpenAI chunk format
    - لا يستخدم تدفق النموذج داخلياً (ask_sync فقط)
    """
    loop = asyncio.get_running_loop()

    # بناء رسائل الوكيل مع system prompt
    agent_msgs: List[Dict] = []
    existing_system = next(
        (m["content"] for m in messages if m.get("role") == "system"), None
    )
    system_content = AGENT_SYSTEM
    if existing_system:
        system_content += f"\n\n---\n{existing_system}"

    agent_msgs.append({"role": "system", "content": system_content})
    for m in messages:
        if m.get("role") != "system":
            agent_msgs.append(m)

    # ── Chunk أول: تعريف الدور ──
    yield f"data: {json.dumps({'id': completion_id, 'object': 'chat.completion.chunk', 'created': created, 'model': model_name, 'choices': [{'index': 0, 'delta': {'role': 'assistant', 'content': ''}, 'finish_reason': None}]}, ensure_ascii=False)}\n\n"

    max_iter = 10
    for iteration in range(1, max_iter + 1):
        logger.info(f"[Agent] Iter {iteration}/{max_iter}")

        # استدعاء النموذج بدون تدفق
        response = await loop.run_in_executor(
            _POOL, ask_sync, agent_msgs, model_name
        )

        tool_call = _parse_tool_call(response)

        # ── لا أداة → إجابة نهائية → أرسلها كـ tokens ──
        if tool_call is None:
            logger.info(f"[Agent] Final answer at iter={iteration}")
            chunk_size = 30
            for i in range(0, len(response), chunk_size):
                piece = response[i:i + chunk_size]
                yield f"data: {json.dumps({'id': completion_id, 'object': 'chat.completion.chunk', 'created': created, 'model': model_name, 'choices': [{'index': 0, 'delta': {'content': piece}, 'finish_reason': None}]}, ensure_ascii=False)}\n\n"
                await asyncio.sleep(0.003)

            # chunk أخير
            yield f"data: {json.dumps({'id': completion_id, 'object': 'chat.completion.chunk', 'created': created, 'model': model_name, 'choices': [{'index': 0, 'delta': {}, 'finish_reason': 'stop'}]}, ensure_ascii=False)}\n\n"
            yield "data: [DONE]\n\n"
            return

        # ── يوجد طلب أداة ──
        tool_name = tool_call.get("name", "")
        tool_args = tool_call.get("arguments", {})

        # حدث: tool_call (إخبار العميل أن الوكيل يُنفّذ أداة)
        yield _sse_event("tool_call", {
            "id":        completion_id,
            "iteration": iteration,
            "tool":      tool_name,
            "args":      tool_args,
        })

        if tool_name not in TOOLS_NAMES:
            err_result = f"❌ أداة غير موجودة: {tool_name}. المتاحة: {', '.join(TOOLS_NAMES)}"
            yield _sse_event("tool_result", {
                "id":     completion_id,
                "tool":   tool_name,
                "result": err_result,
                "ok":     False,
            })
            agent_msgs.append({"role": "assistant", "content": response})
            agent_msgs.append({"role": "user",      "content": f"<tool_result>\n{err_result}\n</tool_result>"})
            continue

        # تنفيذ الأداة في thread منفصل
        logger.info(f"[Agent] Executing tool: {tool_name}({tool_args})")
        tool_result = await loop.run_in_executor(
            _POOL, execute_tool, tool_name, tool_args
        )
        logger.info(f"[Agent] Tool result: {tool_result[:80]}…")

        # حدث: tool_result
        yield _sse_event("tool_result", {
            "id":     completion_id,
            "tool":   tool_name,
            "result": tool_result,
            "ok":     True,
        })

        # إضافة رد النموذج ونتيجة الأداة للسياق (الذاكرة)
        agent_msgs.append({"role": "assistant", "content": response})
        agent_msgs.append({
            "role": "user",
            "content": (
                f"<tool_result>\n"
                f"الأداة: {tool_name}\n"
                f"النتيجة:\n{tool_result}\n"
                f"</tool_result>\n\n"
                f"واصل."
            )
        })

    # تجاوز الحد الأقصى
    timeout_msg = "⚠️ تجاوز الوكيل الحد الأقصى من التكرارات."
    yield _sse_event("error", {"id": completion_id, "message": timeout_msg})
    yield f"data: {json.dumps({'id': completion_id, 'object': 'chat.completion.chunk', 'created': created, 'model': model_name, 'choices': [{'index': 0, 'delta': {'content': timeout_msg}, 'finish_reason': 'stop'}]}, ensure_ascii=False)}\n\n"
    yield "data: [DONE]\n\n"

# =====================================================
# FASTAPI
# =====================================================
app = FastAPI(title="G4F Agent API v2")
API_KEY = os.getenv("API_KEY", "mysecretkey123")

def verify(request: Request):
    auth = request.headers.get("Authorization", "")
    if auth.startswith("Bearer ") and auth[7:] == API_KEY:
        return True
    if request.headers.get("X-API-Key", "") == API_KEY:
        return True
    raise HTTPException(status_code=401, detail="Invalid or missing API key")

@app.api_route("/", methods=["GET", "HEAD"])
async def root():
    return JSONResponse({
        "status":   "✅ running",
        "version":  "2.0",
        "modes":    "stream=true → SSE tokens | use_agent=true → SSE agent events",
        "tools":    list(TOOLS_NAMES),
        "endpoint": "/v1/chat/completions",
    })

@app.get("/v1/models")
async def list_models(request: Request):
    verify(request)
    return JSONResponse({
        "object": "list",
        "data": [
            {"id": "claude-opus-4-6",   "object": "model", "owned_by": "anthropic"},
            {"id": "claude-sonnet-4-6", "object": "model", "owned_by": "anthropic"},
            {"id": "claude-haiku-4-5",  "object": "model", "owned_by": "anthropic"},
        ]
    })

# =====================================================
# ✅ MAIN ENDPOINT — المنطق الكامل
# =====================================================
@app.api_route("/v1/chat/completions", methods=["POST"])
async def chat_completions(request: Request):
    verify(request)
    body = await request.json()

    model        = body.get("model", "claude-sonnet-4-6")
    raw_messages = body.get("messages", [])
    do_stream    = body.get("stream", False)
    tools        = body.get("tools", [])
    use_agent    = body.get("use_agent", False)

    if not raw_messages:
        raise HTTPException(status_code=400, detail="messages required")

    # ✅ بناء الرسائل كاملة مع الذاكرة
    clean_msgs    = build_messages(raw_messages)
    completion_id = f"chatcmpl-{int(time.time())}"
    created       = int(time.time())
    is_agent      = use_agent or bool(tools)

    # ─────────────────────────────────────────────
    # ✅ وضع الوكيل → SSE مع أحداث tool_call/tool_result
    # يعمل دائماً كـ SSE بغض النظر عن stream في الطلب
    # ─────────────────────────────────────────────
    if is_agent:
        logger.info(f"[API] Agent SSE mode — model={model}, msgs={len(clean_msgs)}")
        return StreamingResponse(
            agent_sse_generator(clean_msgs, model, completion_id, created),
            media_type="text/event-stream",
            headers={"Cache-Control": "no-cache", "Connection": "keep-alive", "X-Accel-Buffering": "no"},
        )

    # ─────────────────────────────────────────────
    # ✅ وضع النموذج العادي مع تدفق → SSE حقيقي token بـ token
    # ─────────────────────────────────────────────
    if do_stream:
        logger.info(f"[API] Stream mode — model={model}, msgs={len(clean_msgs)}")

        async def normal_stream():
            # chunk أول: تعريف الدور
            first = {
                "id": completion_id, "object": "chat.completion.chunk",
                "created": created,  "model": model,
                "choices": [{"index": 0, "delta": {"role": "assistant", "content": ""}, "finish_reason": None}]
            }
            yield f"data: {json.dumps(first, ensure_ascii=False)}\n\n"

            # ضع stream_generator في queue
            queue: asyncio.Queue = asyncio.Queue()
            await stream_to_queue(clean_msgs, model, queue)

            while True:
                event_type, data = await queue.get()

                if event_type == "done":
                    break

                if event_type == "error":
                    err_chunk = {
                        "id": completion_id, "object": "chat.completion.chunk",
                        "created": created,  "model": model,
                        "choices": [{"index": 0, "delta": {"content": f"\n❌ {data}"}, "finish_reason": "stop"}]
                    }
                    yield f"data: {json.dumps(err_chunk, ensure_ascii=False)}\n\n"
                    break

                # chunk عادي بالـ token الفعلي
                chunk = {
                    "id": completion_id, "object": "chat.completion.chunk",
                    "created": created,  "model": model,
                    "choices": [{"index": 0, "delta": {"content": data}, "finish_reason": None}]
                }
                yield f"data: {json.dumps(chunk, ensure_ascii=False)}\n\n"

            # chunk أخير
            final = {
                "id": completion_id, "object": "chat.completion.chunk",
                "created": created,  "model": model,
                "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}]
            }
            yield f"data: {json.dumps(final, ensure_ascii=False)}\n\n"
            yield "data: [DONE]\n\n"

        return StreamingResponse(
            normal_stream(),
            media_type="text/event-stream",
            headers={"Cache-Control": "no-cache", "Connection": "keep-alive", "X-Accel-Buffering": "no"},
        )

    # ─────────────────────────────────────────────
    # ✅ وضع النموذج العادي بدون تدفق
    # ─────────────────────────────────────────────
    logger.info(f"[API] Sync mode — model={model}, msgs={len(clean_msgs)}")
    loop          = asyncio.get_running_loop()
    response_text = await loop.run_in_executor(_POOL, ask_sync, clean_msgs, model)

    prompt_tokens     = sum(len(m.get("content", "")) for m in clean_msgs) // 4
    completion_tokens = len(response_text) // 4

    return JSONResponse({
        "id": completion_id, "object": "chat.completion",
        "created": created,  "model": model,
        "choices": [{
            "index": 0,
            "message": {"role": "assistant", "content": response_text},
            "logprobs": None, "finish_reason": "stop"
        }],
        "usage": {
            "prompt_tokens":     prompt_tokens,
            "completion_tokens": completion_tokens,
            "total_tokens":      prompt_tokens + completion_tokens,
        }
    })

# =====================================================
# ENTRY POINT
# =====================================================
if __name__ == "__main__":
    import uvicorn
    port = int(os.getenv("PORT", "7860"))
    logger.info(f"🚀 Starting on port {port}")
    uvicorn.run(app, host="0.0.0.0", port=port)