#!/usr/bin/env python3
"""
OpenAI-compatible API wrapper for v0.dev
Provides a drop-in replacement for OpenAI's API using v0.dev as the backend
"""

import os
import json
import asyncio
from typing import List, Dict, Any, Optional, AsyncGenerator
from datetime import datetime
import uuid

import httpx
from fastapi import FastAPI, HTTPException, Depends, Request
from fastapi.responses import StreamingResponse
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel, Field
from dotenv import load_dotenv

load_dotenv()

# Configuration
V0_API_KEY = os.getenv("V0_API_KEY", "v1:348jgN2Fu5eebFqZDMgEr0qm:u7FitMJwSu8Vi0AhUOgjlo7p")
V0_API_BASE_URL = os.getenv("V0_API_BASE_URL", "https://api.v0.dev/v1")
HOST = os.getenv("HOST", "0.0.0.0")
PORT = int(os.getenv("PORT", 8000))

# FastAPI app
app = FastAPI(
    title="v0.dev OpenAI Compatible API",
    description="Drop-in replacement for OpenAI API using v0.dev as backend",
    version="1.0.0"
)

# CORS middleware
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Pydantic models
class Message(BaseModel):
    role: str = Field(..., description="Role of the message sender (user, assistant, system)")
    content: str = Field(..., description="Content of the message")

class ChatCompletionRequest(BaseModel):
    model: str = Field(..., description="Model to use for completion")
    messages: List[Message] = Field(..., description="List of messages")
    max_tokens: Optional[int] = Field(None, description="Maximum tokens to generate")
    temperature: Optional[float] = Field(0.7, description="Sampling temperature")
    stream: Optional[bool] = Field(False, description="Whether to stream the response")
    project_id: Optional[str] = Field(None, description="v0.dev project ID")

class Choice(BaseModel):
    index: int
    message: Message
    finish_reason: str

class Usage(BaseModel):
    prompt_tokens: int
    completion_tokens: int
    total_tokens: int

class ChatCompletionResponse(BaseModel):
    id: str
    object: str = "chat.completion"
    created: int
    model: str
    choices: List[Choice]
    usage: Usage

class ChatCompletionStreamResponse(BaseModel):
    id: str
    object: str = "chat.completion.chunk"
    created: int
    model: str
    choices: List[Dict[str, Any]]

# v0.dev API client
class V0APIClient:
    def __init__(self, api_key: str, base_url: str):
        self.api_key = api_key
        self.base_url = base_url
        self.headers = {
            "Authorization": f"Bearer {api_key}",
            "Content-Type": "application/json"
        }

    async def create_chat(self, messages: List[Message], model_config: Dict[str, Any], project_id: Optional[str] = None) -> Dict[str, Any]:
        """Create a new chat with v0.dev"""
        url = f"{self.base_url}/chats"
        
        # Use the formatting function to properly format messages
        formatted_messages = [
            {"role": msg.role, "content": msg.content}
            for msg in messages
        ]
        
        # Extract system message and user messages using the formatting function
        system_message = ""
        user_messages = []
        
        for msg in messages:
            if msg.role == "system":
                system_message = msg.content
            else:
                user_messages.append(msg)
        
        # Use the last user message
        if not user_messages:
            raise HTTPException(status_code=400, detail="No user message found")
        
        last_user_message = user_messages[-1].content
        
        payload = {
            "system": system_message,
            "message": last_user_message,
            "modelConfiguration": model_config,
            "projectId": project_id
        }
        
        # Remove None values
        payload = {k: v for k, v in payload.items() if v is not None}
        
        async with httpx.AsyncClient() as client:
            response = await client.post(url, headers=self.headers, json=payload)
            response.raise_for_status()
            return response.json()

# Initialize v0 client
v0_client = V0APIClient(V0_API_KEY, V0_API_BASE_URL)

# Helper functions
def format_prompt(messages: List[Dict[str, Any]], add_special_tokens: bool = False,
                 do_continue: bool = False, include_system: bool = True) -> str:
    """
    Format a series of messages into a single string, optionally adding special tokens.

    Args:
        messages: A list of message dictionaries, each containing 'role' and 'content'.
        add_special_tokens: Whether to add special formatting tokens.
        do_continue: If True, don't add the final "Assistant:" prompt.
        include_system: Whether to include system messages in the formatted output.

    Returns:
        A formatted string containing all messages.
    """
    # Helper function to convert content to string
    def to_string(value) -> str:
        if isinstance(value, str):
            return value
        elif isinstance(value, dict):
            if "text" in value:
                return value.get("text", "")
            return ""
        elif isinstance(value, list):
            return "".join([to_string(v) for v in value])
        return str(value)

    # If there's only one message and no special tokens needed, just return its content
    if not add_special_tokens and len(messages) <= 1:
        return to_string(messages[0]["content"])

    # Filter and process messages
    processed_messages = [
        (message["role"], to_string(message["content"]))
        for message in messages
        if include_system or message.get("role") != "system"
    ]

    # Format each message as "Role: Content"
    formatted = "\n".join([
        f'{role.capitalize()}: {content}'
        for role, content in processed_messages
        if content.strip()
    ])

    # Add final prompt for assistant if needed
    if do_continue:
        return formatted

    return f"{formatted}\nAssistant:"

