from fastapi import APIRouter, HTTPException, Request
from fastapi.responses import StreamingResponse
from pydantic import BaseModel, Field
from typing import Optional
import logging
import uuid
import os
import datetime
import json
import config
import asyncio
import pydicom
from pydicom.uid import generate_uid
from pydicom.dataset import Dataset
from utils import get_redis_client, prepareDataset, saveFile


logger = logging.getLogger("sampacs")

router = APIRouter(prefix="/api/temp-capture", tags=["temp-capture"])

class TempCaptureStartRequest(BaseModel):
    """Request model for starting temporary capture session"""
    equipment_aet: str = Field(..., description="AETitle оборудования")

class TempCaptureStartResponse(BaseModel):
    """Response model for temporary capture session start"""
    session_id: str = Field(..., description="ID сессии для временного захвата")

class SeriesPreviewEvent(BaseModel):
    """Event model for series preview"""
    session_id: str = Field(..., description="ID сессии")
    study_uid: str = Field(..., description="Study Instance UID")
    series_uid: str = Field(..., description="Series Instance UID")
    series_description: Optional[str] = Field(None, description="Описание серии")
    modality: Optional[str] = Field(None, description="Модальность")
    preview_url: str = Field(..., description="URL превью изображения")
    instance_count: int = Field(..., description="Количество инстансов в серии")
    timestamp: datetime.datetime = Field(..., description="Время создания события")

class PatientInfo(BaseModel):
    """Patient information for confirmation"""
    pat_name: str = Field(..., description="ФИО пациента")
    birth_date: Optional[datetime.datetime] = Field(None, description="Дата рождения")
    sex: Optional[str] = Field(None, description="Пол (F-женский, M-мужской)")
    patient_id: str = Field(..., description="ID пациента в системе")

class TempCaptureConfirmRequest(BaseModel):
    """Request model for confirming temporary capture"""
    patient_info: PatientInfo = Field(..., description="Информация о пациенте")
    accession_number: str = Field(..., description="Accession Number")
    req_desc: str = Field(..., description="Описание запрошенной процедуры")

class TempCaptureConfirmResponse(BaseModel):
    """Response model for confirmation"""
    study_uid: str = Field(..., description="Study Instance UID")

class TempCaptureCancelResponse(BaseModel):
    """Response model for cancellation"""
    success: bool = Field(..., description="Успешность операции")

