from openrouter_llm import OpenRouterFreeAdapter, OpenRouterFreeChain
from langchain.schema import Document as LangchainDocument
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.embeddings import HuggingFaceEmbeddings
from langchain.vectorstores import FAISS
import os
import uuid
import shutil
import logging
from typing import List, Optional, Dict, Any
from pathlib import Path

import fitz  # PyMuPDF
import markdown
from fastapi import FastAPI, File, UploadFile, HTTPException, Form, Depends, BackgroundTasks
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse
from pydantic import BaseModel
from dotenv import load_dotenv

# Load environment variables
load_dotenv()

# Import LangChain components for embedding

# Import our free-only OpenRouter adapter

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Initialize FastAPI app
app = FastAPI(title="AskMyDocs API - Free LLM Edition")

# Add CORS middleware for frontend integration
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # Set to specific domain in production
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Configuration
OPENROUTER_API_KEY = os.getenv("OPENROUTER_API_KEY", "")
HF_MODEL_NAME = os.getenv(
    "HF_MODEL_NAME", "sentence-transformers/all-mpnet-base-v2")
UPLOAD_DIR = os.getenv("UPLOAD_DIR", "./uploads")
DB_DIR = os.getenv("DB_DIR", "./vectordb")

print(HF_MODEL_NAME)
# Ensure directories exist
os.makedirs(UPLOAD_DIR, exist_ok=True)
os.makedirs(DB_DIR, exist_ok=True)

# Initialize OpenRouter adapter (singleton)
openrouter_adapter = None

# Pydantic models


class QueryRequest(BaseModel):
    query: str
    collection_id: str


class QueryResponse(BaseModel):
    answer: str
    sources: List[str]


class Document(BaseModel):
    id: str
    filename: str
    content_type: str


class DocumentList(BaseModel):
    documents: List[Document]


class LLMInfo(BaseModel):
    model: str
    is_free: bool = True
    provider: str = "openrouter"


class LLMModelsList(BaseModel):
    current_model: str
    free_models: List[Dict[str, Any]]


# Global variable to store vector databases (in memory for simplicity)
# In production, you would use persistent storage
vector_dbs = {}

# Helper functions


def get_embeddings():
    """Get HuggingFace embedding model."""
    return HuggingFaceEmbeddings(model_name=HF_MODEL_NAME)


def get_openrouter_adapter():
    """Get or initialize the OpenRouter adapter for free models."""
    global openrouter_adapter

    if openrouter_adapter is None:
        openrouter_adapter = OpenRouterFreeAdapter(api_key=OPENROUTER_API_KEY)

    return openrouter_adapter


def extract_text_from_pdf(file_path):
    """Extract text content from PDF files."""
    text = ""
    try:
        doc = fitz.open(file_path)
        for page in doc:
            text += page.get_text()
        return text
    except Exception as e:
        logger.error(f"Error extracting text from PDF: {e}")
        raise HTTPException(
            status_code=500, detail=f"Error processing PDF: {str(e)}")


def extract_text_from_markdown(file_path):
    """Convert Markdown to plain text."""
    try:
        with open(file_path, 'r', encoding='utf-8') as f:
            md_content = f.read()
        html = markdown.markdown(md_content)
        # Simple HTML to text conversion - in production use a more robust method
        text = html.replace('<p>', '\n\n').replace(
            '</p>', '').replace('<br>', '\n')
        text = text.replace('<h1>', '\n\n# ').replace('</h1>', '\n')
        text = text.replace('<h2>', '\n\n## ').replace('</h2>', '\n')
        text = text.replace('<h3>', '\n\n### ').replace('</h3>', '\n')
        # Remove other HTML tags
        import re
        text = re.sub('<[^<]+?>', '', text)
        return text
    except Exception as e:
        logger.error(f"Error processing Markdown: {e}")
        raise HTTPException(
            status_code=500, detail=f"Error processing Markdown: {str(e)}")


def extract_text_from_file(file_path, content_type):
    """Extract text based on file type."""
    if content_type == "application/pdf":
        return extract_text_from_pdf(file_path)
    elif content_type == "text/markdown":
        return extract_text_from_markdown(file_path)
    elif content_type == "text/plain":
        with open(file_path, 'r', encoding='utf-8') as f:
            return f.read()
    else:
        raise HTTPException(
            status_code=400, detail=f"Unsupported file type: {content_type}")


