from fastapi import APIRouter, Query, HTTPException, Depends, UploadFile, File
from typing import List, Dict, Optional
from pydantic import BaseModel, Field
import datetime
from dbmodel import Worklist, pacs_db
from ..dependencies import get_db
from peewee import DoesNotExist
import logging
import tempfile
import os
import zipfile
import uuid
from pydicom import dcmread, Dataset
from pydicom.uid import generate_uid
import shutil
import config
from utils import saveFile, saveIndex, prepareDataset

router = APIRouter(prefix="/api/worklist", tags=["worklist"])

logger = logging.getLogger("sampacs")

class WorklistItem(BaseModel):
    id: int = Field(description="ID записи worklist")
    pat_id: str = Field(description="Patient ID")
    pat_name: str = Field(description="Patient name")
    pat_birth: datetime.date = Field(description="Patient birth date")
    pat_sex: str = Field(description="Patient sex (M/F/O)")
    acc_num: Optional[str] = Field(None, description="Accession number")
    modality: Optional[str] = Field(None, description="Modality")
    station_aet: Optional[str] = Field(None, description="Scheduled station AETitle")
    sched_date: datetime.date = Field(description="Scheduled date")
    sched_time: Optional[datetime.time] = Field(None, description="Scheduled time")
    req_desc: Optional[str] = Field(None, description="Requested procedure description")
    study_uid: Optional[str] = Field(None, description="Study instance UID")
    created_at: datetime.datetime = Field(description="Дата создания записи")
    updated_at: datetime.datetime = Field(description="Дата обновления записи")

    class Config:
        from_attributes = True

class WorklistCreate(BaseModel):
    pat_id: str = Field(..., min_length=1, max_length=64)
    pat_name: str = Field(..., min_length=1, max_length=255)
    pat_birth: datetime.date = Field(...)
    pat_sex: str = Field(..., min_length=1, max_length=1)
    acc_num: Optional[str] = Field(None, max_length=64)
    modality: Optional[str] = Field(None, max_length=16)
    station_aet: Optional[str] = Field(None, max_length=16)
    sched_date: datetime.date = Field(...)
    sched_time: Optional[datetime.time] = Field(None)
    req_desc: Optional[str] = Field(None, max_length=255)

    class Config:
        json_schema_extra = {
            "example": {
                "pat_id": "12345",
                "pat_name": "Иванов Иван Иванович",
                "pat_birth": "1980-01-15",
                "pat_sex": "M",
                "acc_num": "ACC20240001",
                "modality": "CT",
                "station_aet": "CT01",
                "sched_date": "2024-01-20",
                "sched_time": "14:30:00",
                "req_desc": "Контрольное исследование"
            }
        }

class WorklistUpdate(BaseModel):
    pat_id: Optional[str] = Field(None, min_length=1, max_length=64)
    pat_name: Optional[str] = Field(None, min_length=1, max_length=255)
    pat_birth: Optional[datetime.date] = Field(None)
    pat_sex: Optional[str] = Field(None, min_length=1, max_length=1)
    acc_num: Optional[str] = Field(None, max_length=64)
    modality: Optional[str] = Field(None, max_length=16)
    station_aet: Optional[str] = Field(None, max_length=16)
    sched_date: Optional[datetime.date] = Field(None)
    sched_time: Optional[datetime.time] = Field(None)
    req_desc: Optional[str] = Field(None, max_length=255)
    study_uid: Optional[str] = Field(None, max_length=64)

    class Config:
        json_schema_extra = {
            "example": {
                "sched_time": "15:00:00",
                "study_uid": "1.2.840.113619.2.404.3.1234567890.12345.1234567890"
            }
        }

@router.get("", summary="Получить worklist исследований")
async def get_worklist(
    station_aet: Optional[str] = Query(None, description="Filter by modality AETitle"),
    modality: Optional[str] = Query(None, description="Filter by modality"),
    from_date: Optional[datetime.date] = Query(None, description="Filter by scheduled date from"),
    till_date: Optional[datetime.date] = Query(None, description="Filter by scheduled date to"),
    pat_name: Optional[str] = Query(None, description="Filter by patient name"),
    pat_id: Optional[str] = Query(None, description="Filter by patient ID"),
    completed: Optional[bool] = Query(None, description="Filter by completion status"),
    limit: Optional[int] = Query(100, ge=1, le=1000),
    offset: Optional[int] = Query(0, ge=0),
    db: None = Depends(get_db)
) -> List[WorklistItem]:
    try:
        query = Worklist.select()

        if station_aet:
            aet_list = [aet.strip() for aet in station_aet.split(',') if aet.strip()]
            if aet_list:
                query = query.where(
                    (Worklist.station_aet.in_(aet_list)) | (Worklist.station_aet.is_null())
                )

        if modality:
            query = query.where(Worklist.modality == modality)

        if from_date:
            query = query.where(Worklist.sched_date >= from_date)

        if till_date:
            query = query.where(Worklist.sched_date <= till_date)

        if pat_name:
            search_pattern = pat_name.replace('*', '%')
            if '%' not in search_pattern:
                search_pattern = f'%{search_pattern}%'
            query = query.where(Worklist.pat_name ** search_pattern)

        if pat_id:
            query = query.where(Worklist.pat_id == pat_id)

        if completed is not None:
            if completed:
                query = query.where(Worklist.study_uid.is_null(False))
            else:
                query = query.where(Worklist.study_uid.is_null())

        query = query.order_by(Worklist.sched_date, Worklist.sched_time, Worklist.pat_name)
        query = query.limit(limit).offset(offset)

        worklist_items = list(query)

        return [WorklistItem.from_orm(item) for item in worklist_items]

    except Exception as e:
        logger.error(f"Error getting worklist: {e}")
        raise HTTPException(status_code=500, detail="Internal server error")

