import chromadb
from flask import Flask, request, jsonify, send_from_directory
import os
import json
import requests
import logging
from datetime import datetime

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

app = Flask(__name__)

# ChromaDB client setup
try:
    client = chromadb.HttpClient(host="chromadb", port=8000)
    logger.info("Connected to ChromaDB successfully")
except Exception as e:
    logger.error(f"Failed to connect to ChromaDB: {e}")
    client = None

# Company facts database (in-memory for now, could be replaced with persistent storage)
company_facts = {
    "Apple": {
        "products": [
            {"name": "iPhone", "release_date": "2007", "specifications": "Smartphone with iOS"},
            {"name": "MacBook", "release_date": "2006", "specifications": "Laptop with M1 chip"},
            {"name": "iPad", "release_date": "2010", "specifications": "Tablet with iOS"},
            {"name": "Apple Watch", "release_date": "2015", "specifications": "Smartwatch with watchOS"}
        ],
        "founded": "1976",
        "ceo": "Tim Cook",
        "headquarters": "Cupertino, California",
        "ipo_year": "1980",
        "market_cap": "$2.8T (2026)",
        "key_executives": [
            {"name": "Tim Cook", "position": "CEO"},
            {"name": "Johny Srouji", "position": "CFO"},
            {"name": "Katherine Adams", "position": "Chief Design Officer"}
        ],
        "business_segments": ["Consumer Electronics", "Software", "Services"]
    },
    "Microsoft": {
        "products": [
            {"name": "Windows", "release_date": "1985", "specifications": "Operating system"},
            {"name": "Office", "release_date": "1989", "specifications": "Productivity suite"},
            {"name": "Azure", "release_date": "2010", "specifications": "Cloud computing platform"},
            {"name": "Xbox", "release_date": "2001", "specifications": "Gaming console"}
        ],
        "founded": "1975",
        "ceo": "Satya Nadella",
        "headquarters": "Redmond, Washington",
        "ipo_year": "1986",
        "market_cap": "$3.2T (2026)",
        "key_executives": [
            {"name": "Satya Nadella", "position": "CEO"},
            {"name": "Amy Hood", "position": "CFO"},
            {"name": "Kevin Scott", "position": "CTO"}
        ],
        "business_segments": ["Software", "Cloud Services", "Gaming", "Productivity"]
    },
    "Google": {
        "products": [
            {"name": "Search Engine", "release_date": "1998", "specifications": "Web search platform"},
            {"name": "Android", "release_date": "2008", "specifications": "Mobile operating system"},
            {"name": "Gmail", "release_date": "2004", "specifications": "Email service"},
            {"name": "YouTube", "release_date": "2005", "specifications": "Video sharing platform"}
        ],
        "founded": "1998",
        "ceo": "Sundar Pichai",
        "headquarters": "Mountain View, California",
        "ipo_year": "2004",
        "market_cap": "$1.7T (2026)",
        "key_executives": [
            {"name": "Sundar Pichai", "position": "CEO"},
            {"name": "Ruth Porat", "position": "CFO"},
            {"name": "Rajen S. Suri", "position": "Chief Technology Officer"}
        ],
        "business_segments": ["Search", "Advertising", "Cloud", "Mobile"]
    },
    "Amazon": {
        "products": [
            {"name": "Amazon Web Services (AWS)", "release_date": "2006", "specifications": "Cloud computing platform"},
            {"name": "Kindle", "release_date": "2007", "specifications": "E-reader device"},
            {"name": "Alexa", "release_date": "2014", "specifications": "Voice assistant"},
            {"name": "Prime Video", "release_date": "2008", "specifications": "Streaming service"}
        ],
        "founded": "1994",
        "ceo": "Andy Jassy",
        "headquarters": "Seattle, Washington",
        "ipo_year": "1997",
        "market_cap": "$1.5T (2026)",
        "key_executives": [
            {"name": "Andy Jassy", "position": "CEO"},
            {"name": "Brian T. Olsavsky", "position": "CFO"},
            {"name": "Wendy J. Smith", "position": "Chief Technology Officer"}
        ],
        "business_segments": ["E-commerce", "Cloud Computing", "Digital Streaming", "Advertising"]
    }
}

