import json
from pathlib import Path
from fastapi import FastAPI, UploadFile, File, Form, HTTPException, Request
from fastapi.responses import HTMLResponse, StreamingResponse, JSONResponse
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel, HttpUrl

from .rag import RAGManager

# Initialize FastAPI app
app = FastAPI(
    title="Premium RAG Application",
    description="FastAPI + LangChain + OpenRouter owl-alpha RAG Web Interface",
    version="1.0.0"
)

# Enable CORS for local development
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Initialize RAG Manager
try:
    rag_manager = RAGManager()
except Exception as e:
    print(f"Error initializing RAG Manager: {e}")
    rag_manager = None

# Paths
BASE_DIR = Path(__file__).resolve().parent
TEMPLATES_DIR = BASE_DIR / "templates"
TEMPLATES_DIR.mkdir(exist_ok=True)

# Pydantic schemas for requests
class URLUploadRequest(BaseModel):
    url: str
    chunk_size: int = 800
    chunk_overlap: int = 100

class TextUploadRequest(BaseModel):
    text: str
    source_name: str
    chunk_size: int = 800
    chunk_overlap: int = 100

class ChatRequest(BaseModel):
    query: str
    system_prompt: str = None
    temperature: float = 0.7
    k: int = 4

# --- API Endpoints -----------------------------------------------------------

@app.get("/", response_class=HTMLResponse)
async def index():
    """Serves the main web UI page."""
    index_path = TEMPLATES_DIR / "index.html"
    if not index_path.exists():
        raise HTTPException(status_code=404, detail="index.html templates not found.")
    return index_path.read_text(encoding="utf-8")

@app.get("/api/status")
async def status():
    """Checks model availability, configuration details, and returns list of documents."""
    if rag_manager is None:
        return JSONResponse(
            status_code=500,
            content={
                "ok": False,
                "error": "RAG Manager failed to initialize. Please check if OPENROUTER_API_KEY is configured in your .env file."
            }
        )
    
    try:
        sources = rag_manager.get_indexed_sources()
        return {
            "ok": True,
            "model": rag_manager.model_name,
            "embedding_model": rag_manager.embedding_model,
            "document_count": len(sources),
            "documents": sources
        }
    except Exception as e:
        return JSONResponse(
            status_code=500,
            content={"ok": False, "error": f"Failed to fetch status: {str(e)}"}
        )

@app.post("/api/upload/file")
async def upload_file(
    file: UploadFile = File(...),
    chunk_size: int = Form(800),
    chunk_overlap: int = Form(100)
):
    """Handles file uploads (TXT, MD, PDF), splits the text, embeds, and saves to database."""
    if rag_manager is None:
        raise HTTPException(status_code=500, detail="RAG Manager not initialized.")
        
    try:
        content = await file.read()
        result = rag_manager.add_file(content, file.filename, chunk_size, chunk_overlap)
        return {"ok": True, "message": f"Successfully indexed {file.filename}.", "data": result}
    except Exception as e:
        raise HTTPException(status_code=400, detail=f"File upload and indexing failed: {str(e)}")

@app.post("/api/upload/url")
async def upload_url(request: URLUploadRequest):
    """Scrapes a URL, extracts visible text content, embeds, and saves to database."""
    if rag_manager is None:
        raise HTTPException(status_code=500, detail="RAG Manager not initialized.")
        
    try:
        result = await rag_manager.add_url(request.url, request.chunk_size, request.chunk_overlap)
        return {"ok": True, "message": f"Successfully scraped and indexed {request.url}.", "data": result}
    except Exception as e:
        raise HTTPException(status_code=400, detail=f"URL scraping and indexing failed: {str(e)}")

@app.post("/api/upload/text")
async def upload_text(request: TextUploadRequest):
    """Indexes directly pasted raw text content."""
    if rag_manager is None:
        raise HTTPException(status_code=500, detail="RAG Manager not initialized.")
        
    try:
        result = rag_manager.add_text_content(
            request.text, 
            request.source_name, 
            request.chunk_size, 
            request.chunk_overlap
        )
        return {"ok": True, "message": f"Successfully indexed raw text source '{request.source_name}'.", "data": result}
    except Exception as e:
        raise HTTPException(status_code=400, detail=f"Text indexing failed: {str(e)}")

@app.delete("/api/reset")
async def reset_database():
    """Clears all indexed vectors and document caches."""
    if rag_manager is None:
        raise HTTPException(status_code=500, detail="RAG Manager not initialized.")
        
    try:
        rag_manager.clear_database()
        return {"ok": True, "message": "Database and local serialized files cleared successfully."}
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Reset failed: {str(e)}")

@app.post("/api/chat")
async def chat(request: ChatRequest):
    """Streams OpenRouter RAG answers with citations back via Server-Sent Events (SSE)."""
    if rag_manager is None:
        raise HTTPException(status_code=500, detail="RAG Manager not initialized.")

    async def event_stream():
        try:
            async for event_type, data in rag_manager.stream_chat(
                request.query, 
                request.system_prompt, 
                request.temperature, 
                request.k
            ):
                yield f"data: {json.dumps({'event': event_type, 'data': data}, ensure_ascii=False)}\n\n"
        except Exception as e:
            yield f"data: {json.dumps({'event': 'error', 'data': str(e)}, ensure_ascii=False)}\n\n"

    return StreamingResponse(
        event_stream(),
        media_type="text/event-stream",
        headers={
            "Cache-Control": "no-cache",
            "X-Accel-Buffering": "no",  # Disables proxy buffering
            "Connection": "keep-alive",
        }
    )