from typing import Dict, List, Optional

import yaml
import time
from dotenv import load_dotenv
from fastapi import FastAPI, HTTPException, Response
from prometheus_client import Histogram, Counter, generate_latest
from prometheus_client.exposition import CONTENT_TYPE_LATEST
from pydantic import BaseModel

from data.loaders import load_prompt_file, make_user_question
from data.processors import get_few_shot_examples
from models.embeddings import EmbeddingService
from models.llm import LLMService
from models.reranker import RerankerService
from pipeline.context import ContextRetriever
from pipeline.generation import TextGenerator
from storage.milvus_client import MilvusConnectionManager, MilvusTableSearch
from utils.logging import setup_logging

from utils.db_logger import db_logger

load_dotenv()

logger = setup_logging()


class GenerationRequest(BaseModel):
    document_type: str
    project: str
    page_name: str
    jira_text: str
    previous_generations: Optional[List[str]] = []
    reference_answer: Optional[str] = ''
    attached_link: Optional[str] = ''

    @property
    def normalized_project_id(self) -> str:
        """Нормализует project_id для использования в Milvus/GraphRAG."""
        return self.project.replace('-', '_')


class GenerationResponse(BaseModel):
    request_id: str
    result: str
    metrics: Dict[str, float]
    timing: Dict[str, float]
    used_summary: bool
    context_documents: List[Dict[str, str]] = [] 


def load_config():
    with open('config/default_config.yaml', 'r', encoding='utf-8') as f:
        return yaml.safe_load(f)


config = load_config()

dataset = load_prompt_file(config.get('dataset_path'))

milvus_manager = MilvusConnectionManager()
embedding_service = EmbeddingService(milvus_manager)
reranker_service = RerankerService()
llm_service = LLMService()
table_search = MilvusTableSearch(milvus_manager)

context_retriever = ContextRetriever(embedding_service, reranker_service, table_search)
text_generator = TextGenerator(llm_service)

app = FastAPI()

PROCESSING_TIME = Histogram(
    "mlcore_document_processing_seconds",
    "Время выполнения запроса генерации в секундах",
    buckets=(30, 60, 90, 120, 150, 180, 240, 300, 360, 480, 600, 900, 1200)
)
PROCESSING_COUNT = Counter(
    "mlcore_document_processing_count",
    "Количество обработанных файлов",
    ["project_id", "status"]
)


@app.on_event('startup')
async def startup_event():
    milvus_manager.connect()
    load_dotenv()
    db_logger.connect()


@app.on_event('shutdown')
async def shutdown_event():
    milvus_manager.disconnect()
    db_logger.disconnect()


@app.post('/generate', response_model=GenerationResponse)
async def generate_document(request: GenerationRequest) -> GenerationResponse:
    """Генерирует документ на основе запроса."""
    start_time = time.time()
    processing_success = False

    try:    
        request_id = db_logger.log_generation_request(
            document_type=request.document_type,
            project=request.project,
            page_name=request.page_name,
            attached_link=request.attached_link
        )

        normalized_project = request.normalized_project_id

        models_names = config.get('models_names', {})
        models_params = config.get('models_params', {})
        system_message = config.get('system_message', '')

        prompt_data = dataset.get(request.document_type, {})
        if not prompt_data:
            raise HTTPException(status_code=404, detail=f"Document type '{request.document_type}' not found")

        template = prompt_data.get('template', '')

        user_question = make_user_question(request.document_type, request.page_name, request.jira_text)

        example_project = config.get('project_params', {}).get('example_project', '')

        if request.page_name == 'Доработка \"Назначение заместителей\"':
            example_page = 'Интеграция со справочником \"Продуктовые линейки\"'
        else:
            example_page = 'Доработка \"Назначение заместителей\"'

        examples = get_few_shot_examples(
            prompt_data=prompt_data,
            example_project=example_project,
            example_page=example_page,
            is_few_shot=config.get('gen_params', {}).get('is_few_shot', True))

        context, timing_data, top_page_infos = context_retriever.retrieve_full_context(
            query=user_question,
            template=template,
            models_names=models_names,
            models_params=models_params,
            reference_project=normalized_project,
            attached_link=request.attached_link,
            is_context=config.get('gen_params', {}).get('is_context', True),
            is_table=config.get('gen_params', {}).get('is_table', True),
            is_graph=config.get('gen_params', {}).get('is_graph', True),
            request_id=request_id
        )

        previous_generations = request.previous_generations
        previous_summary = text_generator.collect_previous_generations(
            all_answers=previous_generations,
            model_name=models_names.get('LLM'),
            is_summary=bool(previous_generations),
        )
        
        db_logger.log_generation_inputs(
            request_id=request_id,
            user_question=user_question,
            jira_text=request.jira_text,
            few_shot_examples=examples,
            previous_generations_summary=previous_summary
        )

        context_documents = [{"page_link": page_link, "page_title": page_title} 
                           for page_link, page_title in top_page_infos]

        answer, metrics, generation_time, _ = text_generator.generate(
            user_question=user_question,
            template=template,
            context=context,
            system_message=system_message,
            examples=examples,
            model_name=models_names.get('LLM'),
            num_generations=1,
            reference_answer=request.reference_answer,
            previous_generations=previous_summary
        )

        timing_data['generation'] = generation_time

        db_logger.log_generation_result(
            request_id=request_id,
            generated_answer=answer,
            timings=timing_data,
            metrics=metrics,
        )

        processing_success = True

        return GenerationResponse(
            request_id=request_id,
            result=answer,
            metrics=metrics,
            timing=timing_data,
            used_summary=bool(previous_generations),
            context_documents=context_documents
        )

    except HTTPException:
        raise
    except Exception as e:
        logger.error(f'Error during document generation: {str(e)}')
        raise HTTPException(status_code=500, detail=f'Error during document generation: {str(e)}')
    
    finally:
        PROCESSING_TIME.observe(time.time() - start_time)
        status = "success" if processing_success else "error"
        PROCESSING_COUNT.labels(project_id=request.project, status=status).inc()

@app.get("/metrics")
async def metrics():
    """Эндпоинт для экспорта метрик Prometheus."""
    return Response(generate_latest(), media_type=CONTENT_TYPE_LATEST)