@router.post(
    "/start",
    summary="Начать сессию временного захвата DICOM",
    response_description="Информация о созданной сессии",
    responses={
        200: {
            "description": "Сессия успешно создана",
            "content": {
                "application/json": {
                    "example": {
                        "session_id": "550e8400-e29b-41d4-a716-446655440000"
                    }
                }
            }
        },
        400: {"description": "Неверные параметры запроса"},
        500: {"description": "Внутренняя ошибка сервера"}
    }
)
async def start_temp_capture(
    request: TempCaptureStartRequest
) -> TempCaptureStartResponse:
    """
    Начать новую сессию временного захвата DICOM файлов.

    Создает временную сессию для захвата DICOM файлов с указанного оборудования.
    Файлы будут сохраняться во временную папку до подтверждения или отмены.

    ## Параметры:
    - **equipment_aet**: AETitle оборудования, с которого будут поступать DICOM файлы

    ## Возвращает:
    - **session_id**: Уникальный идентификатор сессии

    ## Ошибки:
    - 400: Не указан AETitle оборудования
    - 500: Внутренняя ошибка при создании сессии
    """
    try:
        # Validate input
        if not request.equipment_aet or not request.equipment_aet.strip():
            raise HTTPException(
                status_code=400,
                detail="Equipment AETitle is required"
            )

        # Use Redis for inter-process communication
        try:
            redis_client = get_redis_client()
        except Exception as e:
            logger.error(f"Redis connection failed: {str(e)}")
            raise HTTPException(
                status_code=503,
                detail="Redis storage unavailable"
            )

        # Check for existing sessions for this equipment_aet
        existing_sessions = []
        session_keys = redis_client.keys("temp_capture:*")

        for key in session_keys:
            session_json = redis_client.get(key)
            if session_json:
                session_data = json.loads(session_json)
                if session_data["equipment_aet"] == request.equipment_aet.strip():
                    existing_sessions.append((key, session_data))

        # Check for active session
        for key, session_data in existing_sessions:
            if session_data.get("active") == True:
                raise HTTPException(
                    status_code=400,
                    detail=f"Active session already exists for equipment {request.equipment_aet}"
                )

        # Remove inactive sessions
        for key, session_data in existing_sessions:
            try:
                # Remove temp folder if exists
                temp_folder = session_data.get("temp_folder", "")
                if temp_folder and os.path.exists(temp_folder):
                    import shutil
                    shutil.rmtree(temp_folder)

                # Remove session from Redis
                redis_client.delete(key)

                # Remove events stream
                events_key = f"temp_capture_events:{session_data['session_id']}"
                redis_client.delete(events_key)

                logger.info(f"Removed inactive session {session_data['session_id']} for equipment {request.equipment_aet}")
            except Exception as e:
                logger.warning(f"Error removing inactive session {session_data['session_id']}: {str(e)}")

        # Generate unique session ID with collision check
        max_attempts = 5
        session_id = None

        for attempt in range(max_attempts):
            candidate_id = str(uuid.uuid4())
            redis_key = f"temp_capture:{candidate_id}"

            # Check if session_id already exists
            if not redis_client.exists(redis_key):
                session_id = candidate_id
                break

        if not session_id:
            raise HTTPException(
                status_code=500,
                detail="Failed to generate unique session ID"
            )

        # Create temp folder path
        storage_dir = config.config.storageDir
        temp_folder = os.path.join(storage_dir, "tmp", session_id)

        # Create temp directory
        os.makedirs(temp_folder, exist_ok=True)

        # Create subdirectories for organization
        os.makedirs(os.path.join(temp_folder, "dicom"), exist_ok=True)
        os.makedirs(os.path.join(temp_folder, "preview"), exist_ok=True)

        # Store session info in Redis
        session_data = {
            "session_id": session_id,
            "equipment_aet": request.equipment_aet.strip(),
            "temp_folder": temp_folder,
            "created_at": datetime.datetime.now().isoformat(),
            "active": False
        }

        redis_key = f"temp_capture:{session_id}"
        redis_client.setex(
            redis_key,
            3600,  # 1 hour TTL
            json.dumps(session_data)
        )

        logger.info(f"Created temp capture session {session_id} for equipment {request.equipment_aet}")

        return TempCaptureStartResponse(
            session_id=session_id
        )

    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"Error creating temp capture session: {str(e)}", exc_info=True)
        raise HTTPException(
            status_code=500,
            detail="Internal server error while creating session"
        )

