# ============================================================
# DATETIME FIX — Must be first, before any google.auth import
# ============================================================
import datetime as _dt
import google.auth._helpers as _gah
_gah.utcnow = lambda: _dt.datetime.now(_dt.timezone.utc)
# ============================================================

import os
import json
import asyncio
import logging
import time
import uuid
import httpx

from fastapi import FastAPI, HTTPException, Depends, Request
from fastapi.responses import StreamingResponse, JSONResponse
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from typing import Optional, List, Union
import google.oauth2.credentials
import google.auth.transport.requests

# ── Logging ──────────────────────────────────────────────────
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s %(levelname)s %(message)s",
)
logger = logging.getLogger(__name__)

# ── Config from env ──────────────────────────────────────────
AUTH_PASSWORD   = os.environ.get("GEMINI_AUTH_PASSWORD", "")
RAW_CREDS       = os.environ.get("GEMINI_CREDENTIALS", "")
PORT            = int(os.environ.get("PORT", 7860))

GEMINI_API_BASE = "https://cloudcode-pa.googleapis.com/v1internal/projects/-/locations/-/endpoints/-"

MODELS = [
    "gemini-2.5-pro",
    "gemini-2.5-flash",
    "gemini-2.0-flash",
    "gemini-2.5-pro-search",
    "gemini-2.5-flash-search",
    "gemini-2.5-pro-nothinking",
    "gemini-2.5-flash-nothinking",
    "gemini-2.5-pro-maxthinking",
    "gemini-2.5-flash-maxthinking",
]

# Thinking budgets per model variant
THINKING_BUDGET = {
    "gemini-2.5-pro-nothinking":   0,
    "gemini-2.5-flash-nothinking": 0,
    "gemini-2.5-pro-maxthinking":  32768,
    "gemini-2.5-flash-maxthinking":32768,
}

# Search grounding models
SEARCH_MODELS = {"gemini-2.5-pro-search", "gemini-2.5-flash-search"}

# Base model mapping (strip suffix for API call)
def base_model(model: str) -> str:
    for suffix in ["-search", "-nothinking", "-maxthinking"]:
        if model.endswith(suffix):
            return model[: -len(suffix)]
    return model

# ── Credential management ─────────────────────────────────────
_creds: Optional[google.oauth2.credentials.Credentials] = None
_creds_lock = asyncio.Lock()

def _build_creds() -> google.oauth2.credentials.Credentials:
    if not RAW_CREDS:
        raise RuntimeError("GEMINI_CREDENTIALS env var not set")
    data = json.loads(RAW_CREDS)
    expiry = None
    if "expiry_date" in data:
        # expiry_date is epoch ms from oauth_creds.json
        ts = data["expiry_date"] / 1000.0
        expiry = _dt.datetime.fromtimestamp(ts, tz=_dt.timezone.utc)
    elif "expiry" in data:
        raw = data["expiry"]
        if isinstance(raw, (int, float)):
            expiry = _dt.datetime.fromtimestamp(raw, tz=_dt.timezone.utc)
        else:
            expiry = _dt.datetime.fromisoformat(raw)
            if expiry.tzinfo is None:
                expiry = expiry.replace(tzinfo=_dt.timezone.utc)

    c = google.oauth2.credentials.Credentials(
        token         = data.get("token") or data.get("access_token"),
        refresh_token = data.get("refresh_token"),
        token_uri     = data.get("token_uri", "https://oauth2.googleapis.com/token"),
        client_id     = data.get("client_id"),
        client_secret = data.get("client_secret"),
        scopes        = data.get("scopes", ["https://www.googleapis.com/auth/cloud-platform"]),
    )
    if expiry:
        c.expiry = expiry
    return c


def _refresh(c: google.oauth2.credentials.Credentials):
    """Synchronously refresh credentials if expired."""
    now = _dt.datetime.now(_dt.timezone.utc)

    # Safely check expiry, handle both aware and naive
    needs_refresh = False
    if c.token is None:
        needs_refresh = True
    elif c.expiry is not None:
        expiry = c.expiry
        if expiry.tzinfo is None:
            expiry = expiry.replace(tzinfo=_dt.timezone.utc)
        # refresh 5 minutes early
        needs_refresh = now >= (expiry - _dt.timedelta(minutes=5))
    
    if needs_refresh:
        logger.info("Refreshing Google OAuth token...")
        request = google.auth.transport.requests.Request()
        c.refresh(request)
        logger.info("Token refreshed successfully.")
    return c.token


