import logging
import re
import os
from typing import Sequence, Union, List, Dict, Any
from enum import Enum
import mcp.types as types
from mcp.server import Server
from deepsearch.utils.search_utils import async_fetch_all_documents, async_process_documents_with_openrouter
from deepsearch.utils.pinecone_utils import PineconeManager
from deepsearch.utils.upload_to_cloudflare import CloudflareUploader
from deepsearch.utils.async_utils import log_cancellation
from datetime import datetime
from pathlib import Path
from deepsearch.utils.openrouter_utils import make_api_call

# Set up logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger("deepsearch-mcp")


class ToolName(str, Enum):
    QUICK_SEARCH = "quick-search"
    REMEMBER = "remember"
    DEEP_RESEARCH = "deep-research"


ServerTools = [
    types.Tool(
        name=ToolName.QUICK_SEARCH,
        description="""Quickly searches across database documents to find relevant information for your query.

IMPORTANT QUERY CONSTRUCTION GUIDELINES:
1. Use the FULL CONVERSATION HISTORY to construct a comprehensive query
2. Create a STANDALONE query that would make sense to someone with NO KNOWLEDGE of the conversation
3. Include all relevant context, background, specific details, dates, names, and technical terms
4. Phrase the query as a complete question or set of questions
5. Explain any references to previous conversation elements explicitly

This tool is often useful as a follow-up to the deep-research tool when you need additional information.""",
        inputSchema={
            "type": "object",
            "properties": {
                "query": {
                    "type": "string",
                    "description": "A detailed, self-contained query that thoroughly explains the context, background, and specific aspects of what you're looking for. The query MUST be constructed using the full conversation history and should be written as if explaining the request to someone with no knowledge of the conversation. Include any relevant dates, names, technical terms, or specific details that could help narrow down the search."
                }
            },
            "required": ["query"],
        },
    ),
    types.Tool(
        name=ToolName.REMEMBER,
        description="Saves a memory to the knowledge base. Triggers whenever the user says the word remember.",
        inputSchema={
            "type": "object",
            "properties": {
                "memory_text": {
                    "type": "string",
                    "description": "Text content to be remembered"
                }
            },
            "required": ["memory_text"],
        },
    ),
    types.Tool(
        name=ToolName.DEEP_RESEARCH,
        description="""Performs a streamlined research process with a single clarification step. This tool enables a focused conversation between the user and the search system to refine the query and retrieve the most relevant information.

IMPORTANT: ALWAYS start with stage="CLARIFICATION" for the initial call.

The tool works in two simple stages:
1. CLARIFICATION: Initial stage that asks ONE clarification question to improve search results
2. SEARCH: Used after receiving the user's response to the clarification question. This stage performs the search and generates a comprehensive answer in one step.

Each response will include:
- RESULTS: Information retrieved from the search or the final answer
- NEXT_STEPS: Instructions for Claude on what to do next
- NEXT_STAGE: The stage to use in the next tool call

IMPORTANT FLOW REQUIREMENTS:
- The CLARIFICATION stage asks exactly one high-quality question to improve search results
- After receiving the user's response, the tool proceeds directly to the SEARCH stage
- The SEARCH stage performs the search and generates a comprehensive answer in one step

REFINED QUERY REQUIREMENTS:
- When transitioning to the SEARCH stage, you must generate a comprehensive, detailed refined query
- This refined query should incorporate the full conversation context, user's response to the clarification question, and selected sources
- Include complete source names/paths, not just numbers
- Be highly specific about what information to search for
- Make the query detailed, comprehensive, and standalone

Follow the NEXT_STEPS instructions in each response to determine how to proceed.""",
        inputSchema={
            "type": "object",
            "properties": {
                "query": {
                    "type": "string",
                    "description": "The search query from the user."
                },
                "refined_query": {
                    "type": "string",
                    "description": "A comprehensive, detailed query generated by Claude based on the conversation context, user's response to the clarification question, and selected sources. This should be a standalone query that includes specific source names/paths and all relevant details."
                },
                "stage": {
                    "type": "string",
                    "description": "The current stage of the interaction. IMPORTANT: ALWAYS use 'CLARIFICATION' for the initial call.",
                    "enum": ["CLARIFICATION", "SEARCH"],
                    "default": "CLARIFICATION"
                }
            },
            "required": ["query"],
        },
    ),
]


