# from fastapi import FastAPI, HTTPException, status
# from pydantic import BaseModel
# from pymongo import MongoClient
# from pymongo.errors import PyMongoError
# from typing import Dict
# from fastapi.responses import JSONResponse
# from typing import Optional
# import re

# # Initialize FastAPI
# app = FastAPI(title="Multi-Source Company Screener API")

# # MongoDB connection
# # client = MongoClient("mongodb://jenya:DJenya$Mongo%40St0ckDB@172.105.59.175:27017/")
# client = MongoClient("mongodb://localhost:27017/")
# db = client["test2"]
# tijori_collection = db["tijori_data"]
# investing_collection = db["investing_data"]
# screener_collection = db["screener_data"]
# moneycontrol_collection = db["money_data"]

# # Helper to clean MongoDB documents
# def clean_document(doc):
#     if doc is None:
#         return {}
#     doc = dict(doc)
#     doc['_id'] = str(doc['_id'])  # Convert ObjectId to string
#     return doc

# # Request model for POST
# class CompanyRequest(BaseModel):
#     # symbol: str
#     # company_name: str
#     # collection_type: str
#     symbol: Optional[str] = None
#     company_name: Optional[str] = None
#     collection_type: str

# # @app.post("/companies", response_model=Dict[str, object])
# # async def post_company_data(request: CompanyRequest):

# #     try:
# #         symbol = request.symbol.upper()
# #         collection_type = request.collection_type.lower()

# #         if collection_type == "tijori":
# #             data = clean_document(tijori_collection.find_one({"symbol": symbol}))
# #             return {"tijori_data": data}
# #         elif collection_type == "investing":
# #             data = clean_document(investing_collection.find_one({"symbol": symbol}))
# #             return {"investing_data": data}
# #         elif collection_type == "screener":
# #             data = clean_document(screener_collection.find_one({"symbol": symbol}))
# #             return {"screener_data": data}
# #         elif collection_type == "moneycontrol":
# #             data = clean_document(moneycontrol_collection.find_one({"symbol": symbol}))
# #             return {"moneycontrol_data": data}
# #         else:
# #             raise HTTPException(status_code=400, detail="Invalid collection type.")

# #     except PyMongoError as e:
# #         raise HTTPException(
# #             status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
# #             detail=f"Database error: {str(e)}"
# #         )


# @app.post("/companies", response_model=Dict[str, object])
# async def post_company_data(request: CompanyRequest):
#     try:
#         collection_type = request.collection_type.lower()
#         symbol = request.symbol.upper() if request.symbol else None
#         company_name = request.company_name.upper() if request.company_name else None

#         # If both symbol and company_name are missing
#         if not symbol and not company_name:
#             raise HTTPException(status_code=400, detail="Please provide either 'symbol' or 'company_name'.")

#         # Prepare query based on availability
#         query = {}
#         if symbol:
#             query["symbol"] = {"$regex": f"^{symbol}$", "$options": "i"}
#         if company_name:
#             query["company_name"] = {"$regex": f"^{company_name}$", "$options": "i"}

#         # Select collection and query
#         if collection_type == "tijori":
#             data = clean_document(tijori_collection.find_one(query))
#             return {"tijori_data": data}
#         elif collection_type == "investing":
#             # Investing collection might not use company_name
#             if "symbol" in query:
#                 data = clean_document(investing_collection.find_one({"symbol": symbol}))
#             elif "company_name" in query:
#                 data = clean_document(investing_collection.find_one({"company_name": company_name}))
#             else:
#                 data = None
#             return {"investing_data": data}
#         elif collection_type == "screener":
#             data = clean_document(screener_collection.find_one(query))
#             return {"screener_data": data}
#         elif collection_type == "moneycontrol":
#             if "symbol" in query:
#                 data = clean_document(moneycontrol_collection.find_one({"symbol": symbol}))
#             elif "company_name" in query:
#                 data = clean_document(moneycontrol_collection.find_one({"company_name": company_name}))
#             else:
#                 data = None
#             return {"moneycontrol_data": data}
#         else:
#             raise HTTPException(status_code=400, detail="Invalid collection type.")