async def _token() -> str:
    global _creds
    async with _creds_lock:
        if _creds is None:
            _creds = _build_creds()
        token = await asyncio.to_thread(_refresh, _creds)
    return token


# ── FastAPI app ───────────────────────────────────────────────
app = FastAPI(title="geminicli2api", version="1.0.0")

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)


# ── Auth dependency ───────────────────────────────────────────
async def verify_auth(request: Request):
    if not AUTH_PASSWORD:
        return
    auth = request.headers.get("Authorization", "")
    if auth.startswith("Bearer "):
        token = auth[7:]
    else:
        token = auth
    if token != AUTH_PASSWORD:
        raise HTTPException(status_code=401, detail="Unauthorized")


# ── Pydantic models ───────────────────────────────────────────
class Message(BaseModel):
    role: str
    content: Union[str, list]


class ChatRequest(BaseModel):
    model: str = "gemini-2.5-flash"
    messages: List[Message]
    stream: bool = False
    max_tokens: Optional[int] = None
    temperature: Optional[float] = None
    top_p: Optional[float] = None


# ── Conversion helpers ────────────────────────────────────────
def openai_messages_to_gemini(messages: List[Message]):
    """Convert OpenAI messages to Gemini contents format."""
    system_parts = []
    contents = []

    for msg in messages:
        role = msg.role
        content = msg.content

        if isinstance(content, str):
            parts = [{"text": content}]
        elif isinstance(content, list):
            parts = []
            for item in content:
                if isinstance(item, dict):
                    if item.get("type") == "text":
                        parts.append({"text": item["text"]})
                    elif item.get("type") == "image_url":
                        url = item["image_url"]["url"]
                        if url.startswith("data:"):
                            mime, b64 = url[5:].split(";base64,", 1)
                            parts.append({
                                "inlineData": {"mimeType": mime, "data": b64}
                            })
                        else:
                            parts.append({"text": f"[Image: {url}]"})
                else:
                    parts.append({"text": str(item)})
        else:
            parts = [{"text": str(content)}]

        if role == "system":
            system_parts.extend(parts)
        elif role == "user":
            contents.append({"role": "user", "parts": parts})
        elif role == "assistant":
            contents.append({"role": "model", "parts": parts})

    return system_parts, contents


def build_gemini_payload(req: ChatRequest) -> dict:
    system_parts, contents = openai_messages_to_gemini(req.messages)

    payload: dict = {"contents": contents}

    if system_parts:
        payload["systemInstruction"] = {"parts": system_parts}

    gen_config: dict = {}
    if req.max_tokens:
        gen_config["maxOutputTokens"] = req.max_tokens
    if req.temperature is not None:
        gen_config["temperature"] = req.temperature
    if req.top_p is not None:
        gen_config["topP"] = req.top_p

    model = req.model
    if model in THINKING_BUDGET:
        gen_config["thinkingConfig"] = {
            "thinkingBudget": THINKING_BUDGET[model],
            "includeThoughts": THINKING_BUDGET[model] > 0,
        }
    elif model not in {"gemini-2.0-flash"} and "flash" not in model:
        # Default thinking for pro models
        gen_config["thinkingConfig"] = {
            "thinkingBudget": -1,
            "includeThoughts": False,
        }

    if gen_config:
        payload["generationConfig"] = gen_config

    if model in SEARCH_MODELS:
        payload["tools"] = [{"googleSearch": {}}]

    return payload