@log_cancellation
async def perform_search_analysis(
    query: str,
    pinecone_client: PineconeManager,
    cloudflare_uploader: CloudflareUploader,
    request_id: str,
    inflight_requests: dict,
    model: str = "openai/o3-mini-high"
) -> Sequence[Union[types.TextContent, types.ImageContent, types.EmbeddedResource]]:
    """
    Perform search analysis using Pinecone and Cloudflare
    """
    try:
        inflight_requests[request_id] = "running"

        if not query:
            logger.error("Empty query provided")
            raise ValueError("Query parameter is required")

        logger.info("Starting document retrieval and analysis...")

        # Step 1: Get search results and fetch documents
        search_results = pinecone_client.search_documents(
            query, min_normalized_score=0.2)
        logger.info(f"Retrieved {len(search_results)} documents")

        if inflight_requests.get(request_id) == "cancelled":
            return []

        # Step 2: Fetch documents
        documents = await async_fetch_all_documents(search_results, cloudflare_uploader)

        if inflight_requests.get(request_id) == "cancelled":
            return []

        # Step 3: Process documents with OpenRouter
        processed_results = await async_process_documents_with_openrouter(
            query=query,
            documents=documents,
            model=model
        )

        if inflight_requests.get(request_id) == "cancelled":
            return []

        # Utility to strip out repeated "No relevant information..." lines
        def _clean_extracted_info(info: str) -> str:
            """
            Remove repeated "No relevant information found in the document." statements
            and return what's left. If nothing is left, it means there's no real info.
            """
            cleaned = info.replace(
                "No relevant information found in the document.", "")
            return cleaned.strip()

        # Filter out documents that end up with no actual content
        filtered_processed_results = {}
        for doc_path, raw_info in processed_results.items():
            cleaned_info = _clean_extracted_info(raw_info)
            # If there's still something left besides whitespace, we keep it; otherwise we skip
            if cleaned_info:
                filtered_processed_results[doc_path] = cleaned_info

        # If no documents had relevant information, return empty list
        if not filtered_processed_results:
            logger.info("No documents contained relevant information")
            inflight_requests[request_id] = "done"
            return [types.TextContent(type="text", text="No results found for the given query.")]

        # Combine results from filtered documents
        all_results = []
        for doc_path, info in filtered_processed_results.items():
            # Get the matching normalized score from search_results
            score = next(
                entry['normalized_score']
                for entry in search_results
                if entry['cloudflare_path'] == doc_path
            )
            # Get the raw document text
            raw_document = documents.get(
                doc_path, "Document text not available")
            all_results.append({
                'source': doc_path,
                'score': score,
                'raw_document': raw_document,
                'extracted_info': info
            })

        # Sort by score in descending order
        results = sorted(all_results, key=lambda x: x['score'], reverse=True)

        # Format output for MCP
        formatted_output = []
        for result in results:
            section = [
                f"\nSource: {result['source']}",
                f"Score: {result['score']:.3f}",
                "Raw Document:",
                f"{result['raw_document']}",
                "\nExtracted Information:",
                f"{result['extracted_info']}",
                "=" * 80
            ]
            formatted_output.append("\n".join(section))

        inflight_requests[request_id] = "done"
        return [types.TextContent(type="text", text="\n".join(formatted_output))]

    except Exception as e:
        logger.error(f"Error in perform_search_analysis: {str(e)}")
        inflight_requests[request_id] = "error"
        raise