def get_embedding(text):
    """
    Get embedding using the OpenAI-compatible server at http://example.com:4000
    """
    try:
        # Using the OpenAI-compatible endpoint for embeddings
        response = requests.post(
            "http://example.com:4000/v1/embeddings",
            json={
                "input": text,
                "model": "text-embedding-3-small"  # or whatever model you're using
            },
            timeout=30
        )
        response.raise_for_status()
        embedding = response.json()['data'][0]['embedding']
        return embedding
    except Exception as e:
        logger.error(f"Error getting embedding: {e}")
        return None

def get_diverse_articles(articles, max_diverse=5):
    """
    Filter articles to ensure diversity in content, sources, and perspectives
    """
    if len(articles) <= max_diverse:
        return articles
    
    # More sophisticated diversity algorithm
    diverse_articles = []
    source_count = {}
    topic_count = {}
    
    # First pass: try to get articles from different sources
    for article in articles:
        source = article.get('metadata', {}).get('source', 'unknown')
        topic = article.get('metadata', {}).get('topic', 'unknown')
        
        # If we haven't reached max diversity and this source is new, add it
        if len(diverse_articles) < max_diverse and source not in source_count:
            diverse_articles.append(article)
            source_count[source] = 1
            topic_count[topic] = topic_count.get(topic, 0) + 1
    
    # Second pass: fill remaining slots with different topics if possible
    if len(diverse_articles) < max_diverse:
        for article in articles:
            if len(diverse_articles) >= max_diverse:
                break
            source = article.get('metadata', {}).get('source', 'unknown')
            topic = article.get('metadata', {}).get('topic', 'unknown')
            
            # Add article if it's from a different topic and we haven't seen too many from this topic
            if source not in source_count and topic_count.get(topic, 0) < 2:
                diverse_articles.append(article)
                source_count[source] = 1
                topic_count[topic] = topic_count.get(topic, 0) + 1
    
    # If we still don't have enough, just return first few
    if len(diverse_articles) < max_diverse:
        return articles[:max_diverse]
    
    return diverse_articles

def query_chroma(question, n_results=10):
    """
    Query ChromaDB for articles related to the question
    """
    if not client:
        return {"error": "ChromaDB connection failed"}
    
    try:
        # Get embedding for the question
        query_embedding = get_embedding(question)
        if not query_embedding:
            return {"error": "Failed to get embedding"}
        
        # Query the collection
        results = client.get_or_create_collection("news").query(
            query_embeddings=[query_embedding],
            n_results=n_results,
        )
        
        return results
    except Exception as e:
        logger.error(f"Error querying ChromaDB: {e}")
        return {"error": f"Query failed: {str(e)}"}

# New endpoints implementation based on the OpenAPI specification

@app.route("/query", methods=["POST"])
def query_vector_database():
    """Query the vector database"""
    data = request.get_json()
    question = data.get("question")
    
    if not question:
        return jsonify({"error": "Missing 'question' in request body"}), 400
    
    # Query ChromaDB for relevant articles
    results = query_chroma(question, n_results=10)
    
    if "error" in results:
        return jsonify({"error": results["error"]}), 500
    
    # Process results to create diverse article set
    mcp_results = []
    for doc, score, meta in zip(
            results.get("documents", [[]])[0],
            results.get("distances", [[]])[0],
            results.get("metadatas", [[]])[0]):
        mcp_results.append({
            "document": doc,
            "score": float(score),
            "metadata": meta
        })
    
    # Apply diversity filtering
    diverse_results = get_diverse_articles(mcp_results, 5)
    
    return jsonify({"results": diverse_results})

@app.route("/articles/query", methods=["POST"])
def query_articles():
    """Query articles based on a question with diversity filtering"""
    data = request.get_json()
    question = data.get("question")
    max_results = data.get("max_results", 5)
    
    if not question:
        return jsonify({"error": "Missing 'question' in request body"}), 400
    
    # Query ChromaDB for relevant articles
    results = query_chroma(question, n_results=max_results)
    
    if "error" in results:
        return jsonify({"error": results["error"]}), 500
    
    # Process results to create diverse article set
    mcp_results = []
    for doc, score, meta in zip(
            results.get("documents", [[]])[0],
            results.get("distances", [[]])[0],
            results.get("metadatas", [[]])[0]):
        mcp_results.append({
            "document": doc,
            "score": float(score),
            "metadata": meta
        })
    
    # Apply diversity filtering
    diverse_results = get_diverse_articles(mcp_results, max_results)
    
    return jsonify({
        "results": diverse_results,
        "query": question
    })

