"""
Native RAG service using Google Gemini API.
This provides NotebookLM-like functionality without requiring browser automation.
"""
import os
import sys
import pysqlite3
sys.modules["sqlite3"] = sys.modules.pop("pysqlite3")
from flask import current_app


def create_rag_notebook(report_id, user_id):
    """
    Create a RAG-powered notebook for a report using Google Gemini API.

    This is an alternative to NotebookLM that works server-side without
    requiring browser authentication.

    Args:
        report_id: The ID of the report to create a notebook for
        user_id: The ID of the user creating the notebook

    Returns:
        notebook_url: URL to access the notebook
    """
    try:
        from app.models.report import Report
        from app.extensions import db
        import chromadb
        from google.generativeai import GenerativeModel
        import google.generativeai as genai

        # Get the report
        report = Report.query.get(report_id)
        if not report:
            return None

        # Initialize Gemini
        api_key = os.getenv("GOOGLE_GEMINI_API_KEY")
        if not api_key:
            current_app.logger.error("GOOGLE_GEMINI_API_KEY not set")
            return None

        genai.configure(api_key=api_key)

        # Create a collection for this user
        client = chromadb.PersistentClient(path=f"/tmp/chroma_{user_id}")
        collection_name = f"report_{report_id}"

        try:
            # Delete existing collection if any
            client.delete_collection(name=collection_name)
        except:
            pass

        collection = client.get_or_create_collection(
            name=collection_name,
            metadata={"hnsw:space": "cosine"}
        )

        # Prepare document content from the report
        documents = []
        metadata = []

        # Add report description
        if report.project_description:
            documents.append(f"Project Description: {report.project_description}")
            metadata.append({"type": "description", "report_id": report_id})

        # Add market analysis
        if report.market_analysis:
            import json
            try:
                market_data = json.loads(report.market_analysis)
                for key, value in market_data.items():
                    if isinstance(value, (str, int, float)):
                        documents.append(f"{key}: {value}")
                        metadata.append({"type": "market", "key": key, "report_id": report_id})
            except:
                documents.append(f"Market Analysis: {report.market_analysis}")
                metadata.append({"type": "market", "report_id": report_id})

        # Add synthesis/result
        if report.synthesis_result:
            import json
            try:
                synthesis_data = json.loads(report.synthesis_result)
                verdict = synthesis_data.get("verdict", "")
                explanation = synthesis_data.get("verdict_explanation", "")
                documents.append(f"Verdict: {verdict}\nExplanation: {explanation}")
                metadata.append({"type": "verdict", "report_id": report_id})
            except:
                documents.append(f"Analysis Result: {report.synthesis_result}")
                metadata.append({"type": "synthesis", "report_id": report_id})

        # Add documents to collection
        if documents:
            collection.add(
                documents=documents,
                metadatas=metadata,
                ids=[f"doc_{i}" for i in range(len(documents))]
            )

        # Store notebook reference in database
        notebook_url = f"/study/notebook/{report_id}?rag=true"

        # Update report with RAG enabled flag
        report.report_data = report.report_data or {}
        report.report_data['rag_enabled'] = True
        report.report_data['rag_collection'] = collection_name
        db.session.commit()

        current_app.logger.info(f"Created RAG notebook for report {report_id}")
        return notebook_url

    except Exception as e:
        current_app.logger.error(f"Error creating RAG notebook: {e}")
        import traceback
        current_app.logger.error(traceback.format_exc())
        return None


def query_rag_notebook(report_id, user_id, question):
    """
    Query a RAG notebook with a question.

    Args:
        report_id: The ID of the report
        user_id: The ID of the user
        question: The question to ask

    Returns:
        answer: The AI-generated answer
    """
    try:
        import chromadb
        from google.generativeai import GenerativeModel
        import google.generativeai as genai

        # Initialize Gemini
        api_key = os.getenv("GOOGLE_GEMINI_API_KEY")
        genai.configure(api_key=api_key)

        # Get the collection
        client = chromadb.PersistentClient(path=f"/tmp/chroma_{user_id}")
        collection_name = f"report_{report_id}"

        try:
            collection = client.get_collection(name=collection_name)
        except:
            return "عذراً، لم يتم العثور على دفتر المعرفة. يرجى إنشاء دفتر أولاً."

        # Query the collection
        results = collection.query(
            query_texts=[question],
            n_results=5
        )

        # Prepare context from retrieved documents
        context_parts = []
        for i, doc in enumerate(results['documents'][0]):
            context_parts.append(f"[Source {i+1}]: {doc}")

        context = "\n\n".join(context_parts)

        # Generate answer using Gemini
        model = GenerativeModel('gemini-2.5-flash')
        prompt = f"""You are a helpful assistant analyzing a business feasibility study.

Based on the following information from the report, answer the user's question:

CONTEXT:
{context}

QUESTION: {question}

Provide a helpful, accurate answer in Arabic. If the information is not in the context, say so honestly."""

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

        return answer

    except Exception as e:
        current_app.logger.error(f"Error querying RAG notebook: {e}")
        import traceback
        current_app.logger.error(traceback.format_exc())
        return f"عذراً، حدث خطأ: {str(e)}"