@router.get(
    "/stream/{session_id}",
    summary="SSE поток событий временного захвата",
    response_description="Поток событий о полученных DICOM файлах",
    responses={
        200: {
            "description": "SSE поток событий",
            "content": {"text/event-stream": {}}
        },
        404: {"description": "Сессия не найдена"},
        500: {"description": "Внутренняя ошибка сервера"}
    }
)
async def stream_temp_capture_events(
    session_id: str,
    request: Request
):
    """
    SSE поток событий для временного захвата DICOM файлов.

    При подключении активирует сессию. При разрыве соединения деактивирует сессию.
    Отправляет события о полученных DICOM сериях с превью изображениями.

    ## Параметры:
    - **session_id**: ID сессии временного захвата

    ## Возвращает:
    - SSE поток событий в формате JSON

    ## Ошибки:
    - 404: Сессия с указанным ID не найдена
    - 500: Внутренняя ошибка сервера или Redis недоступен
    """
    try:
        redis_client = get_redis_client()
    except Exception as e:
        logger.error(f"Redis connection failed: {str(e)}")
        raise HTTPException(
            status_code=503,
            detail="Redis storage unavailable"
        )

    # Check if session exists
    redis_key = f"temp_capture:{session_id}"
    session_json = redis_client.get(redis_key)

    if not session_json:
        raise HTTPException(
            status_code=404,
            detail=f"Session {session_id} not found"
        )

    # Parse session data
    session_data = json.loads(session_json)

    # Activate session
    session_data["active"] = True
    redis_client.setex(
        redis_key,
        3600,  # Refresh TTL
        json.dumps(session_data)
    )
    logger.info(f"Session {session_id} activated via SSE connection")

    async def event_generator():
        """Generate SSE events from Redis Stream"""
        last_id = "0"  # Start from beginning

        try:
            while True:
                # Check if client is still connected
                if await request.is_disconnected():
                    logger.info(f"SSE client disconnected for session {session_id}")
                    break

                # Read events from Redis Stream
                try:
                    # Redis Stream key for DICOM events
                    stream_key = f"temp_capture_events:{session_id}"

                    # Read events (block for 5 seconds)
                    events = redis_client.xread(
                        {stream_key: last_id},
                        count=10,
                        block=5000
                    )

                    if events:
                        for stream, messages in events:
                            for message_id, message_data in messages:
                                # Filter events for this session
                                if message_data.get(b"session_id", b"").decode() == session_id:
                                    # Create event object
                                    event = SeriesPreviewEvent(
                                        session_id=session_id,
                                        study_uid=message_data.get(b"study_uid", b"").decode(),
                                        series_uid=message_data.get(b"series_uid", b"").decode(),
                                        series_description=message_data.get(b"series_description", b"").decode() or None,
                                        modality=message_data.get(b"modality", b"").decode() or None,
                                        preview_url=message_data.get(b"preview_url", b"").decode(),
                                        instance_count=int(message_data.get(b"instance_count", b"1")),
                                        timestamp=datetime.datetime.fromisoformat(
                                            message_data.get(b"timestamp", b"").decode()
                                        )
                                    )

                                    # Send as SSE event
                                    event_data = {
                                        "event": "series_preview",
                                        "data": json.loads(event.json())
                                    }
                                    yield f"data: {json.dumps(event_data, ensure_ascii=False)}\n\n"

                                last_id = message_id

                except Exception as e:
                    logger.error(f"Error reading from Redis Stream: {str(e)}")

                # Small delay to prevent busy waiting
                await asyncio.sleep(0.1)

        except asyncio.CancelledError:
            logger.info(f"SSE generator cancelled for session {session_id}")
            raise
        except Exception as e:
            logger.error(f"Error in SSE generator for session {session_id}: {str(e)}")
            # Send error event
            error_data = {
                "event": "error",
                "data": {
                    "message": f"Stream error: {str(e)}",
                    "session_id": session_id
                }
            }
            yield f"data: {json.dumps(error_data, ensure_ascii=False)}\n\n"

        finally:
            # Deactivate session on disconnect
            try:
                session_json = redis_client.get(redis_key)
                if session_json:
                    session_data = json.loads(session_json)
                    session_data["active"] = False
                    redis_client.setex(
                        redis_key,
                        3600,  # Keep TTL
                        json.dumps(session_data)
                    )
                    logger.info(f"Session {session_id} deactivated after SSE disconnect")
            except Exception as e:
                logger.error(f"Error deactivating session {session_id}: {str(e)}")

    return StreamingResponse(
        event_generator(),
        media_type="text/event-stream",
        headers={
            "Cache-Control": "no-cache",
            "Connection": "keep-alive",
            "X-Accel-Buffering": "no"
        }
    )