@router.post("", summary="Создать запись worklist")
async def create_worklist(
    worklist_data: WorklistCreate,
    db: None = Depends(get_db)
) -> WorklistItem:
    try:
        if worklist_data.pat_sex.upper() not in ['M', 'F', 'O']:
            raise HTTPException(status_code=400, detail="pat_sex must be one of: M, F, O")

        worklist_item = Worklist(
            pat_id=worklist_data.pat_id,
            pat_name=worklist_data.pat_name,
            pat_birth=worklist_data.pat_birth,
            pat_sex=worklist_data.pat_sex.upper(),
            acc_num=worklist_data.acc_num,
            modality=worklist_data.modality,
            station_aet=worklist_data.station_aet,
            sched_date=worklist_data.sched_date,
            sched_time=worklist_data.sched_time,
            req_desc=worklist_data.req_desc,
            study_uid=None
        )

        worklist_item.save()

        logger.info(f"Worklist item created: ID={worklist_item.id}, Patient={worklist_data.pat_name}")

        return WorklistItem.from_orm(worklist_item)

    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"Error creating worklist item: {e}")
        raise HTTPException(status_code=500, detail="Internal server error")

@router.get("/{item_id}", summary="Получить запись worklist по ID")
async def get_worklist_item(
    item_id: int,
    db: None = Depends(get_db)
) -> WorklistItem:
    try:
        worklist_item = Worklist.get_by_id(item_id)
        return WorklistItem.from_orm(worklist_item)

    except DoesNotExist:
        raise HTTPException(status_code=404, detail="Worklist item not found")
    except Exception as e:
        logger.error(f"Error getting worklist item {item_id}: {e}")
        raise HTTPException(status_code=500, detail="Internal server error")

@router.put("/{item_id}", summary="Обновить запись worklist")
async def update_worklist_item(
    item_id: int,
    worklist_data: WorklistUpdate,
    db: None = Depends(get_db)
) -> WorklistItem:
    try:
        worklist_item = Worklist.get_by_id(item_id)

        update_data = worklist_data.dict(exclude_unset=True)

        if 'pat_sex' in update_data:
            if update_data['pat_sex'].upper() not in ['M', 'F', 'O']:
                raise HTTPException(status_code=400, detail="pat_sex must be one of: M, F, O")
            update_data['pat_sex'] = update_data['pat_sex'].upper()

        for field, value in update_data.items():
            setattr(worklist_item, field, value)

        worklist_item.save()

        logger.info(f"Worklist item updated: ID={item_id}")

        return WorklistItem.from_orm(worklist_item)

    except DoesNotExist:
        raise HTTPException(status_code=404, detail="Worklist item not found")
    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"Error updating worklist item {item_id}: {e}")
        raise HTTPException(status_code=500, detail="Internal server error")

@router.delete("/{item_id}", summary="Удалить запись worklist")
async def delete_worklist_item(
    item_id: int,
    db: None = Depends(get_db)
) -> Dict[str, bool]:
    try:
        worklist_item = Worklist.get_by_id(item_id)
        worklist_item.delete_instance()

        logger.info(f"Worklist item deleted: ID={item_id}")

        return {"success": True}

    except DoesNotExist:
        raise HTTPException(status_code=404, detail="Worklist item not found")
    except Exception as e:
        logger.error(f"Error deleting worklist item {item_id}: {e}")
        raise HTTPException(status_code=500, detail="Internal server error")

@router.post("/{item_id}/complete", summary="Отметить исследование как выполненное")
async def complete_worklist_item(
    item_id: int,
    study_uid: str,
    db: None = Depends(get_db)
) -> WorklistItem:
    try:
        worklist_item = Worklist.get_by_id(item_id)

        if worklist_item.study_uid:
            raise HTTPException(status_code=400, detail="Worklist item already completed")

        worklist_item.study_uid = study_uid
        worklist_item.save()

        logger.info(f"Worklist item marked as completed: ID={item_id}, StudyUID={study_uid}")

        return WorklistItem.from_orm(worklist_item)

    except DoesNotExist:
        raise HTTPException(status_code=404, detail="Worklist item not found")
    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"Error completing worklist item {item_id}: {e}")
        raise HTTPException(status_code=500, detail="Internal server error")

