from fastapi import FastAPI
from fastapi.responses import StreamingResponse, JSONResponse
from pydantic import BaseModel
from typing import List, Optional
import time
import json
import os
import base64
import httpx
from dotenv import load_dotenv

load_dotenv()

# Dummy imports for example
from models import AVAILABLE_MODELS, MODEL_ALIASES  # Ensure these are defined properly

# Env variables
IMAGE_API_URL = os.environ["IMAGE_API_URL"]
SNAPZION_UPLOAD_URL = "https://upload.snapzion.com/api/public-upload"
SNAPZION_API_KEY = os.environ["SNAP"]

app = FastAPI()

def unix_id():
    return str(int(time.time() * 1000))

# === Models ===
@app.get("/v1/models")
async def list_models():
    return {"object": "list", "data": AVAILABLE_MODELS}

# === Chat Completion ===

class Message(BaseModel):
    role: str
    content: str

class ChatRequest(BaseModel):
    messages: List[Message]
    model: str
    stream: Optional[bool] = False

@app.post("/v1/chat/completions")
async def chat_completion(request: ChatRequest):
    model_id = MODEL_ALIASES.get(request.model, request.model)

    headers = {
        'accept': 'text/event-stream',
        'content-type': 'application/json',
        'origin': 'https://www.chatwithmono.xyz',
        'referer': 'https://www.chatwithmono.xyz/',
        'user-agent': 'Mozilla/5.0',
    }

    payload = {
        "messages": [{"role": msg.role, "content": msg.content} for msg in request.messages],
        "model": model_id
    }

    if request.stream:
        async def event_stream():
            chat_id = f"chatcmpl-{unix_id()}"
            created = int(time.time())
            sent_done = False

            async with httpx.AsyncClient(timeout=120) as client:
                async with client.stream("POST", "https://www.chatwithmono.xyz/api/chat", headers=headers, json=payload) as response:
                    async for line in response.aiter_lines():
                        if line.startswith("0:"):
                            try:
                                content_piece = json.loads(line[2:])
                                chunk_data = {
                                    "id": chat_id,
                                    "object": "chat.completion.chunk",
                                    "created": created,
                                    "model": model_id,
                                    "choices": [{
                                        "delta": {"content": content_piece},
                                        "index": 0,
                                        "finish_reason": None
                                    }]
                                }
                                yield f"data: {json.dumps(chunk_data)}\n\n"
                            except:
                                continue
                        elif line.startswith(("e:", "d:")) and not sent_done:
                            sent_done = True
                            done_chunk = {
                                "id": chat_id,
                                "object": "chat.completion.chunk",
                                "created": created,
                                "model": model_id,
                                "choices": [{
                                    "delta": {},
                                    "index": 0,
                                    "finish_reason": "stop"
                                }]
                            }
                            yield f"data: {json.dumps(done_chunk)}\n\ndata: [DONE]\n\n"
        return StreamingResponse(event_stream(), media_type="text/event-stream")
    else:
        assistant_response = ""
        usage_info = {}

        async with httpx.AsyncClient(timeout=120) as client:
            async with client.stream("POST", "https://www.chatwithmono.xyz/api/chat", headers=headers, json=payload) as response:
                async for chunk in response.aiter_lines():
                    if chunk.startswith("0:"):
                        try:
                            piece = json.loads(chunk[2:])
                            assistant_response += piece
                        except:
                            continue
                    elif chunk.startswith(("e:", "d:")):
                        try:
                            data = json.loads(chunk[2:])
                            usage_info = data.get("usage", {})
                        except:
                            continue

        return JSONResponse(content={
            "id": f"chatcmpl-{unix_id()}",
            "object": "chat.completion",
            "created": int(time.time()),
            "model": model_id,
            "choices": [{
                "index": 0,
                "message": {
                    "role": "assistant",
                    "content": assistant_response
                },
                "finish_reason": "stop"
            }],
            "usage": {
                "prompt_tokens": usage_info.get("promptTokens", 0),
                "completion_tokens": usage_info.get("completionTokens", 0),
                "total_tokens": usage_info.get("promptTokens", 0) + usage_info.get("completionTokens", 0),
            }
        })

# === Image Generation ===

class ImageGenerationRequest(BaseModel):
    prompt: str
    aspect_ratio: Optional[str] = "1:1"
    n: Optional[int] = 1
    user: Optional[str] = None
    model: Optional[str] = "default"

@app.post("/v1/images/generations")
async def generate_images(request: ImageGenerationRequest):
    results = []

    async with httpx.AsyncClient(timeout=60) as client:
        for _ in range(request.n):
            model = request.model or "default"

            if model in ["gpt-image-1", "dall-e-3", "dall-e-2", "nextlm-image-1"]:
                headers = {
                    'Content-Type': 'application/json',
                    'User-Agent': 'Mozilla/5.0',
                    'Referer': 'https://www.chatwithmono.xyz/',
                    'sec-ch-ua-platform': '"Windows"',
                    'sec-ch-ua': '"Not)A;Brand";v="8", "Chromium";v="138", "Google Chrome";v="138"',
                    'sec-ch-ua-mobile': '?0',
                }

                payload = {
                    "prompt": request.prompt,
                    "model": model
                }

                resp = await client.post("https://www.chatwithmono.xyz/api/image", headers=headers, json=payload)
                if resp.status_code != 200:
                    return JSONResponse(
                        status_code=502,
                        content={"error": f"{model} generation failed", "details": resp.text}
                    )

                data = resp.json()
                b64_image = data.get("image")
                if not b64_image:
                    return JSONResponse(status_code=502, content={"error": "Missing base64 image in response"})

                data_uri = f"data:image/png;base64,{b64_image}"

                upload_headers = {"Authorization": SNAPZION_API_KEY}
                upload_files = {
                    'file': ('image.png', base64.b64decode(b64_image), 'image/png')
                }

                upload_resp = await client.post(SNAPZION_UPLOAD_URL, headers=upload_headers, files=upload_files)
                if upload_resp.status_code != 200:
                    return JSONResponse(
                        status_code=502,
                        content={"error": "Upload failed", "details": upload_resp.text}
                    )

                upload_data = upload_resp.json()
                results.append({
                    "url": upload_data.get("url"),
                    "b64_json": b64_image,
                    "data_uri": data_uri,
                    "model": model
                })
            else:
                params = {
                    "prompt": request.prompt,
                    "aspect_ratio": request.aspect_ratio,
                    "link": "typegpt.net"
                }

                resp = await client.get(IMAGE_API_URL, params=params)
                if resp.status_code != 200:
                    return JSONResponse(
                        status_code=502,
                        content={"error": "Image generation failed", "details": resp.text}
                    )

                data = resp.json()
                results.append({
                    "url": data.get("image_link"),
                    "b64_json": data.get("base64_output"),
                    "retries": data.get("attempt"),
                    "model": "default"
                })

    return {
        "created": int(time.time()),
        "data": results
    }