def process_documents(collection_id: str, file_paths: List[tuple]):
    """Process documents and create vector store."""
    try:
        # Create text splitter
        text_splitter = RecursiveCharacterTextSplitter(
            chunk_size=1000,
            chunk_overlap=100,
            length_function=len,
        )

        all_docs = []
        for file_path, content_type, filename in file_paths:
            text_content = extract_text_from_file(file_path, content_type)
            chunks = text_splitter.split_text(text_content)

            # Create Document objects with metadata
            docs = [
                LangchainDocument(
                    page_content=chunk,
                    metadata={"source": filename, "chunk": i}
                )
                for i, chunk in enumerate(chunks)
            ]
            all_docs.extend(docs)

        # Create vector store
        embeddings = get_embeddings()
        vector_db = FAISS.from_documents(all_docs, embeddings)

        # Save vector store
        collection_path = os.path.join(DB_DIR, collection_id)
        os.makedirs(collection_path, exist_ok=True)
        vector_db.save_local(collection_path)

        # Store in memory (would be replaced by database lookup in production)
        vector_dbs[collection_id] = vector_db

        logger.info(
            f"Successfully processed {len(all_docs)} chunks from {len(file_paths)} documents")
    except Exception as e:
        logger.error(f"Error processing documents: {e}")
        raise HTTPException(
            status_code=500, detail=f"Error processing documents: {str(e)}")


@app.get("/")
async def index():
    return {"message": "Welcome to ask my doc"}


@app.get("/health")
async def health_check():
    return {"status": "healthy"}


@app.post("/upload", response_model=Document)
async def upload_file(
    background_tasks: BackgroundTasks,
    collection_id: str = Form(...),
    file: UploadFile = File(...),
):
    """Upload a document and process it for querying."""
    try:
        # Generate a unique ID for the document
        doc_id = str(uuid.uuid4())

        # Create collection directory if it doesn't exist
        collection_dir = os.path.join(UPLOAD_DIR, collection_id)
        os.makedirs(collection_dir, exist_ok=True)

        # Define the file path
        file_path = os.path.join(collection_dir, file.filename)

        # Determine content type
        content_type = file.content_type
        if not content_type:
            if file.filename.endswith('.pdf'):
                content_type = "application/pdf"
            elif file.filename.endswith('.md'):
                content_type = "text/markdown"
            elif file.filename.endswith('.txt'):
                content_type = "text/plain"
            else:
                raise HTTPException(
                    status_code=400, detail="Unsupported file type")

        # Save the file
        with open(file_path, "wb") as f:
            shutil.copyfileobj(file.file, f)

        # Process the document in the background
        background_tasks.add_task(
            process_documents,
            collection_id,
            [(file_path, content_type, file.filename)]
        )

        return Document(
            id=doc_id,
            filename=file.filename,
            content_type=content_type
        )
    except Exception as e:
        logger.error(f"Error uploading file: {e}")
        raise HTTPException(
            status_code=500, detail=f"Error uploading file: {str(e)}")


@app.get("/collections/{collection_id}/documents", response_model=DocumentList)
async def list_documents(collection_id: str):
    """List all documents in a collection."""
    try:
        collection_dir = os.path.join(UPLOAD_DIR, collection_id)
        if not os.path.exists(collection_dir):
            return DocumentList(documents=[])

        documents = []
        for filename in os.listdir(collection_dir):
            file_path = os.path.join(collection_dir, filename)
            if os.path.isfile(file_path):
                content_type = "application/octet-stream"
                if filename.endswith('.pdf'):
                    content_type = "application/pdf"
                elif filename.endswith('.md'):
                    content_type = "text/markdown"
                elif filename.endswith('.txt'):
                    content_type = "text/plain"

                documents.append(Document(
                    # In production, store and retrieve actual IDs
                    id=str(uuid.uuid4()),
                    filename=filename,
                    content_type=content_type
                ))

        return DocumentList(documents=documents)
    except Exception as e:
        logger.error(f"Error listing documents: {e}")
        raise HTTPException(
            status_code=500, detail=f"Error listing documents: {str(e)}")