@router.post("/{item_id}/upload", summary="Загрузить файл исследования")
async def upload_worklist_file(
    item_id: int,
    file: UploadFile = File(...)
) -> Dict[str, bool]:
    """
    Загрузить DICOM или ZIP файл для worklist записи.

    Подменяет данные в DICOM файлах на данные из worklist:
    - PatientID, PatientName, PatientBirthDate, PatientSex
    - AccessionNumber
    - StudyDate, StudyTime (текущие)
    - StudyInstanceUID (генерируется если нет в worklist)

    Сохраняет файлы в хранилище PACS через saveFile и saveIndex.

    Возвращает StudyInstanceUID.
    """
    try:
        pacs_db.connect()
        worklist_item = Worklist.get_by_id(item_id)
        pacs_db.close()

        if not worklist_item.study_uid:
            study_uid_to_use = generate_uid()
        else:
            study_uid_to_use = worklist_item.study_uid

        # Создаем временную директорию
        filename = file.filename or 'temp.dcm'
        with tempfile.TemporaryDirectory() as temp_dir:
            file_path = os.path.join(temp_dir, filename)

            # Сохраняем загруженный файл
            with open(file_path, "wb") as f:
                content = await file.read()
                f.write(content)

            # Обрабатываем ZIP или DICOM
            processed_files = []

            if filename.lower().endswith('.zip'):
                # Распаковываем ZIP
                with zipfile.ZipFile(file_path, 'r') as zip_ref:
                    zip_ref.extractall(temp_dir)

                # Ищем все DICOM файлы
                for root, _, files in os.walk(temp_dir):
                    for filename in files:
                        if filename.lower().endswith('.dcm'):
                            processed_files.append(os.path.join(root, filename))
            else:
                # Одиночный DICOM файл
                processed_files.append(file_path)

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

            # Обрабатываем каждый DICOM файл
            successfully_saved = 0
            for dicom_file in processed_files:
                try:
                    ds = dcmread(dicom_file, force=True)

                    # Подменяем данные пациента
                    ds.PatientID = worklist_item.pat_id
                    ds.PatientName = worklist_item.pat_name
                    ds.PatientBirthDate = worklist_item.pat_birth.strftime("%Y%m%d")
                    ds.PatientSex = worklist_item.pat_sex

                    # AccessionNumber
                    if worklist_item.acc_num:
                        ds.AccessionNumber = worklist_item.acc_num

                    # Study дата и время (текущие)
                    now = datetime.datetime.now()
                    ds.StudyDate = now.strftime("%Y%m%d")
                    ds.StudyTime = now.strftime("%H%M%S")

                    # StudyInstanceUID
                    ds.StudyInstanceUID = study_uid_to_use

                    # SeriesInstanceUID и SOPInstanceUID должны быть уникальными
                    # Генерируем новые если нужно сохранить структуру
                    if 'SeriesInstanceUID' in ds:
                        # Сохраняем оригинальный SeriesInstanceUID как базовый
                        # но добавляем суффикс от worklist для уникальности
                        original_series_uid = ds.SeriesInstanceUID
                        ds.SeriesInstanceUID = f"{original_series_uid}.{worklist_item.id}"

                    if 'SOPInstanceUID' in ds:
                        # Генерируем новый SOPInstanceUID
                        ds.SOPInstanceUID = generate_uid()

                    # Подготавливаем dataset (как при приеме от аппарата)
                    #(prepared_ds, error) = prepareDataset(ds, registrateOptions=config.config.registrateOptions)

                    # if error:
                    #     logger.warning(f"Error preparing dataset for {dicom_file}: {error}")
                    #     # Продолжаем с оригинальным ds если ошибка подготовки
                    #     prepared_ds = ds

                    # Сохраняем файл в хранилище PACS
                    # Используем aetitle = "worklist" для идентификации источника
                    saveFile(ds, "worklist", saveFile=True)
                    successfully_saved += 1

                except Exception as e:
                    logger.error(f"Error processing DICOM file {dicom_file}: {e}")
                    continue

            if successfully_saved == 0:
                raise HTTPException(status_code=400, detail="Failed to save any DICOM files")

            # Сохраняем study_uid только после успешной обработки всех файлов
            if not worklist_item.study_uid:
                worklist_item.study_uid = study_uid_to_use
                # Переподключаемся если соединение закрыто saveIndex
                if pacs_db.is_closed():
                    pacs_db.connect()
                worklist_item.save()

            logger.info(f"File uploaded for worklist item {item_id}: {file.filename}, StudyUID={study_uid_to_use}, Saved {successfully_saved}/{len(processed_files)} files")

            # Закрываем соединение после всей работы
            if not pacs_db.is_closed():
                pacs_db.close()

            return {
                "success": True
            }

    except DoesNotExist:
        raise HTTPException(status_code=404, detail="Worklist item not found")
    except zipfile.BadZipFile:
        raise HTTPException(status_code=400, detail="Invalid ZIP file")
    except Exception as e:
        logger.error(f"Error uploading file for worklist item {item_id}: {e}")
        raise HTTPException(status_code=500, detail="Internal server error")
