import os
import requests
from dotenv import load_dotenv
from fastapi import FastAPI, File, Form, UploadFile
from fastapi.responses import StreamingResponse, JSONResponse
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from typing import List, Optional
import uvicorn

# Load environment variables
load_dotenv()
HF_TOKEN = os.environ.get("CHAT_MATE")

# ---------------- FastAPI setup ----------------
app = FastAPI(
    title="ChatMate API",
    description="Stream replies from Hugging Face Space and upload requirement docs.",
    version="1.0",
    docs_url="/apidocs",   # Swagger UI
    redoc_url="/redoc"     # ReDoc UI
)

# Allow all CORS (for browser testing)
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# ---------------- Models ----------------
class HistoryItem(BaseModel):
    role: str
    content: str

class ChatRequest(BaseModel):
    message: str
    history: Optional[List[HistoryItem]] = []

# ---------------- Endpoints ----------------
@app.post(
    "/chat-stream",
    summary="Stream assistant reply or image",
    description="Send a message and history, receive either a streamed text reply or base64-encoded image.",
    response_description="Streamed reply or image base64",
    responses={
        200: {"content": {"text/plain": {}}}
    },
    tags=["Chat"]
)
async def chat_stream(request_data: ChatRequest):
    headers = {"Authorization": f"Bearer {HF_TOKEN}"}

    response = requests.post(
        "https://fredericksundeep-chatmateapi.hf.space/chat-stream",
        json=request_data.dict(),
        headers=headers,
        stream=True
    )

    return StreamingResponse(response.iter_content(decode_unicode=True), media_type="text/plain")

@app.post(
    "/chat-stream-doc",
    summary="Upload requirement doc & generate downloadable project code (ZIP)",
    description="Upload a PDF/TXT requirement document with stack preferences, and receive the scaffolded project code as a downloadable .zip file.",
    responses={
        200: {"content": {"application/zip": {}}}
    },
    tags=["Chat"]
)
async def chat_stream_doc(
    file: UploadFile = File(..., description="Requirement document (PDF or TXT)"),
    frontend: str = Form(..., description="Frontend tech (React, Angular, etc.)"),
    backend: str = Form(..., description="Backend tech (Flask, Node.js, etc.)"),
    database: str = Form(..., description="Database (MongoDB, PostgreSQL, etc.)")
):
    headers = {"Authorization": f"Bearer {HF_TOKEN}"}

    files = {
        "file": (file.filename, file.file, file.content_type)
    }
    data = {
        "frontend": frontend,
        "backend": backend,
        "database": database
    }

    try:
        response = requests.post(
            "https://fredericksundeep-aimateapis.hf.space/chat-stream-doc",
            headers=headers,
            files=files,
            data=data
        )

        return StreamingResponse(
            iter([response.content]),
            media_type=response.headers.get("Content-Type", "application/octet-stream"),
            headers={
                "Content-Disposition": response.headers.get(
                    "Content-Disposition",
                    "attachment; filename=generated_project.zip"
                )
            }
        )

    except Exception as e:
        return JSONResponse(content={"error": str(e)}, status_code=500)

# ---------------- Startup warm-up ----------------
@app.on_event("startup")
async def warmup():
    print("🔧 Warming up...")

# ---------------- Run ----------------
if __name__ == "__main__":
    port = int(os.environ.get("PORT", 7860))
    uvicorn.run("app:app", host="0.0.0.0", port=port, reload=False)