@app.post("/query", response_model=QueryResponse)
async def query_documents(request: QueryRequest):
    """Query documents using natural language."""
    try:
        collection_id = request.collection_id

        # Check if vector DB exists in memory
        if collection_id in vector_dbs:
            vector_db = vector_dbs[collection_id]
        else:
            # Load from disk
            collection_path = os.path.join(DB_DIR, collection_id)
            if not os.path.exists(collection_path):
                raise HTTPException(
                    status_code=404, detail=f"Collection {collection_id} not found")

            embeddings = get_embeddings()
            vector_db = FAISS.load_local(collection_path, embeddings)
            vector_dbs[collection_id] = vector_db

        # Get the retriever
        retriever = vector_db.as_retriever(search_kwargs={"k": 3})

        # Get relevant documents
        docs = retriever.get_relevant_documents(request.query)

        # Extract sources
        sources = []
        for doc in docs:
            if doc.metadata.get("source") not in sources:
                sources.append(doc.metadata.get("source"))

        # Get context from documents
        context = [doc.page_content for doc in docs]

        # Get OpenRouter adapter for free LLMs
        adapter = get_openrouter_adapter()
        chain = OpenRouterFreeChain(adapter)

        # Generate answer
        answer = chain.run(request.query, context)

        return QueryResponse(
            answer=answer,
            sources=sources
        )
    except Exception as e:
        logger.error(f"Error querying documents: {e}")
        raise HTTPException(
            status_code=500, detail=f"Error querying documents: {str(e)}")


@app.delete("/collections/{collection_id}/documents/{filename}")
async def delete_document(collection_id: str, filename: str):
    """Delete a document from a collection."""
    try:
        file_path = os.path.join(UPLOAD_DIR, collection_id, filename)
        if not os.path.exists(file_path):
            raise HTTPException(
                status_code=404, detail=f"Document {filename} not found")

        os.remove(file_path)

        # Rebuild vector store if needed
        collection_path = os.path.join(DB_DIR, collection_id)
        if os.path.exists(collection_path):
            # In production, you would selectively remove documents rather than rebuilding
            shutil.rmtree(collection_path)

            # If there are still documents, rebuild the vector store
            collection_dir = os.path.join(UPLOAD_DIR, collection_id)
            if os.path.exists(collection_dir) and os.listdir(collection_dir):
                file_paths = []
                for fname in os.listdir(collection_dir):
                    fpath = os.path.join(collection_dir, fname)
                    if os.path.isfile(fpath):
                        content_type = "application/octet-stream"
                        if fname.endswith('.pdf'):
                            content_type = "application/pdf"
                        elif fname.endswith('.md'):
                            content_type = "text/markdown"
                        elif fname.endswith('.txt'):
                            content_type = "text/plain"
                        file_paths.append((fpath, content_type, fname))

                if file_paths:
                    process_documents(collection_id, file_paths)

            # Remove from in-memory cache
            if collection_id in vector_dbs:
                del vector_dbs[collection_id]

        return JSONResponse(content={"message": f"Document {filename} deleted"})
    except Exception as e:
        logger.error(f"Error deleting document: {e}")
        raise HTTPException(
            status_code=500, detail=f"Error deleting document: {str(e)}")


@app.get("/llm/info", response_model=LLMInfo)
async def get_llm_info():
    """Get the current LLM information."""
    adapter = get_openrouter_adapter()

    return LLMInfo(
        model=adapter.model,
        is_free=True,
        provider="openrouter"
    )


@app.get("/llm/models", response_model=LLMModelsList)
async def list_free_models():
    """List all available free models."""
    adapter = get_openrouter_adapter()
    free_models = adapter.list_free_models()

    # Create a simplified list for the frontend
    model_list = []
    for model in free_models:
        model_info = {
            "id": model.get("id"),
            "name": model.get("name", model.get("id")),
            "context_length": model.get("context_length", 4096),
            "provider": model.get("id").split("/")[0] if "/" in model.get("id") else "unknown"
        }
        model_list.append(model_info)

    return LLMModelsList(
        current_model=adapter.model,
        free_models=model_list
    )


@app.post("/llm/change-model")
async def change_model(model_info: LLMInfo):
    """Change the LLM model (only to another free model)."""
    adapter = get_openrouter_adapter()

    # Make sure the model has the :free suffix if it doesn't already
    model_id = model_info.model
    if not model_id.endswith(":free") and ":free" not in model_id:
        model_id = f"{model_id}:free"

    # Set the new model
    adapter.model = model_id

    return JSONResponse(content={"message": f"Model changed to {model_id}"})