@log_cancellation
async def deep_research(
    query: str,
    stage: str,
    request_id: str,
    pinecone_client: PineconeManager,
    cloudflare_uploader: CloudflareUploader,
    inflight_requests: dict,
    refined_query: str = ""
) -> Sequence[Union[types.TextContent, types.ImageContent, types.EmbeddedResource]]:
    """
    Perform an interactive search with a single clarification step.

    The function maintains state between calls using request_id and handles
    different stages of the interaction process.
    """
    try:
        from deepsearch.utils.prompts import (
            clarification_assessment_prompt,
            search_refinement_prompt,
            final_answer_prompt
        )

        inflight_requests[request_id] = "running"

        # Store state in a module-level variable
        if not hasattr(deep_research, 'state'):
            deep_research.state = {}

        # Check if this is the first call for this request_id and force CLARIFICATION stage
        if request_id not in deep_research.state:
            # This is the first call for this request_id
            if stage != "CLARIFICATION":
                logger.warning(
                    f"First call for request {request_id} should use CLARIFICATION stage, but got {stage}. Forcing CLARIFICATION.")
                stage = "CLARIFICATION"

            # Initialize state for this request
            deep_research.state[request_id] = {
                "original_query": query,
                "refined_query": query,
                "conversation_history": [],
                "search_history": [],
                "clarification_asked": False  # Track if we've already asked a clarification
            }

        state = deep_research.state[request_id]

        # Update conversation history with refined query if provided
        if refined_query:
            state["conversation_history"].append({
                "role": "user",
                "content": refined_query
            })

        # Format conversation history for prompts
        conversation_text = "\n".join([
            f"{'Assistant' if item['role'] == 'assistant' else 'User'}: {item['content']}"
            for item in state["conversation_history"]
        ])

        # Handle different stages
        if stage == "CLARIFICATION":
            # If we've already asked a clarification and received a response, proceed to SEARCH
            if state["clarification_asked"] and refined_query:
                # Use the provided refined query directly
                # Store refined query
                state["refined_query"] = refined_query

                # Store this interaction in conversation history
                state["conversation_history"].append({
                    "role": "assistant",
                    "content": f"Thank you for the additional information. I'll search based on your response."
                })

                # Format response to proceed to SEARCH
                response = (
                    f"RESULTS:\n"
                    f"I have received your clarification and will now proceed with the search.\n\n"
                    f"NEXT_STEPS:\n"
                    f"1. Inform the user that I'm searching based on their query and clarification\n"
                    f"2. Generate a comprehensive, detailed refined query based on the conversation context\n"
                    f"3. Call this tool again with ONLY the stage='SEARCH' and your refined_query parameters\n\n"
                    f"NEXT_STAGE: SEARCH\n\n"
                    f"IMPORTANT - REFINED QUERY FORMAT:\n"
                    f"When calling this tool for the SEARCH stage, use this exact format:\n"
                    f"{{\"stage\": \"SEARCH\", \"refined_query\": \"Your detailed query here\"}}\n\n"
                    f"DO NOT include the original 'query' parameter.\n\n"
                    f"Your refined_query MUST be comprehensive and detailed, following this format:\n"
                    f"'Search for [specific information] in source [full source name/path]. [Detailed context including specific details, dates, technical terms, etc.].'\n\n"
                    f"Example of a good refined query:\n"
                    f"\"Search for detailed raw materials cost calculations in source 'zenfold_autoupload_docs/dig_comparision_iosynth_3_3_2025_1_2503032355.txt'. Focus specifically on calculating the cost per kilogram for each of the 20 individual chemical compounds listed in the comparison table, with particular attention to nitrogen-based compounds as mentioned by the user. Include any cost variations over the past quarter if available, and prioritize the most recent pricing data from March 2025.\"\n\n"
                    f"DO NOT simply restate the original query. The refined query must be standalone, comprehensive, and explicitly mention the full source names/paths based on the user's selection."
                )

                return [types.TextContent(type="text", text=response)]

            # First time in CLARIFICATION stage - ask one clarification question
            # Get source summaries based on query
            search_results = pinecone_client.search_documents(
                query, min_normalized_score=0.2)

            # Format results with source and summary only
            summaries = []
            for result in search_results:
                summaries.append({
                    'source': result['cloudflare_path'],
                    'score': result['normalized_score'],
                    'summary': result['summary']
                })

            # Format summaries for prompt
            summaries_text = "\n".join([
                f"Source: {s['source']}\nScore: {s['score']:.3f}\nSummary: {s['summary']}"
                # Limit to top 100 for more comprehensive context
                for s in summaries[:100]
            ])

            # Get a single clarification question
            assessment = await make_api_call(
                clarification_assessment_prompt.format(
                    query=query,
                    summaries=summaries_text,
                    conversation_history=conversation_text
                ),
                model="anthropic/claude-3.7-sonnet:thinking"
            )

            # Store the assessment and summaries
            state["last_assessment"] = assessment
            state["current_summaries"] = summaries
            state["clarification_asked"] = True

            # Extract the clarification question, explanation, and relevant sources
            question_match = re.search(r"CLARIFICATION_QUESTION:(.*?)(?:EXPLANATION:|$)",
                                       assessment, re.DOTALL)
            explanation_match = re.search(
                r"EXPLANATION:(.*?)(?:RELEVANT_SOURCES:|$)", assessment, re.DOTALL)
            sources_match = re.search(
                r"RELEVANT_SOURCES:(.*?)(?:POTENTIAL_REFINED_QUERY:|$)", assessment, re.DOTALL)

            question = question_match.group(1).strip(
            ) if question_match else "Could you provide more details about your query?"
            explanation = explanation_match.group(
                1).strip() if explanation_match else ""
            relevant_sources = sources_match.group(
                1).strip() if sources_match else "No specific sources identified."

            # Format relevant sources with better structure
            formatted_sources = "RELEVANT_SOURCES:\n"
            if relevant_sources and relevant_sources != "No specific sources identified.":
                # Try to split the sources by newlines or commas
                source_list = [s.strip() for s in re.split(
                    r'[\n,]', relevant_sources) if s.strip()]
                if source_list:
                    formatted_sources += "\n".join(
                        [f"{i+1}. {source}" for i, source in enumerate(source_list)])
                else:
                    formatted_sources += relevant_sources
            else:
                formatted_sources += "No specific sources identified."

            # Store this interaction in conversation history
            state["conversation_history"].append({
                "role": "assistant",
                "content": f"To provide better results, I need to ask: {question}\n\nReason: {explanation}\n\nBased on sources: {relevant_sources}"
            })

            # Format response with clear sections
            response = (
                f"RESULTS:\n"
                f"I found some potentially relevant sources, but I need one piece of information to provide the best results.\n\n"
                f"{explanation}\n\n"
                f"{formatted_sources}\n\n"
                f"CLARIFICATION_QUESTION:\n{question}\n\n"
                f"NEXT_STEPS:\n"
                f"1. Present the list of RELEVANT_SOURCES to the user\n"
                f"2. Ask the user EXPLICITLY: 'Which of these sources would you like me to focus on? Please specify by their full names or numbers.'\n"
                f"3. Present the clarification question to the user\n"
                f"4. Collect their response to BOTH the source selection AND the clarification question\n"
                f"5. Generate a comprehensive, detailed refined query that follows this format:\n"
                f"   'Search for [specific information] in source [full source name/path]. [Detailed context including specific details, dates, technical terms, etc.].'\n"
                f"6. Call this tool again with ONLY the stage='CLARIFICATION' and your refined_query parameters\n\n"
                f"IMPORTANT REFINED QUERY REQUIREMENTS:\n"
                f"- Include the FULL source names/paths, not just numbers\n"
                f"- Be HIGHLY SPECIFIC about what information to search for\n"
                f"- Make the query DETAILED and COMPREHENSIVE (aim for 3-5 sentences)\n"
                f"- Include ALL relevant context from the conversation\n"
                f"- Specify any relevant date ranges, technical terms, or specific details\n"
                f"- The refined query MUST be standalone and not require additional context\n\n"
                f"NEXT_STAGE: CLARIFICATION"
            )

            return [types.TextContent(type="text", text=response)]

        elif stage == "SEARCH":
            # Use the refined query provided in the request
            refined_query_for_search = refined_query if refined_query else state.get(
                "refined_query", query)

            # Perform search with refined query
            search_results = await perform_search_analysis(
                query=refined_query_for_search,
                pinecone_client=pinecone_client,
                cloudflare_uploader=cloudflare_uploader,
                request_id=request_id,
                inflight_requests=inflight_requests,
                model="openai/o3-mini-high"
            )

            # Extract text content
            search_content = "No relevant results found."
            if search_results and hasattr(search_results[0], 'text'):
                search_content = search_results[0].text

            # Store search results
            state["last_search_results"] = search_content

            # Generate final answer directly
            final_answer = await make_api_call(
                final_answer_prompt.format(
                    query=query,
                    conversation_history=conversation_text,
                    search_results=search_content
                ),
                model="openai/o3-mini"
            )

            # Store this interaction in conversation history
            state["conversation_history"].append({
                "role": "assistant",
                "content": final_answer
            })

            # Format response with clear sections
            response = (
                f"RESULTS:\n"
                f"{final_answer}\n\n"
                f"SOURCE_INFORMATION:\n"
                f"{search_content}\n\n"
                f"NEXT_STEPS:\n"
                f"1. Present the answer to the user\n"
                f"2. If the user needs additional information, consider using the quick-search tool as a follow-up\n"
                f"   When using quick-search, construct a standalone query that incorporates the full conversation context\n"
                f"   Example: Instead of 'Tell me more about that chemical compound', use 'Tell me more about the nitrogen-based compound mentioned in the dig_comparision_iosynth file that has a raw materials cost of X per kilogram'\n"
                f"3. The deep research process is now complete\n\n"
                f"NEXT_STAGE: COMPLETE"
            )

            # Clean up state
            if request_id in deep_research.state:
                del deep_research.state[request_id]

            return [types.TextContent(type="text", text=response)]

        else:
            return [types.TextContent(type="text", text="Invalid stage specified. Please use 'CLARIFICATION' or 'SEARCH'.")]

    except Exception as e:
        logger.error(f"Error in deep_research: {str(e)}")
        inflight_requests[request_id] = "error"
        raise