def gemini_response_to_openai(gemini_resp: dict, model: str, stream: bool = False) -> dict:
    """Convert Gemini response to OpenAI format."""
    candidates = gemini_resp.get("candidates", [])
    text = ""
    finish_reason = "stop"

    if candidates:
        candidate = candidates[0]
        parts = candidate.get("content", {}).get("parts", [])
        for part in parts:
            if "text" in part and not part.get("thought", False):
                text += part["text"]
        fr = candidate.get("finishReason", "STOP")
        finish_reason = {
            "STOP": "stop",
            "MAX_TOKENS": "length",
            "SAFETY": "content_filter",
        }.get(fr, "stop")

    usage = gemini_resp.get("usageMetadata", {})
    prompt_tokens = usage.get("promptTokenCount", 0)
    completion_tokens = usage.get("candidatesTokenCount", 0)

    resp_id = f"chatcmpl-{uuid.uuid4().hex[:8]}"
    created = int(time.time())

    if stream:
        return {
            "id": resp_id,
            "object": "chat.completion.chunk",
            "created": created,
            "model": model,
            "choices": [{
                "index": 0,
                "delta": {"content": text},
                "finish_reason": finish_reason,
            }],
        }

    return {
        "id": resp_id,
        "object": "chat.completion",
        "created": created,
        "model": model,
        "choices": [{
            "index": 0,
            "message": {"role": "assistant", "content": text},
            "finish_reason": finish_reason,
        }],
        "usage": {
            "prompt_tokens": prompt_tokens,
            "completion_tokens": completion_tokens,
            "total_tokens": prompt_tokens + completion_tokens,
        },
    }


# ── Routes ────────────────────────────────────────────────────
@app.get("/")
async def root():
    return {"status": "ok", "models": MODELS}


@app.get("/v1/models")
async def list_models(_=Depends(verify_auth)):
    return {
        "object": "list",
        "data": [
            {
                "id": m,
                "object": "model",
                "created": 1700000000,
                "owned_by": "google",
            }
            for m in MODELS
        ],
    }


@app.post("/v1/chat/completions")
async def chat(req: ChatRequest, _=Depends(verify_auth)):
    tok = await _token()
    model = req.model
    api_model = base_model(model)
    payload = build_gemini_payload(req)

    headers = {
        "Authorization": f"Bearer {tok}",
        "Content-Type": "application/json",
    }

    if req.stream:
        url = f"{GEMINI_API_BASE}:streamGenerateContent?alt=sse&model={api_model}"

        async def generate():
            async with httpx.AsyncClient(timeout=120) as client:
                async with client.stream("POST", url, headers=headers, json=payload) as resp:
                    if resp.status_code != 200:
                        body = await resp.aread()
                        err = body.decode(errors="replace")
                        logger.error(f"Gemini API error {resp.status_code}: {err}")
                        yield f"data: {json.dumps({'error': err})}\n\n"
                        return

                    buffer = ""
                    async for chunk in resp.aiter_text():
                        buffer += chunk
                        while "\n\n" in buffer:
                            event, buffer = buffer.split("\n\n", 1)
                            for line in event.splitlines():
                                if line.startswith("data: "):
                                    data_str = line[6:]
                                    if data_str.strip() == "[DONE]":
                                        yield "data: [DONE]\n\n"
                                        return
                                    try:
                                        gemini_data = json.loads(data_str)
                                        openai_chunk = gemini_response_to_openai(
                                            gemini_data, model, stream=True
                                        )
                                        yield f"data: {json.dumps(openai_chunk)}\n\n"
                                    except json.JSONDecodeError:
                                        pass

                    yield "data: [DONE]\n\n"

        return StreamingResponse(generate(), media_type="text/event-stream")

    else:
        url = f"{GEMINI_API_BASE}:generateContent?model={api_model}"
        async with httpx.AsyncClient(timeout=120) as client:
            resp = await client.post(url, headers=headers, json=payload)

        if resp.status_code != 200:
            logger.error(f"Gemini API error {resp.status_code}: {resp.text}")
            raise HTTPException(status_code=resp.status_code, detail=resp.text)

        gemini_data = resp.json()
        return gemini_response_to_openai(gemini_data, model)


# ── Startup ───────────────────────────────────────────────────
@app.on_event("startup")
async def startup():
    print(f"\n===== Application Startup at {_dt.datetime.now().strftime('%Y-%m-%d %H:%M:%S')} =====\n")
    logger.info(f"Proxy ready — {len(MODELS)} models")


# ── Main ──────────────────────────────────────────────────────
if __name__ == "__main__":
    import uvicorn
    uvicorn.run("app:app", host="0.0.0.0", port=PORT, log_level="info")