@router.post(
    "/confirm/{session_id}",
    summary="Подтвердить сохранение временно захваченных DICOM файлов",
    response_description="Study UID сохраненного исследования",
    responses={
        200: {
            "description": "Файлы успешно сохранены",
            "content": {
                "application/json": {
                    "example": {
                        "study_uid": "1.2.840.113619.2.5.1762583153.215519.978957063.78"
                    }
                }
            }
        },
        400: {"description": "Неверные параметры запроса"},
        404: {"description": "Сессия не найдена или нет файлов"},
        500: {"description": "Внутренняя ошибка сервера"}
    }
)
async def confirm_temp_capture(
    session_id: str,
    request: TempCaptureConfirmRequest
) -> TempCaptureConfirmResponse:
    """
    Подтвердить сохранение временно захваченных DICOM файлов.

    Заменяет идентификаторы пациента и accession number во всех DICOM файлах сессии
    и сохраняет их в основное хранилище через стандартные функции PACS.

    ## Параметры:
    - **session_id**: ID сессии временного захвата
    - **patient_info**: Информация о пациенте (pat_name в формате 'Фамилия^Имя^Отчество', ID, пол, дата рождения)
    - **accession_number**: Accession Number для исследования
    - **req_desc**: Описание запрошенной процедуры

    ## Возвращает:
    - **study_uid**: Study Instance UID сохраненного исследования

    ## Ошибки:
    - 400: Не указаны обязательные параметры или нет файлов для сохранения
    - 404: Сессия не найдена
    - 500: Внутренняя ошибка при сохранении файлов
    """
    try:
        redis_client = get_redis_client()
    except Exception as e:
        logger.error(f"Redis connection failed: {str(e)}")
        raise HTTPException(
            status_code=500,
            detail="Redis storage unavailable"
        )

    # Check if session exists
    redis_key = f"temp_capture:{session_id}"
    session_json = redis_client.get(redis_key)

    if not session_json:
        raise HTTPException(
            status_code=404,
            detail=f"Session {session_id} not found"
        )

    # Parse session data
    session_data = json.loads(session_json)
    temp_folder = session_data["temp_folder"]
    equipment_aet = session_data["equipment_aet"]

    # Check if temp folder exists and has files
    dicom_folder = os.path.join(temp_folder, "dicom")
    if not os.path.exists(dicom_folder):
        raise HTTPException(
            status_code=400,
            detail="No DICOM files found in session"
        )

    # Find all DICOM files
    dicom_files = []
    for root, dirs, files in os.walk(dicom_folder):
        for file in files:
            if file.lower().endswith('.dcm'):
                dicom_files.append(os.path.join(root, file))

    logger.info(f"Found {len(dicom_files)} DICOM files to process in session {session_id}")

    if not dicom_files:
        raise HTTPException(
            status_code=400,
            detail="No DICOM files found in session"
        )

    # Get modality config for prepareDataset
    modality_config = None
    for modality in config.config.modalities:
        if modality.name == equipment_aet or modality.aetitle == equipment_aet:
            modality_config = modality
            break

    # Process DICOM files
    generated_study_uid = generate_uid()
    files_processed = 0
    files_failed = 0

    for dicom_file in dicom_files:
        try:
            # Load DICOM file
            ds = pydicom.dcmread(dicom_file, force=True)

            # Replace StudyInstanceUID with generated one
            ds.StudyInstanceUID = generated_study_uid

            # Replace patient information
            # Patient ID (required)
            ds.PatientID = request.patient_info.patient_id

            # Patient Name (required)
            ds.PatientName = request.patient_info.pat_name

            # Accession Number (required)
            ds.AccessionNumber = request.accession_number

            # Optional fields
            if request.patient_info.birth_date:
                ds.PatientBirthDate = request.patient_info.birth_date.strftime("%Y%m%d")

            if request.patient_info.sex is not None:
                ds.PatientSex = request.patient_info.sex

            # Replace study description with requested procedure description
            if request.req_desc:
                ds.StudyDescription = request.req_desc

            # Prepare dataset using existing utility function

            (prepared_ds, error) = prepareDataset(
                ds,
                registrateOptions=config.config.registrateOptions,
                modality=modality_config
            )

            if error:
                logger.error(f"Error preparing dataset {dicom_file}: {error}")
                files_failed += 1
                continue

            # Save file using existing utility function
            saveFile(prepared_ds, equipment_aet, saveFile=True)
            files_processed += 1

        except Exception as e:
            logger.error(f"Error processing file {dicom_file}: {str(e)}", exc_info=True)
            files_failed += 1

    if files_processed == 0:
        raise HTTPException(
            status_code=500,
            detail=f"Failed to process any files. Success: {files_processed}, Failed: {files_failed}"
        )

    if files_failed > 0:
        logger.warning(f"Some files failed to process. Success: {files_processed}, Failed: {files_failed}")

    # Cleanup
    try:
        # Remove temp folder
        import shutil
        shutil.rmtree(temp_folder)

        # Remove session from Redis
        redis_client.delete(redis_key)

        # Remove events stream
        events_key = f"temp_capture_events:{session_id}"
        redis_client.delete(events_key)

        logger.info(f"Session {session_id} confirmed and cleaned up. Files processed: {files_processed}")

    except Exception as e:
        logger.error(f"Error cleaning up session {session_id}: {str(e)}")
        # Don't fail the request if cleanup fails

    if not generated_study_uid:
        raise HTTPException(
            status_code=500,
            detail="Failed to determine study UID"
        )

    return TempCaptureConfirmResponse(
        study_uid=generated_study_uid
    )