#     except PyMongoError as e:
#         raise HTTPException(
#             status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
#             detail=f"Database error: {str(e)}"
#         )
# # Note: This block is not needed when using uvicorn to run the app
# if __name__ == "__main__":
#     import uvicorn
#     uvicorn.run(app, host="0.0.0.0", port=8001, reload=True)




from fastapi import FastAPI, HTTPException, status
from pydantic import BaseModel
from pymongo import MongoClient
from pymongo.errors import PyMongoError
from typing import Dict, Optional
from fastapi.responses import JSONResponse
import re

# Initialize FastAPI
app = FastAPI(title="Multi-Source Company Screener API")

# MongoDB connection
client = MongoClient("mongodb://jenya:DJenya$Mongo%40St0ckDB@172.105.59.175:27017/")
db = client["trade_iq"]
tijori_collection = db["tijori_data"]
investing_collection = db["investing_data"]
screener_collection = db["screener_data"]
moneycontrol_collection = db["moneycontrol_data"]

# Helper to clean MongoDB documents
def clean_document(doc):
    if doc is None:
        return {}
    doc = dict(doc)
    doc['_id'] = str(doc['_id'])  # Convert ObjectId to string
    return doc

# Request model for POST
class CompanyRequest(BaseModel):
    symbol: Optional[str] = None
    company_name: Optional[str] = None
    collection_type: str

@app.post("/companies", response_model=Dict[str, object])
async def post_company_data(request: CompanyRequest):
    try:
        collection_type = request.collection_type.lower()
        symbol = request.symbol.strip().upper() if request.symbol else None
        company_name = request.company_name.strip().upper() if request.company_name else None

        # If both symbol and company_name are missing
        if not symbol and not company_name:
            raise HTTPException(status_code=400, detail="Please provide either 'symbol' or 'company_name'.")

        # Prepare query with flexible regex
        query = {}
        if symbol:
            query["symbol"] = {"$regex": re.escape(symbol), "$options": "i"}
        if company_name:
            query["company_name"] = {"$regex": re.escape(company_name), "$options": "i"}

        # Debug print (optional)
        print(f"MongoDB Query: {query}")

        # Select collection and query
        if collection_type == "tijori":
            if "symbol" in query:
                data = clean_document(tijori_collection.find_one({"symbol": {"$regex": re.escape(symbol), "$options": "i"}}))
            elif "company_name" in query:
                data = clean_document(tijori_collection.find_one({"company_name": {"$regex": re.escape(company_name), "$options": "i"}}))
            else:
                data = {}
            return {"tijori_data": data}

        elif collection_type == "investing":
            if "symbol" in query:
                data = clean_document(investing_collection.find_one({"symbol": {"$regex": re.escape(symbol), "$options": "i"}}))
            elif "company_name" in query:
                data = clean_document(investing_collection.find_one({"company_name": {"$regex": re.escape(company_name), "$options": "i"}}))
            else:
                data = {}
            return {"investing_data": data}

        elif collection_type == "screener":
            if "symbol" in query:
                data = clean_document(screener_collection.find_one({"symbol": {"$regex": re.escape(symbol), "$options": "i"}}))
            elif "company_name" in query:
                data = clean_document(screener_collection.find_one({"company_name": {"$regex": re.escape(company_name), "$options": "i"}}))
            else:
                data = {}
            return {"screener_data": data}

        elif collection_type == "moneycontrol":
            if "symbol" in query:
                data = clean_document(moneycontrol_collection.find_one({"symbol": {"$regex": re.escape(symbol), "$options": "i"}}))
            elif "company_name" in query:
                data = clean_document(moneycontrol_collection.find_one({"company_name": {"$regex": re.escape(company_name), "$options": "i"}}))
            else:
                data = {}
            return {"moneycontrol_data": data}

        else:
            raise HTTPException(status_code=400, detail="Invalid collection type.")

    except PyMongoError as e:
        raise HTTPException(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
            detail=f"Database error: {str(e)}"
        )

# Only for running with `python filename.py`
if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8001, reload=True)