@log_cancellation
async def save_memory(
    memory_text: str,
    pinecone_client: PineconeManager,
    cloudflare_uploader: CloudflareUploader,
    request_id: str,
    inflight_requests: dict
) -> Sequence[Union[types.TextContent, types.ImageContent, types.EmbeddedResource]]:
    """Save a memory to Cloudflare and Pinecone"""
    try:
        inflight_requests[request_id] = "running"

        # Create temporary file with memory content
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        temp_file = Path(f"memory_{timestamp}.txt")
        temp_file.write_text(memory_text)

        try:
            # Upload to Cloudflare
            cloudflare_path, metadata = cloudflare_uploader.upload_document(
                str(temp_file), {})

            if not cloudflare_path:
                raise Exception("Failed to upload memory to Cloudflare")

            # Create embedding for the memory text (using same text as summary)
            embedding = pinecone_client.create_embedding(memory_text)

            # Upsert to Pinecone
            success = pinecone_client.upsert_vector(
                vector_id=f"memory_{timestamp}",
                vector_values=embedding,
                metadata={
                    "summary": memory_text,
                    "cloudflare_path": cloudflare_path
                }
            )

            if not success:
                raise Exception("Failed to save memory to Pinecone")

            inflight_requests[request_id] = "done"
            return [types.TextContent(type="text", text="Memory successfully saved.")]

        finally:
            # Clean up temp file
            temp_file.unlink(missing_ok=True)

    except Exception as e:
        logger.error(f"Error in save_memory: {str(e)}")
        inflight_requests[request_id] = "error"
        raise


