from __future__ import annotations

import json
import os
from pathlib import Path
from typing import Any, Dict, List

import httpx
from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import FileResponse, JSONResponse
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel, Field

BASE_DIR = Path(__file__).resolve().parent
PROMPT_BUNDLE = BASE_DIR / "prompts" / "prompt_bundle.json"
INDEX_FILE = BASE_DIR / "index.html"

app = FastAPI(title="THE Z AI/AGENT", version="1.0.0")
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

app.mount("/static", StaticFiles(directory=str(BASE_DIR)), name="static")


class FilePayload(BaseModel):
    name: str
    content: str


class ChatRequest(BaseModel):
    messages: List[Dict[str, str]]
    tier: str = Field(default="NORMAL")
    mode: str = Field(default="chat")
    attachments: List[FilePayload] = Field(default_factory=list)
    lang: str = Field(default="ar")


def load_prompts() -> Dict[str, Any]:
    if not PROMPT_BUNDLE.exists():
        raise RuntimeError("prompt_bundle.json not found")
    return json.loads(PROMPT_BUNDLE.read_text(encoding="utf-8"))


def provider_config() -> Dict[str, str]:
    # List of Groq API keys to try in order
    groq_keys = [
        os.getenv("GROQ_API_KEY", "").strip(),
        "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",
        "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",
        "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",
    ]

    openrouter_key = os.getenv("OPENROUTER_API_KEY", "").strip()

    # Try each Groq key until one works
    for key in groq_keys:
        if key:
            return {
                "provider": "groq",
                "api_key": key,
                "base_url": os.getenv("GROQ_BASE_URL", "https://api.groq.com/openai/v1/chat/completions"),
                "model": os.getenv("GROQ_MODEL", "llama-3.3-70b-versatile"),
            }

    if openrouter_key:
        return {
            "provider": "openrouter",
            "api_key": openrouter_key,
            "base_url": os.getenv("OPENROUTER_BASE_URL", "https://openrouter.ai/api/v1/chat/completions"),
            "model": os.getenv("OPENROUTER_MODEL", "meta-llama/llama-3.1-70b-instruct:free"),
        }
    raise RuntimeError("Set GROQ_API_KEY or OPENROUTER_API_KEY")


def build_system_prompt(payload: ChatRequest) -> str:
    prompts = load_prompts()
    tier = str(payload.tier or "NORMAL").upper()
    mode = str(payload.mode or "chat").lower()
    tier_prompt = prompts.get("tiers", {}).get(tier, {}).get("prompt", "")
    system = prompts.get("system", "")
    coder = prompts.get("coder", "") if mode in {"coder", "agent"} else ""
    parts = [system]
    if tier_prompt:
        parts.append(f"[TIER {tier}]\n{tier_prompt}")
    if coder:
        parts.append(coder)
    return "\n\n".join(parts).strip()


def attachments_context(attachments: List[FilePayload]) -> str:
    if not attachments:
        return ""
    parts = ["[ATTACHED FILES]\nRead these files carefully before answering."]
    for idx, item in enumerate(attachments, start=1):
        parts.append(f"File {idx}: {item.name}\n{item.content}")
    return "\n\n".join(parts)


@app.get("/")
def home():
    return FileResponse(INDEX_FILE)


@app.get("/api/health")
def health():
    return {"ok": True}


@app.post("/api/chat")
async def chat(req: ChatRequest):
    try:
        cfg = provider_config()
    except RuntimeError as exc:
        raise HTTPException(status_code=500, detail=str(exc)) from exc

    system_prompt = build_system_prompt(req)
    extra = attachments_context(req.attachments)
    messages = list(req.messages)
    if extra:
        messages.append({"role": "system", "content": extra})

    body = {
        "model": cfg["model"],
        "messages": [{"role": "system", "content": system_prompt}] + messages,
        "temperature": 0.7,
        "max_tokens": 3500,
    }

    headers = {
        "Authorization": f"Bearer {cfg['api_key']}",
        "Content-Type": "application/json",
    }
    if cfg["provider"] == "openrouter":
        headers["HTTP-Referer"] = os.getenv("OPENROUTER_SITE_URL", "https://example.com")
        headers["X-Title"] = os.getenv("OPENROUTER_SITE_NAME", "THE Z AI/AGENT")

    async with httpx.AsyncClient(timeout=120) as client:
        try:
            response = await client.post(cfg["base_url"], headers=headers, json=body)
        except httpx.HTTPError as exc:
            raise HTTPException(status_code=502, detail=f"Upstream request failed: {exc}") from exc

    if response.status_code >= 400:
        raise HTTPException(status_code=502, detail=f"Upstream error {response.status_code}: {response.text[:500]}")

    data = response.json()
    try:
        reply = data["choices"][0]["message"]["content"]
    except Exception:
        reply = data.get("output_text") or data.get("response") or json.dumps(data, ensure_ascii=False)

    return JSONResponse({"reply": reply, "raw": data})


if __name__ == "__main__":
    import uvicorn

    uvicorn.run("app:app", host="0.0.0.0", port=int(os.getenv("PORT", "7860")), reload=False)