@router.post(
    "/cancel/{session_id}",
    summary="Отменить сессию временного захвата DICOM",
    response_description="Результат отмены сессии",
    responses={
        200: {
            "description": "Сессия успешно отменена",
            "content": {
                "application/json": {
                    "example": {
                        "success": True,
                        "message": "Сессия отменена, временные файлы удалены",
                        "session_id": "550e8400-e29b-41d4-a716-446655440000"
                    }
                }
            }
        },
        404: {"description": "Сессия не найдена"},
        500: {"description": "Внутренняя ошибка сервера"}
    }
)
async def cancel_temp_capture(
    session_id: str
) -> TempCaptureCancelResponse:
    """
    Отменить сессию временного захвата DICOM файлов.

    Удаляет все временные файлы и очищает данные сессии из Redis без сохранения.

    ## Параметры:
    - **session_id**: ID сессии временного захвата

    ## Возвращает:
    - **success**: Успешность операции

    ## Ошибки:
    - 404: Сессия не найдена
    - 500: Внутренняя ошибка при отмене сессии
    """
    try:
        redis_client = get_redis_client()
    except Exception as e:
        logger.error(f"Redis connection failed: {str(e)}")
        raise HTTPException(
            status_code=500,
            detail="Redis storage unavailable"
        )

    # Check if session exists
    redis_key = f"temp_capture:{session_id}"
    session_json = redis_client.get(redis_key)

    if not session_json:
        raise HTTPException(
            status_code=404,
            detail=f"Session {session_id} not found"
        )

    # Parse session data to get temp folder
    session_data = json.loads(session_json)
    temp_folder = session_data.get("temp_folder", "")

    # Cleanup
    try:
        # Remove temp folder if it exists
        if temp_folder and os.path.exists(temp_folder):
            import shutil
            shutil.rmtree(temp_folder)
            logger.info(f"Removed temp folder: {temp_folder}")

        # Remove session from Redis
        redis_client.delete(redis_key)

        # Remove events stream
        events_key = f"temp_capture_events:{session_id}"
        redis_client.delete(events_key)

        logger.info(f"Session {session_id} cancelled and cleaned up")

        return TempCaptureCancelResponse(
            success=True
        )

    except Exception as e:
        logger.error(f"Error cancelling session {session_id}: {str(e)}", exc_info=True)
        raise HTTPException(
            status_code=500,
            detail=f"Internal server error while cancelling session: {str(e)}"
        )
