from langchain.vectorstores import Chroma
from langchain.embeddings import OpenAIEmbeddings
from database import SessionLocal

@celery_app.task()
def answer_with_context(question: str, document_ids: list, user_id: int):
    """Answer questions using your own documents"""

    db = SessionLocal()
    try:
        # Retrieve relevant documents from database
        documents = db.query(Document).filter(
            Document.id.in_(document_ids),
            Document.user_id == user_id
        ).all()

        # Create embeddings and search
        embeddings = OpenAIEmbeddings()
        vectorstore = Chroma(embedding_function=embeddings)

        # Add documents to vector store
        texts = [doc.content for doc in documents]
        vectorstore.add_texts(texts)

        # Get relevant context
        docs = vectorstore.similarity_search(question, k=3)
        context = "\n".join([doc.page_content for doc in docs])

        # Generate answer with context
        client = OpenAI()
        prompt = f"""
        Answer this question based on the provided context.
        If the context doesn't contain the answer, say so.

        Context: {context}
        Question: {question}
        Answer:
        """

        response = client.chat.completions.create(
            model="gpt-4o-mini",
            messages=[{"role": "user", "content": prompt}],
            temperature=0.3
        )

        # Save Q&A to database
        qa_record = QASession(
            user_id=user_id,
            question=question,
            answer=response.choices[0].message.content,
            context_docs=document_ids,
            model_used="gpt-4o-mini"
        )
        db.add(qa_record)
        db.commit()

        return {
            "answer": response.choices[0].message.content,
            "session_id": qa_record.id
        }
    finally:
        db.close()

import base64

@celery_app.task()
def analyze_image_and_text(image_path: str, text_prompt: str):
    """Analyze images with text context"""

    # Encode image
    with open(image_path, "rb") as image_file:
        base64_image = base64.b64encode(image_file.read()).decode('utf-8')

    client = OpenAI()
    response = client.chat.completions.create(
        model="gpt-4o",  # Supports vision
        messages=[
            {
                "role": "user",
                "content": [
                    {"type": "text", "text": text_prompt},
                    {
                        "type": "image_url",
                        "image_url": {
                            "url": f"data:image/jpeg;base64,{base64_image}"
                        }
                    }
                ]
            }
        ],
        max_tokens=1000
    )

    return response.choices[0].message.content

from sqlalchemy import Column, Integer, String, Text, DateTime, JSON, Float
from sqlalchemy.ext.declarative import declarative_base
from datetime import datetime

Base = declarative_base()

class AIAnalysis(Base):
    __tablename__ = "ai_analysis"

    id = Column(Integer, primary_key=True)
    user_id = Column(Integer, nullable=False)
    input_text = Column(Text, nullable=False)
    analysis_type = Column(String(50), nullable=False)
    result = Column(Text, nullable=False)
    model_used = Column(String(50), nullable=False)
    tokens_used = Column(Integer, default=0)
    cost = Column(Float, default=0.0)
    created_at = Column(DateTime, default=datetime.utcnow)

class QASession(Base):
    __tablename__ = "qa_sessions"

    id = Column(Integer, primary_key=True)
    user_id = Column(Integer, nullable=False)
    question = Column(Text, nullable=False)
    answer = Column(Text, nullable=False)
    context_docs = Column(JSON)  # Store document IDs used
    model_used = Column(String(50), nullable=False)
    created_at = Column(DateTime, default=datetime.utcnow)

class AICache(Base):
    __tablename__ = "ai_cache"
    
    id = Column(Integer, primary_key=True)
    cache_key = Column(String(255), unique=True, nullable=False)
    result = Column(Text, nullable=False)
    expires_at = Column(DateTime, nullable=False)
    created_at = Column(DateTime, default=datetime.utcnow)

class Document(Base):
    __tablename__ = "documents"

    id = Column(Integer, primary_key=True)
    user_id = Column(Integer, nullable=False)
    title = Column(String(255), nullable=False)
    content = Column(Text, nullable=False)
    embedding = Column(JSON)  # Store vector embeddings
    created_at = Column(DateTime, default=datetime.utcnow)

from anthropic import Anthropic

@celery_app.task()
def analyze_with_claude(text: str):
    """Use Claude for analysis tasks"""
    client = Anthropic(api_key=os.getenv("ANTHROPIC_API_KEY"))

    response = client.messages.create(
        model="claude-3-5-sonnet-20241022",
        max_tokens=1000,
        messages=[{
            "role": "user",
            "content": f"Analyze this text for key insights: {text}"
        }]
    )

    return response.content[0].text

import google.generativeai as genai

@celery_app.task()
def process_with_gemini(prompt: str):
    """Use Gemini for fast processing"""
    genai.configure(api_key=os.getenv("GOOGLE_API_KEY"))
    model = genai.GenerativeModel('gemini-pro')

    response = model.generate_content(prompt)
    return response.text