import json
import traceback

from fastapi import APIRouter, HTTPException
from pymilvus import connections, Collection, FieldSchema, CollectionSchema, \
    DataType, utility

from app.config import logger
from app.models import ImportRequest, ImportResponse, \
    AvailableCollectionsResponse
from app.routers.utils.generate_embeddings import generate_embeddings
from app.routers.utils.read_content_from_minio import read_content_from_minio

import_router = APIRouter(prefix="/import", tags=["Import"])

@import_router.post("/", response_model=ImportResponse)
async def import_data(request: ImportRequest):
    try:
        logger.info(f"Starting import process for {request.source_type}")

        # Check Milvus connection
        try:
            connections.connect("default", host="milvus", port="19530")
            logger.info("Successfully connected to Milvus")
        except Exception as e:
            logger.error(f"Failed to connect to Milvus: {str(e)}")
            return ImportResponse(status="error",
                                  message="Failed to connect to Milvus")

        # Fetch data from MinIO
        try:
            data = read_content_from_minio(request.object_name)
            logger.info(
                f"Successfully fetched data from MinIO: {request.object_name}")
        except Exception as e:
            logger.error(f"Failed to fetch data from MinIO: {str(e)}")
            return ImportResponse(status="error",
                                  message="Failed to fetch data from MinIO")

        # Process data
        processed_data = json.loads(data)
        logger.info("Data processed successfully")

        # Generate embeddings and insert into Milvus
        collection_name = f"{request.source_type}_collection"
        if not utility.has_collection(collection_name):
            create_collection(collection_name)

        collection = Collection(collection_name)

        total_items = len(processed_data)
        for i, item in enumerate(processed_data, 1):
            try:
                item["embedding"] = generate_embeddings(item)
                filtered_item = {
                    "content": item.get("content", ""),
                    "embedding": item["embedding"],
                    "creation_timestamp": int(
                        item.get("creation_timestamp", 0)),
                    "index": item.get("index", ""),
                    "type": item.get("type", ""),
                    "network": item.get("network", ""),
                    "url": item.get("url", "")
                }
                _ = collection.insert([filtered_item])
                logger.info(
                    f"Inserted item {i}/{total_items} into Milvus collection {collection_name}")
            except Exception as e:
                logger.error(f"Failed to process item {i}: {str(e)}")

        logger.info(f"Import completed for {request.source_type}")
        return ImportResponse(status="success",
                              message="Import completed successfully")

    except Exception as e:
        logger.error(f"Unexpected error during import: {str(e)}")
        logger.error(traceback.format_exc())
        return ImportResponse(status="error",
                              message=f"Unexpected error: {str(e)}")


def create_collection(collection_name: str):
    fields = [
        FieldSchema(name="id", dtype=DataType.INT64, is_primary=True,
                    auto_id=True),
        FieldSchema(name="content", dtype=DataType.VARCHAR, max_length=65535),
        FieldSchema(name="embedding", dtype=DataType.FLOAT_VECTOR, dim=1024),
        FieldSchema(name="creation_timestamp", dtype=DataType.INT64),
        FieldSchema(name="index", dtype=DataType.VARCHAR, max_length=255),
        FieldSchema(name="type", dtype=DataType.VARCHAR, max_length=255),
        FieldSchema(name="network", dtype=DataType.VARCHAR, max_length=255),
        FieldSchema(name="url", dtype=DataType.VARCHAR, max_length=2083),
    ]
    schema = CollectionSchema(fields, "A collection for storing embeddings")
    collection = Collection(collection_name, schema)

    index_params = {
        "metric_type": "L2",
        "index_type": "IVF_FLAT",
        "params": {"nlist": 1024}
    }
    collection.create_index("embedding", index_params)
    logger.info(f"Created new collection: {collection_name}")


@import_router.get("/available_collections",
                   response_model=AvailableCollectionsResponse)
def get_available_collections():
    logger.info("Getting available collections from Milvus")

    try:
        if not utility.has_collection("default"):
            connections.connect("default", host="milvus", port="19530")

        collections = utility.list_collections()

        logger.info(f"Found {len(collections)} collections")
        return AvailableCollectionsResponse(collections=collections)

    except Exception as e:
        logger.error(f"Error getting collections from Milvus: {str(e)}")
        raise HTTPException(status_code=500,
                            detail=f"Error getting collections from Milvus: {str(e)}")