@app.route("/articles/latest/<field>", methods=["GET"])
def get_latest_articles(field):
    """Get latest articles about a specific field with diversity"""
    # This would typically query the database for latest articles about the field
    # For now, we'll return some sample data
    sample_articles = [
        {
            "document": f"Latest article about {field}",
            "score": 0.95,
            "metadata": {
                "source": "Sample Source",
                "topic": field,
                "date": "2026-01-31"
            }
        }
    ]
    
    return jsonify({
        "results": sample_articles,
        "field": field
    })

@app.route("/company/<company_name>/facts", methods=["GET"])
def get_company_facts(company_name):
    """Get facts about a specific company"""
    facts = company_facts.get(company_name, {})
    if not facts:
        return jsonify({"error": f"Company {company_name} not found"}), 404
    
    return jsonify({
        "company": company_name,
        "facts": facts
    })

@app.route("/company/<company_name>/products", methods=["GET"])
def get_company_products(company_name):
    """Get products information for a company"""
    facts = company_facts.get(company_name, {})
    products = facts.get("products", [])
    
    if not products:
        return jsonify({"error": f"No products found for company {company_name}"}), 404
    
    return jsonify({
        "company": company_name,
        "products": products
    })

@app.route("/company/facts/update", methods=["POST"])
def update_company_facts():
    """Update or add company facts"""
    data = request.get_json()
    company_name = data.get("company_name")
    facts = data.get("facts")
    
    if not company_name or not facts:
        return jsonify({"error": "Missing 'company_name' or 'facts' in request body"}), 400
    
    # Update or add company facts
    company_facts[company_name] = facts
    
    return jsonify({
        "message": "Facts updated successfully",
        "company": company_name,
        "facts": facts
    })

# Restore the original /facts endpoint
@app.route("/facts", methods=["POST"])
def get_facts():
    """
    Get the best set of facts about a question
    """
    data = request.get_json()
    question = data.get("question")

    if not question:
        return jsonify({"error": "Missing 'question' in request body"}), 400

    # Query ChromaDB for relevant articles
    results = query_chroma(question, n_results=10)

    if "error" in results:
        return jsonify({"error": results["error"]}), 500

    # Process results to create diverse article set
    mcp_results = []
    for doc, score, meta in zip(
            results.get("documents", [[]])[0],
            results.get("distances", [[]])[0],
            results.get("metadatas", [[]])[0]):
        mcp_results.append({
            "document": doc,
            "score": float(score),
            "metadata": meta
        })

    # Apply diversity filtering
    diverse_results = get_diverse_articles(mcp_results, 5)

    # Combine with company facts if question mentions a company
    company_facts_result = {}
    question_lower = question.lower()

    # Check if question mentions any known company
    for company_name in company_facts.keys():
        if company_name.lower() in question_lower:
            company_facts_result = company_facts[company_name]
            break

    # Return combined results
    response_data = {
        "question": question,
        "articles": diverse_results,
        "company_facts": company_facts_result,
        "timestamp": datetime.now().isoformat()
    }

    return jsonify(response_data)

@app.route("/health", methods=["GET"])
def health():
    """Health check endpoint"""
    return jsonify({"status": "healthy", "timestamp": datetime.now().isoformat()})

@app.route("/info", methods=["GET"])
def info():
    """Service information endpoint"""
    return jsonify({
        "provider": "StockDoc",
        "service": "MCP Server",
        "version": "1.0.0",
        "collection": "news",
        "embedding_service": "http://example.com:4000"
    })

@app.route("/openapi.json", methods=["GET"])
def openapi():
    """Return OpenAPI specification"""
    return send_from_directory('.', 'openapi.json')

if __name__ == "__main__":
    app.run(host="0.0.0.0", port=5005, threaded=True, debug=True)