def register_tools(server: Server, pinecone_client: PineconeManager, cloudflare_uploader: CloudflareUploader, inflight_requests: dict):
    @server.list_tools()
    @log_cancellation
    async def handle_list_tools() -> list[types.Tool]:
        return ServerTools

    @server.call_tool()
    @log_cancellation
    async def handle_call_tool(
        name: str, arguments: dict | None
    ) -> Sequence[Union[types.TextContent, types.ImageContent, types.EmbeddedResource]]:
        try:
            request_id = arguments.get("__request_id__")
            logger.info(f"Calling tool: {name} for request {request_id}")

            if name == ToolName.QUICK_SEARCH:
                return await perform_search_analysis(
                    query=arguments.get("query"),
                    pinecone_client=pinecone_client,
                    cloudflare_uploader=cloudflare_uploader,
                    request_id=request_id,
                    inflight_requests=inflight_requests,
                    model="openai/o3-mini-high"
                )
            elif name == ToolName.REMEMBER:
                return await save_memory(
                    memory_text=arguments.get("memory_text"),
                    pinecone_client=pinecone_client,
                    cloudflare_uploader=cloudflare_uploader,
                    request_id=request_id,
                    inflight_requests=inflight_requests
                )
            elif name == ToolName.DEEP_RESEARCH:
                return await deep_research(
                    query=arguments.get("query", ""),
                    # Default to CLARIFICATION if not provided
                    stage=arguments.get("stage", "CLARIFICATION"),
                    request_id=request_id,  # Always use the MCP request_id
                    pinecone_client=pinecone_client,
                    cloudflare_uploader=cloudflare_uploader,
                    inflight_requests=inflight_requests,
                    refined_query=arguments.get("refined_query", "")
                )
            else:
                logger.error(f"Unknown tool: {name}")
                raise ValueError(f"Unknown tool: {name}")
        except Exception as e:
            logger.error(f"Error calling tool {name}: {str(e)}")
            raise


__all__ = [
    "register_tools",
]
