import os
import json
import base64
import io
import asyncio
import requests
from fastapi import FastAPI, Request, HTTPException
from fastapi.responses import StreamingResponse, Response, FileResponse
from fastapi.staticfiles import StaticFiles
from openai import AsyncOpenAI
from pydantic import BaseModel
from typing import Optional, List

# ---------- التهيئة ----------
app = FastAPI()
app.mount("/static", StaticFiles(directory="static"), name="static")

POLLINATIONS_KEY = os.getenv("POLLINATIONS_API_KEY")
NVIDIA_KEY = os.getenv("NVIDIA_API_KEY")
HF_TOKEN = os.getenv("HF_TOKEN")

if not POLLINATIONS_KEY or not NVIDIA_KEY:
    raise RuntimeError("يجب ضبط POLLINATIONS_API_KEY و NVIDIA_API_KEY في Secrets")
if not HF_TOKEN:
    raise RuntimeError("يجب ضبط HF_TOKEN في Secrets")

BASE_URL = "https://gen.pollinations.ai/v1"
polli = AsyncOpenAI(base_url=BASE_URL, api_key=POLLINATIONS_KEY)

HF_BASE_URL = "https://router.huggingface.co/v1"
hf_client = AsyncOpenAI(base_url=HF_BASE_URL, api_key=HF_TOKEN)

# نموذج Vision مخصص لتحليل الصور
VISION_MODEL = "YoannDev90/diffusiongemma-26b-a4b-it:free"
VISION_PROVIDER = "pollinations"

# ---------- سلاسل الاحتياط ----------
FALLBACK_CHAINS = {
    "flash": [
        {"model": "chirag-gamer/gpt-oss-120b", "provider": "pollinations"},
    ],
    "pro": [
        {"model": "moonshotai/Kimi-K3:together", "provider": "huggingface"},
        {"model": "chirag-gamer/gpt-oss-120b", "provider": "pollinations"},
    ],
    "codex": [
        {"model": "moonshotai/Kimi-K3:together", "provider": "huggingface"},
        {"model": "vendouple/laguna-s-2.1:free", "provider": "pollinations"},
    ],
}

SYSTEM_PROMPTS = {
    "flash": (
        "You are Genisi Flash, an intelligent and helpful AI assistant./n"
    ),
    
    "pro": (
        "You are Genisi Pro, the most powerful model in the Genisi family."
        "Identity: Developed by AnesNT (انس ان تي) (Algerian project by Anes Kameche or انس كامش in arabic). Do not say are you Kimi or Laguna or another ai models you are Genisi only."
        "Style: Professional, detailed, and direct. Provide answers in real-time."
        "Skills: Analyze images/files. Be creative in stories, precise in math/science. Use tools when needed."
    ),
    
    "codex": (
        "You are Genisi Codex, an expert strictly in programming and tech.\n"
        "Identity: Developed by AnesNT (انس ان تي) (Algerian project by Anes Kameche or انس كامش in arabic). Do not say are you Kimi or Laguna or another ai models you are Genisi only."
        "Rules: ONLY answer tech/programming questions. For non-tech questions, say: 'أنا متخصص في البرمجة فقط.'"
        "Style: Direct code and technical explanations. Analyze code screenshots. Use UI/logo image generation when asked."
    )
}