def create_openai_response(v0_response: Dict[str, Any], model: str) -> ChatCompletionResponse:
    """Convert v0.dev response to OpenAI format"""
    messages = v0_response.get("messages", [])
    assistant_message = None
    
    for msg in messages:
        if msg.get("role") == "assistant":
            assistant_message = msg
            break
    
    if not assistant_message:
        raise HTTPException(status_code=500, detail="No assistant message found in v0 response")
    
    # Generate a unique ID
    response_id = f"chatcmpl-{uuid.uuid4().hex}"
    created = int(datetime.now().timestamp())
    
    # Create choices
    choice = Choice(
        index=0,
        message=Message(
            role="assistant",
            content=assistant_message.get("content", "")
        ),
        finish_reason="stop"
    )
    
    # Create usage (estimated)
    content = assistant_message.get("content", "")
    usage = Usage(
        prompt_tokens=len(str(v0_response)),
        completion_tokens=len(content),
        total_tokens=len(str(v0_response)) + len(content)
    )
    
    return ChatCompletionResponse(
        id=response_id,
        created=created,
        model=model,
        choices=[choice],
        usage=usage
    )

async def create_streaming_response(v0_response: Dict[str, Any], model: str) -> AsyncGenerator[str, None]:
    """Create streaming response in OpenAI format"""
    response_id = f"chatcmpl-{uuid.uuid4().hex}"
    created = int(datetime.now().timestamp())
    
    messages = v0_response.get("messages", [])
    assistant_message = ""
    
    for msg in messages:
        if msg.get("role") == "assistant":
            assistant_message = msg.get("content", "")
            break
    
    # Simulate streaming by breaking the response into chunks
    words = assistant_message.split()
    
    # Send initial response
    initial_chunk = {
        'id': response_id,
        'object': 'chat.completion.chunk',
        'created': created,
        'model': model,
        'choices': [{
            'index': 0,
            'delta': {'role': 'assistant'},
            'finish_reason': None
        }]
    }
    yield f"data: {json.dumps(initial_chunk)}\n\n"
    
    # Send content in chunks
    current_text = ""
    for word in words:
        current_text += word + " "
        content_chunk = {
            'id': response_id,
            'object': 'chat.completion.chunk',
            'created': created,
            'model': model,
            'choices': [{
                'index': 0,
                'delta': {'content': word + " "},
                'finish_reason': None
            }]
        }
        yield f"data: {json.dumps(content_chunk)}\n\n"
        await asyncio.sleep(0.01)  # Small delay for streaming effect
    
    # Send final response
    final_chunk = {
        'id': response_id,
        'object': 'chat.completion.chunk',
        'created': created,
        'model': model,
        'choices': [{
            'index': 0,
            'delta': {},
            'finish_reason': 'stop'
        }]
    }
    yield f"data: {json.dumps(final_chunk)}\n\n"
    
    yield "data: [DONE]\n\n"

# API endpoints
@app.get("/")
async def root():
    return {"message": "v0.dev OpenAI Compatible API", "version": "1.0.0"}

@app.get("/v1/models")
async def list_models():
    """List available models (mock OpenAI format)"""
    return {
        "object": "list",
        "data": [
            {
                "id": "v0-gpt-5",
                "object": "model",
                "created": int(datetime.now().timestamp()),
                "owned_by": "v0.dev"
            },
            {
                "id": "v0-gpt-4",
                "object": "model",
                "created": int(datetime.now().timestamp()),
                "owned_by": "v0.dev"
            }
        ]
    }

@app.post("/v1/chat/completions")
async def create_chat_completion(request: ChatCompletionRequest):
    """Create chat completion (OpenAI compatible)"""
    try:
        # Map OpenAI model to v0.dev model
        model_config = {
            "modelId": request.model,
            "imageGenerations": True,
            "thinking": True
        }
        
        # Create chat with v0.dev
        v0_response = await v0_client.create_chat(
            messages=request.messages,
            model_config=model_config,
            project_id=request.project_id
        )
        
        if request.stream:
            return StreamingResponse(
                create_streaming_response(v0_response, request.model),
                media_type="text/plain"
            )
        else:
            return create_openai_response(v0_response, request.model)
            
    except httpx.HTTPStatusError as e:
        raise HTTPException(status_code=e.response.status_code, detail=str(e))
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/health")
async def health_check():
    return {"status": "healthy", "timestamp": datetime.now().isoformat()}

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