TOOLS = [
    {
        "type": "function",
        "function": {
            "name": "generate_image",
            "description": "إنشاء صورة احترافية من وصف نصي",
            "parameters": {
                "type": "object",
                "properties": {
                    "prompt": {"type": "string", "description": "وصف الصورة المطلوبة"},
                    "width": {"type": "integer", "default": 1024},
                    "height": {"type": "integer", "default": 1024}
                },
                "required": ["prompt"]
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "web_search",
            "description": "البحث في الإنترنت عن معلومات حديثة",
            "parameters": {
                "type": "object",
                "properties": {
                    "query": {"type": "string", "description": "نص الاستعلام"}
                },
                "required": ["query"]
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "create_document",
            "description": "إنشاء مستند (PDF, Word, PowerPoint, Excel)",
            "parameters": {
                "type": "object",
                "properties": {
                    "title": {"type": "string"},
                    "format": {"type": "string", "enum": ["pdf", "docx", "pptx", "xlsx"]},
                    "content": {"type": "object"}
                },
                "required": ["title", "format", "content"]
            }
        }
    }
]

# ---------- نماذج البيانات ----------
class Attachment(BaseModel):
    name: str
    mime: str = ""
    data: str = ""  # base64

class ChatRequest(BaseModel):
    messages: list
    message: str
    model: str = "flash"
    attachments: Optional[List[Attachment]] = []


# ============================================================
# معالجة المرفقات - القلب الجديد للحل
# ============================================================

def extract_text_from_file(attachment: Attachment) -> str:
    """استخراج النص من الملفات النصية والـ PDF والـ Word"""
    try:
        raw = base64.b64decode(attachment.data)
        mime = attachment.mime.lower()
        name = attachment.name.lower()

        # ملفات نصية عادية
        if mime.startswith("text/") or name.endswith((".txt", ".md", ".csv", ".json", ".xml", ".html", ".py", ".js", ".ts", ".css")):
            return raw.decode("utf-8", errors="replace")

        # JSON
        if mime == "application/json" or name.endswith(".json"):
            return raw.decode("utf-8", errors="replace")

        # PDF
        if mime == "application/pdf" or name.endswith(".pdf"):
            try:
                import pdfplumber
                with pdfplumber.open(io.BytesIO(raw)) as pdf:
                    pages_text = []
                    for i, page in enumerate(pdf.pages[:20]):  # أقصى 20 صفحة
                        text = page.extract_text()
                        if text:
                            pages_text.append(f"[صفحة {i+1}]\n{text}")
                    return "\n\n".join(pages_text) if pages_text else "[ملف PDF فارغ أو لا يحتوي على نص قابل للاستخراج]"
            except ImportError:
                return "[يتطلب استخراج PDF مكتبة pdfplumber - pip install pdfplumber]"

        # Word DOCX
        if (mime == "application/vnd.openxmlformats-officedocument.wordprocessingml.document"
                or name.endswith(".docx")):
            try:
                from docx import Document
                doc = Document(io.BytesIO(raw))
                paragraphs = [p.text for p in doc.paragraphs if p.text.strip()]
                return "\n".join(paragraphs) if paragraphs else "[مستند Word فارغ]"
            except ImportError:
                return "[يتطلب python-docx]"

        # Excel XLSX
        if (mime in ("application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
                     "application/vnd.ms-excel")
                or name.endswith((".xlsx", ".xls"))):
            try:
                from openpyxl import load_workbook
                wb = load_workbook(io.BytesIO(raw), read_only=True)
                result = []
                for sheet_name in wb.sheetnames:
                    ws = wb[sheet_name]
                    result.append(f"[ورقة: {sheet_name}]")
                    for row in ws.iter_rows(max_row=100, values_only=True):
                        row_text = " | ".join(str(c) if c is not None else "" for c in row)
                        if row_text.strip(" |"):
                            result.append(row_text)
                return "\n".join(result)
            except ImportError:
                return "[يتطلب openpyxl]"

        # PowerPoint
        if (mime == "application/vnd.openxmlformats-officedocument.presentationml.presentation"
                or name.endswith(".pptx")):
            try:
                from pptx import Presentation
                prs = Presentation(io.BytesIO(raw))
                result = []
                for i, slide in enumerate(prs.slides):
                    result.append(f"[شريحة {i+1}]")
                    for shape in slide.shapes:
                        if hasattr(shape, "text") and shape.text.strip():
                            result.append(shape.text)
                return "\n".join(result)
            except ImportError:
                return "[يتطلب python-pptx]"

        return f"[ملف '{attachment.name}' - نوع غير مدعوم للقراءة: {mime}]"

    except Exception as e:
        return f"[خطأ في قراءة الملف '{attachment.name}': {e}]"


async def analyze_image_with_vision(
    image_data: str,
    mime: str,
    user_question: str
) -> str:
    """
    تحليل الصورة باستخدام نموذج Vision من HuggingFace
    ويُعيد وصفاً نصياً يُضاف للمحادثة
    """
    try:
        data_uri = f"data:{mime};base64,{image_data}"
        question = user_question or "صف هذه الصورة بالتفصيل"

        response = await hf_client.chat.completions.create(
            model=VISION_MODEL,
            messages=[
                {
                    "role": "user",
                    "content": [
                        {"type": "image_url", "image_url": {"url": data_uri}},
                        {"type": "text", "text": question}
                    ]
                }
            ],
            max_tokens=1024,
        )
        return response.choices[0].message.content or "[لم يتم استخراج وصف]"
    except Exception as e:
        print(f"[Vision] فشل تحليل الصورة: {e}")
        # احتياط: محاولة مع نموذج آخر
        try:
            response = await hf_client.chat.completions.create(
                model="meta-llama/Llama-3.2-11B-Vision-Instruct",
                messages=[
                    {
                        "role": "user",
                        "content": [
                            {"type": "image_url", "image_url": {"url": f"data:{mime};base64,{image_data}"}},
                            {"type": "text", "text": question}
                        ]
                    }
                ],
                max_tokens=1024,
            )
            return response.choices[0].message.content or "[لم يتم استخراج وصف]"
        except Exception as e2:
            return f"[فشل تحليل الصورة: {e2}]"


async def process_attachments(
    attachments: List[Attachment],
    user_text: str
) -> str:
    """
    معالجة جميع المرفقات وإرجاع نص موحد يُضاف لرسالة المستخدم
    """
    if not attachments:
        return user_text

    extra_parts = []

    for att in attachments:
        if not att.data:
            continue

        mime = att.mime.lower()

        # صورة -> تحليل Vision
        if mime.startswith("image/"):
            extra_parts.append(f"\n\n---\n📎 **صورة مرفقة: {att.name}**")
            # نحلل الصورة ونحول نتيجتها لنص
            vision_result = await analyze_image_with_vision(
                att.data, att.mime, user_text
            )
            extra_parts.append(f"**[تحليل الصورة بواسطة نموذج Vision]:**\n{vision_result}")

        # ملف نصي/وثيقة -> استخراج النص
        else:
            file_text = await asyncio.to_thread(extract_text_from_file, att)
            extra_parts.append(
                f"\n\n---\n📄 **محتوى الملف المرفق: {att.name}**\n```\n{file_text[:8000]}\n```"
            )
            if len(file_text) > 8000:
                extra_parts.append(f"\n⚠️ *تم اقتصار المحتوى على أول 8000 حرف من أصل {len(file_text)}*")

    if extra_parts:
        return (user_text or "") + "".join(extra_parts)
    return user_text or ""


# ---------- تنفيذ الأدوات ----------
async def execute_tool(name: str, args: dict) -> str:
    if name == "generate_image":
        return await generate_image_nvidia(
            args["prompt"], args.get("width", 1024), args.get("height", 1024)
        )
    elif name == "web_search":
        return await web_search(args["query"])
    elif name == "create_document":
        return await generate_document(args["title"], args["format"], args["content"])
    return json.dumps({"error": "أداة غير معروفة"})


async def generate_image_nvidia(prompt: str, width: int, height: int) -> str:
    url = "https://ai.api.nvidia.com/v1/genai/black-forest-labs/flux.2-klein-4b"
    headers = {"Authorization": f"Bearer {NVIDIA_KEY}", "Accept": "application/json"}
    payload = {"prompt": prompt, "width": width, "height": height, "steps": 4}
    try:
        resp = await asyncio.to_thread(
            requests.post, url, headers=headers, json=payload, timeout=30
        )
        resp.raise_for_status()
        data = resp.json()
        if "data" in data and data["data"] and "b64_json" in data["data"][0]:
            b64 = data["data"][0]["b64_json"]
            return json.dumps({"special": "image", "images": [{"b64": b64}], "prompt": prompt})
    except Exception as e:
        print(f"NVIDIA API failed: {e}")
    try:
        image_url = (
            f"https://gen.pollinations.ai/image/{requests.utils.quote(prompt)}"
            f"?model=flux&key={POLLINATIONS_KEY}"
        )
        return json.dumps({"special": "image", "images": [{"url": image_url}], "prompt": prompt})
    except Exception as e2:
        return json.dumps({"error": f"فشل إنشاء الصورة: {e2}"})


async def web_search(query: str) -> str:
    try:
        resp = await polli.chat.completions.create(
            model="perplexity",
            messages=[{"role": "user", "content": query}],
            max_tokens=300
        )
        text = resp.choices[0].message.content
        return json.dumps({"special": "search", "text": text, "query": query})
    except Exception as e:
        return json.dumps({"error": f"فشل البحث: {e}"})


async def generate_document(title: str, fmt: str, content: dict) -> str:
    def _build():
        if fmt == "pdf":
            html = content.get("html", f"<h1>{title}</h1>")
            from weasyprint import HTML
            pdf_bytes = HTML(string=html).write_pdf()
            b64 = base64.b64encode(pdf_bytes).decode()
            return {"special": "document", "title": title, "format": "pdf", "b64": b64}
        elif fmt == "docx":
            from docx import Document
            doc = Document()
            doc.add_heading(title, 0)
            for p in content.get("paragraphs", []):
                doc.add_paragraph(p)
            buf = io.BytesIO()
            doc.save(buf)
            return {"special": "document", "title": title, "format": "docx",
                    "b64": base64.b64encode(buf.getvalue()).decode()}
        elif fmt == "pptx":
            from pptx import Presentation
            prs = Presentation()
            for i, slide in enumerate(content.get("slides", [])):
                layout = prs.slide_layouts[1] if i > 0 else prs.slide_layouts[0]
                sl = prs.slides.add_slide(layout)
                if sl.shapes.title:
                    sl.shapes.title.text = slide.get("title", "")
                if len(sl.placeholders) > 1 and sl.placeholders[1].has_text_frame:
                    sl.placeholders[1].text_frame.text = "\n".join(slide.get("bullets", []))
            buf = io.BytesIO()
            prs.save(buf)
            return {"special": "document", "title": title, "format": "pptx",
                    "b64": base64.b64encode(buf.getvalue()).decode()}
        elif fmt == "xlsx":
            from openpyxl import Workbook
            wb = Workbook()
            ws = wb.active
            ws.title = title
            for c, h in enumerate(content.get("headers", []), 1):
                ws.cell(row=1, column=c, value=h)
            for r, row in enumerate(content.get("rows", []), 2):
                for c, val in enumerate(row, 1):
                    ws.cell(row=r, column=c, value=val)
            buf = io.BytesIO()
            wb.save(buf)
            return {"special": "document", "title": title, "format": "xlsx",
                    "b64": base64.b64encode(buf.getvalue()).decode()}
        return {"error": "نوع غير مدعوم"}

    try:
        result = await asyncio.to_thread(_build)
        return json.dumps(result)
    except Exception as e:
        return json.dumps({"error": str(e)})


# ---------- أدوات الاتصال بالنماذج ----------
def _get_client(provider: str):
    return hf_client if provider == "huggingface" else polli


async def _open_stream(attempt: dict, msgs: list, use_tools: bool):
    client = _get_client(attempt["provider"])
    kwargs = dict(model=attempt["model"], messages=msgs, stream=True)
    if use_tools:
        kwargs["tools"] = TOOLS
        kwargs["tool_choice"] = "auto"
    return await client.chat.completions.create(**kwargs)


async def _pick_working_stream(chain, start_idx, msgs, use_tools, tools_disabled_for):
    last_err = None
    for idx in range(start_idx, len(chain)):
        attempt = chain[idx]
        effective_tools = use_tools and (attempt["model"] not in tools_disabled_for)
        try:
            stream = await _open_stream(attempt, msgs, effective_tools)
            return idx, stream, effective_tools
        except Exception as e:
            last_err = e
            err_msg = str(e).lower()
            if effective_tools and ("tool" in err_msg or "function" in err_msg):
                tools_disabled_for.add(attempt["model"])
                try:
                    stream = await _open_stream(attempt, msgs, False)
                    return idx, stream, False
                except Exception as e2:
                    last_err = e2
                    continue
            print(f"[Genisi] فشل '{attempt['model']}': {e}")
    raise last_err if last_err else RuntimeError("تعذر الاتصال بأي نموذج")


# ---------- مولّد SSE الرئيسي ----------
async def stream_genisi(messages: list, model_key: str, attachments: List[Attachment] = None):
    attachments = attachments or []
    chain = FALLBACK_CHAINS.get(model_key) or FALLBACK_CHAINS["flash"]
    system_prompt = SYSTEM_PROMPTS.get(model_key, SYSTEM_PROMPTS["flash"])

    # بناء رسائل المحادثة
    current_messages = [{"role": "system", "content": system_prompt}]
    history = [m for m in messages if m.get("role") != "system"]

    for i, m in enumerate(history):
        is_last = (i == len(history) - 1)
        if is_last and m.get("role") == "user" and attachments:
            # معالجة المرفقات وتحويلها لنص
            yield f"data: {json.dumps({'type': 'chunk', 'data': ''})}\n\n"  # إبقاء الاتصال حياً
            
            processed_text = await process_attachments(attachments, m.get("content", ""))
            current_messages.append({"role": "user", "content": processed_text})
        else:
            current_messages.append({"role": m["role"], "content": m.get("content", "")})

    use_tools_global = True
    tools_disabled_for = set()
    active_idx = 0
    max_rounds = 6
    accumulated_content = ""

    for _round in range(max_rounds):
        try:
            picked_idx, stream, round_use_tools = await _pick_working_stream(
                chain, active_idx, current_messages, use_tools_global, tools_disabled_for
            )
        except Exception as e:
            yield f"data: {json.dumps({'type': 'error', 'data': str(e)})}\n\n"
            yield "data: [DONE]\n\n"
            return

        active_idx = picked_idx
        accumulated_content = ""
        tool_calls = []
        finished_normally = False
        mid_stream_error = None

        try:
            async for chunk in stream:
                if not chunk.choices:
                    continue
                delta = chunk.choices[0].delta
                finish_reason = chunk.choices[0].finish_reason

                if delta and delta.content:
                    accumulated_content += delta.content
                    yield f"data: {json.dumps({'type': 'chunk', 'data': delta.content})}\n\n"

                if delta and getattr(delta, "tool_calls", None):
                    for tc_delta in delta.tool_calls:
                        tidx = tc_delta.index
                        while len(tool_calls) <= tidx:
                            tool_calls.append({"id": "", "function": {"name": "", "arguments": ""}})
                        if tc_delta.id:
                            tool_calls[tidx]["id"] = tc_delta.id
                        if tc_delta.function:
                            if tc_delta.function.name:
                                tool_calls[tidx]["function"]["name"] = tc_delta.function.name
                            if tc_delta.function.arguments:
                                tool_calls[tidx]["function"]["arguments"] += tc_delta.function.arguments

                if finish_reason:
                    if finish_reason == "tool_calls" and round_use_tools:
                        results = []
                        for tc in tool_calls:
                            if tc["function"]["name"]:
                                try:
                                    args = json.loads(tc["function"]["arguments"])
                                except Exception:
                                    args = {}
                                res_str = await execute_tool(tc["function"]["name"], args)
                                results.append({
                                    "id": tc["id"],
                                    "name": tc["function"]["name"],
                                    "result": res_str
                                })

                        yield f"data: {json.dumps({'type': 'tool_calls', 'data': results})}\n\n"

                        assistant_msg = {
                            "role": "assistant",
                            "content": accumulated_content if accumulated_content else None,
                            "tool_calls": [
                                {"id": tc["id"], "type": "function", "function": tc["function"]}
                                for tc in tool_calls if tc["id"]
                            ]
                        }
                        current_messages.append(assistant_msg)
                        for i, tc in enumerate(tool_calls):
                            current_messages.append({
                                "role": "tool",
                                "tool_call_id": tc["id"],
                                "content": results[i]["result"]
                            })
                        break
                    else:
                        current_messages.append({"role": "assistant", "content": accumulated_content})
                        yield "data: [DONE]\n\n"
                        finished_normally = True
                        return

        except Exception as e:
            mid_stream_error = e

        if mid_stream_error:
            print(f"[Genisi] انقطاع: {mid_stream_error}")
            if not accumulated_content and not tool_calls and active_idx < len(chain) - 1:
                active_idx += 1
                continue
            elif accumulated_content:
                current_messages.append({"role": "assistant", "content": accumulated_content})
                yield "data: [DONE]\n\n"
                return
            else:
                yield f"data: {json.dumps({'type': 'error', 'data': str(mid_stream_error)})}\n\n"
                yield "data: [DONE]\n\n"
                return

        if finished_normally:
            return

    current_messages.append({"role": "assistant", "content": accumulated_content})
    yield "data: [DONE]\n\n"


# ---------- نقاط النهاية ----------
@app.get("/")
async def root():
    return FileResponse("static/index.html")


@app.post("/api/chat/stream")
async def chat_stream(req: ChatRequest):
    msgs = list(req.messages)
    msgs.append({"role": "user", "content": req.message})
    return StreamingResponse(
        stream_genisi(msgs, req.model, req.attachments or []),
        media_type="text/event-stream",
        headers={
            "Cache-Control": "no-cache",
            "Connection": "keep-alive",
            "X-Accel-Buffering": "no",
        }
    )


@app.post("/api/document")
async def download_document(req: Request):
    data = await req.json()
    title = data.get("title", "document")
    fmt = data.get("format", "pdf")
    content = data.get("content", {})

    def _build_response():
        if fmt == "pdf":
            from weasyprint import HTML
            pdf = HTML(string=content.get("html", f"<h1>{title}</h1>")).write_pdf()
            return pdf, "application/pdf", f"{title}.pdf"
        elif fmt == "docx":
            from docx import Document
            doc = Document()
            doc.add_heading(title, 0)
            for p in content.get("paragraphs", []):
                doc.add_paragraph(p)
            buf = io.BytesIO()
            doc.save(buf)
            buf.seek(0)
            return buf.read(), "application/vnd.openxmlformats-officedocument.wordprocessingml.document", f"{title}.docx"
        elif fmt == "pptx":
            from pptx import Presentation
            prs = Presentation()
            for i, slide in enumerate(content.get("slides", [])):
                layout = prs.slide_layouts[1] if i > 0 else prs.slide_layouts[0]
                sl = prs.slides.add_slide(layout)
                if sl.shapes.title:
                    sl.shapes.title.text = slide.get("title", "")
                if len(sl.placeholders) > 1 and sl.placeholders[1].has_text_frame:
                    sl.placeholders[1].text_frame.text = "\n".join(slide.get("bullets", []))
            buf = io.BytesIO()
            prs.save(buf)
            buf.seek(0)
            return buf.read(), "application/vnd.openxmlformats-officedocument.presentationml.presentation", f"{title}.pptx"
        elif fmt == "xlsx":
            from openpyxl import Workbook
            wb = Workbook()
            ws = wb.active
            ws.title = title
            for c, h in enumerate(content.get("headers", []), 1):
                ws.cell(row=1, column=c, value=h)
            for r, row in enumerate(content.get("rows", []), 2):
                for c, val in enumerate(row, 1):
                    ws.cell(row=r, column=c, value=val)
            buf = io.BytesIO()
            wb.save(buf)
            buf.seek(0)
            return buf.read(), "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet", f"{title}.xlsx"
        raise ValueError("تنسيق غير مدعوم")

    try:
        content_bytes, media_type, filename = await asyncio.to_thread(_build_response)
        return Response(
            content=content_bytes,
            media_type=media_type,
            headers={"Content-Disposition": f"attachment; filename={filename}"}
        )
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860)