import asyncio
from contextlib import asynccontextmanager
from dataclasses import asdict
import dataclasses
import enum
from json import dumps, loads
import logging
import time
from typing import Any, Optional

#
from asyncpg import UniqueViolationError
from sqlalchemy.exc import IntegrityError
from fastapi import APIRouter, WebSocket, WebSocketDisconnect, status
from sqlalchemy.ext.asyncio import AsyncSession

from src.repository.sheet_repository import SheetValidationError
from src.services.rf_project_report_service import RfProjectReportService
from src.services.rf_project_report_line_service import RfProjectReportLineService
from src.repository.user_repository import UserRepository
from src.services.budget_line_service import BudgetLineService
from src.api.v1.deps import get_user_by_token
from src.db.models.app_user import AppUser
from src.db.models.form_type import FormTypeEnum
from src.services.budget_form_service import BudgetFormService
from src.services.project_service import ProjectService
from src.services.sheet_service import SheetService
from src.core.errors import BasicAppException, ValidationsError
from src.db.session import SessionLocal

from src.services.user_service import UserService


logger = logging.getLogger()


@asynccontextmanager
async def get_db_session(user_id: int | None = None):
    """Контекстный менеджер для получения сессии базы данных."""
    db = SessionLocal()
    try:
        if user_id is not None:
            user_repo = UserRepository(db)
            await user_repo.set_app_user_id(user_id=user_id)
        yield db
    finally:
        await db.close()


router = APIRouter(prefix="/ws", tags=["websocket"])


@dataclasses.dataclass
class ConnectionInfo:
    ws: WebSocket
    user_id: int | None


@dataclasses.dataclass
class FormConnectionInfo(ConnectionInfo):
    form_id: int | None = None
    sheet: str | None = None
    direction: str | None = None


@dataclasses.dataclass
class ProjectConnectionInfo(ConnectionInfo):
    project_id: int | None = None
    year: int | None = None
    report_type: str | None = None


class ConnectionKeyEnum(str, enum.Enum):
    FORM = "FORM"
    PROJECT = "PROJECT"


class ConnectionManager:
    con_info_mapping = {
        ConnectionKeyEnum.FORM: FormConnectionInfo,
        ConnectionKeyEnum.PROJECT: ProjectConnectionInfo,
    }

    def __init__(self):
        self.connections: dict[int, dict[ConnectionKeyEnum, list[ConnectionInfo]]] = {
            ConnectionKeyEnum.FORM: {},
            ConnectionKeyEnum.PROJECT: {},
        }
        self.cell_locks: dict[Any, int] = {}
        

    async def connect(
            self, 
            websocket: WebSocket,
            con_key: ConnectionKeyEnum,
            **kwargs,
            # form_id: int, 
            # sheet: str, 
            # direction: str | None = None,
    ):
        await websocket.accept()

        key = frozenset(kwargs.items())
        cls = self.con_info_mapping[con_key]

        if key not in self.connections[con_key]:
            self.connections[con_key][key] = [
                cls(
                    ws=websocket, 
                    user_id=None,
                    **kwargs,
                )
            ]
        else:
            self.connections[con_key][key].append(
                cls(
                    ws=websocket,
                    user_id=None,
                    **kwargs,
                )
            
            )
    
    def set_user(
            self, 
            websocket: WebSocket, 
            user_id: int, 
            con_key: ConnectionKeyEnum,
            **kwargs,
    ):
        key = frozenset(kwargs.items())

        if key not in self.connections[con_key]:
            return
        for con_info in self.connections[con_key][key]:
            if con_info.ws == websocket:
                con_info.user_id = user_id

    async def disconnect(
            self, 
            websocket: WebSocket, 
            con_key: ConnectionKeyEnum,
            code: int = status.WS_1008_POLICY_VIOLATION,
            reason: str = "Ошибка",
            **kwargs,
    ):
        key = frozenset(kwargs.items())

        if key not in self.connections[con_key]:
            return
        for con_info in self.connections[con_key][key]:
            if con_info.ws == websocket:
                self.connections[con_key][key].remove(con_info)
        if not len(self.connections[con_key][key]):
            del self.connections[con_key][key]
        
        self.release_locks_for_user(con_key=con_key, user_id=con_info.user_id, form_key=kwargs)
        await websocket.close(
            code=code, reason=reason,
        )
    
    async def broadcast_to_other(
            self, 
            data: dict, 
            websocket: WebSocket,
            con_key: ConnectionKeyEnum,
            **kwargs,
    ):
        key = frozenset(kwargs.items())

        if key not in self.connections[con_key]:
            return
        
        data = data.copy()
        data["is_self"] = False
        message = dumps(data, default=_convert_error, ensure_ascii=False)

        for con_info in self.connections[con_key][key]:
            if con_info.ws != websocket:
                try:
                    await con_info.ws.send_text(message)
                except:
                    pass  # Игнорируем недоступные соединения

    async def broadcast_to_all(
            self, 
            data: dict,  
            con_key: ConnectionKeyEnum,
            websocket: WebSocket | None = None,
            **kwargs,
    ):
        key = frozenset(kwargs.items())
        if key not in self.connections[con_key]:
            return
        
        data = data.copy()
        if websocket is None:
            data["is_self"] = None
            default_message = dumps(data, default=_convert_error, ensure_ascii=False)
            self_message = None
        else:
            data["is_self"] = False
            default_message = dumps(data, default=_convert_error, ensure_ascii=False)
            data["is_self"] = True
            self_message = dumps(data, default=_convert_error, ensure_ascii=False)

        for con_info in self.connections[con_key][key]:
            try:
                if websocket is not None and con_info.ws == websocket:
                    await con_info.ws.send_text(self_message)
                else:
                    await con_info.ws.send_text(default_message)


            except:
                pass  # Игнорируем недоступные соединения

    async def send_back(self, data: dict, websocket: WebSocket):
        data = data.copy()
        data["is_self"] = True
        message = dumps(data, default=_convert_error, ensure_ascii=False)
        await websocket.send_text(message)
    
    async def get_data(self, websocket: WebSocket) -> dict:
        return loads(await websocket.receive_text())

    def acquire_cell_lock(
            self, 
            con_key: ConnectionKeyEnum, 
            user_id: int,
            cell_key: dict,
            form_key: dict,
    ) -> bool:
        if not cell_key:
            return False
        cell_key = frozenset(cell_key.items())
        
        lock_key = (con_key, frozenset(form_key.items()), cell_key)
        lock_owner_id = self.cell_locks.get(lock_key)
        if lock_owner_id is not None and lock_owner_id != user_id:
            return False
        self.cell_locks[lock_key] = user_id
        return True

    def release_cell_lock(
            self, 
            con_key: ConnectionKeyEnum, 
            user_id: int,
            cell_key: dict,
            form_key: dict,
    ):
        if not cell_key:
            return
        cell_key = frozenset(cell_key.items())

        lock_key = (con_key, frozenset(form_key.items()), cell_key)
        lock_owner_id = self.cell_locks.get(lock_key)
        if lock_owner_id == user_id:
            self.cell_locks.pop(lock_key, None)

    def release_locks_for_user(
            self, 
            con_key: ConnectionKeyEnum, 
            user_id: int,
            form_key: dict
        ):
        filter_key = (con_key, frozenset(form_key.items()))
        lock_keys = [
            lock_key
            for lock_key, lock_owner_id in self.cell_locks.items()
            if lock_key[:-1] == filter_key and lock_owner_id == user_id
        ]
        for lock_key in lock_keys:
            self.cell_locks.pop(lock_key, None)

    def is_cell_available_for_edit(
            self, 
            con_key: ConnectionKeyEnum, 
            user_id: int,
            cell_key: dict,
            form_key,
    ) -> bool:
        if not cell_key:
            return True
        cell_key = frozenset(cell_key.items())
        lock_owner_id = self.cell_locks.get((con_key, frozenset(form_key.items()), cell_key))
        return lock_owner_id is not None and lock_owner_id == user_id

    def is_row_locked_by_other(
            self, 
            con_key: ConnectionKeyEnum, 
            user_id: int,
            line_id: int,
    ) -> bool:
        if line_id is None:
            return False

        for (dict_con_key, form_key, cell_key), lock_owner_id in self.cell_locks.items():
            if dict_con_key != con_key:
                continue
            
            if form_key != form_key or lock_owner_id == user_id:
                continue

            for k, v in cell_key:
                if k == "line_id" and v == line_id:
                    return True
        return False

    def release_locks_for_row(
            self, 
            con_key: ConnectionKeyEnum, 
            user_id: int,
            line_id: int,
    ) -> bool:
        
        if line_id is None:
            return False
        
        lock_keys = []
        for lock_key, lock_owner_id in self.cell_locks.items():
            dict_con_key, form_key, cell_key = lock_key
            if dict_con_key != con_key:
                continue
            
            if form_key != form_key or lock_owner_id == user_id:
                continue

            for k, v in cell_key.items():
                if k == "line_id" and v == line_id:
                    lock_keys.append(lock_key)
        
        for lock_key in lock_keys:
            self.cell_locks.pop(lock_key, None)


class FormEventProcess:
    def __init__(self, db: AsyncSession):
        self.sheet_service: SheetService = SheetService(db)
        self.bf_service: BudgetFormService = BudgetFormService(db)
        self.bl_service: BudgetLineService = BudgetLineService(db)
        self.user_service: UserService = UserService(db)


    async def process(
        self, 
        event_data: dict, 
        form_id: int, 
        user_id: int,
        sheet: str,
        direction: str | None = None,
    ) -> int | bool | dict:
        curr_user: AppUser = await self.user_service.get(user_id)

        match event_data["event"]:
            case "cell_updated":
                return await self.__update_cell(
                    event_data=event_data["data"], 
                    form_id=form_id, 
                    user=curr_user,
                    sheet=sheet,
                    direction=direction,
                )
            case "row_added":
                return await self.__add_row(
                    event_data=event_data["data"], 
                    form_id=form_id, 
                    user=curr_user,
                    sheet=sheet,
                    direction=direction,
                )
            case "row_deleted":
                return await self.__del_row(
                    event_data=event_data["data"], 
                    form_id=form_id, 
                    user=curr_user,
                    sheet=sheet,
                    direction=direction,
                )
            case "program_added":
                return await self.__add_program(
                    event_data=event_data["data"],
                    form_id=form_id, 
                    user=curr_user,
                    sheet=sheet,
                )
            case "project_added":
                return await self.__add_project(
                    event_data=event_data["data"],
                    user=curr_user,
                )
            case _:
                pass

    async def __update_cell(
            self, 
            event_data: dict, 
            form_id: int, 
            sheet: str, 
            user: AppUser,
            direction: str | None = None,
    ) -> list[tuple]:
        """
        event_data: {
            "line_id": int,
            "column": str,
            "value": any,
        }
        """
        form = await self.bf_service.get(budget_form_id=form_id, user=user)
        if not form:
            return None

        return await self.sheet_service.update_cell(
            form_id=form_id, 
            sheet=sheet,
            direction=direction,
            sections=None,
            line_id=event_data["line_id"], 
            column=event_data["column"], 
            value=event_data["value"],
            user=user,
        )

    async def __add_row(
            self, 
            event_data: dict, 
            form_id: int, 
            sheet: str,
            user: AppUser,
            direction: str | None,
    ) -> list[tuple]:
        """
            data = {
                expense_item_id: Optional[int] = None
                item_id: Optional[str] = None
                section_code: Optional[str] = None
                name: Optional[str] = None
                internal_order: Optional[str] = None
                vsp_id: Optional[int] = None
                project_id: Optional[int] = None
                justification: Optional[str] = None
                contract_number: Optional[str] = None
                contract_end_date: Optional[datetime] = None
            }
        """
        form = await self.bf_service.get(budget_form_id=form_id, user=user)
        if not form:
            return None

        ids = await self.bl_service.get_ids(budget_form_id=form_id)
        result = await self.sheet_service.add_line(
            form_id=form_id,
            sheet=sheet,
            expense_item_id=event_data.get("expense_item_id"),
            item_id=event_data.get("item_id"),
            section_code=event_data.get("section_code"),
            direction=direction,
            name=event_data.get("name"),
            internal_order=event_data.get("internal_order"),
            vsp_id=event_data.get("vsp_id"),
            project_id=event_data.get("project_id"),
            justification=event_data.get("justification"),
            contract_number=event_data.get("contract_number"),
            contract_end_date=event_data.get("contract_end_date"),
            user=user,
        )
        final_result = {
            "data": result,
            "new_line_id": None,
        }
        for el in result:
            if el[3]["line_id"] and el[3]["line_id"] not in ids:
                final_result["new_line_id"] = el[3]["line_id"]
                return final_result
        return final_result
    
    async def __add_program(
            self, 
            event_data: dict, 
            form_id: int, 
            sheet: str,
            user: AppUser,
    ) -> list[tuple]:
        """
            data = {
                name: str
            }
        """
        return await self.sheet_service.add_program(
            form_id=form_id,
            sheet=sheet,
            name=event_data["name"],
            user=user,
        )

    async def __add_project(
            self, 
            event_data: dict, 
            user: AppUser,
    ) -> list[tuple]:
        """
            data = {
                name: str,
                program_id: int
            }
        """
        return await self.sheet_service.add_project(
            name=event_data["name"],
            program_id=event_data["program_id"],
            user=user,
        )
    
    async def __del_row(
            self, 
            event_data: dict, 
            form_id: int, 
            sheet: str,
            user: AppUser,
            direction: str | None,
    ) -> list[tuple]:
        """
            data = {
                "row_id": int,
            }
        """
        form = await self.bf_service.get(budget_form_id=form_id, user=user)
        if not form:
            return None
        result = await self.sheet_service.delete_line(
            form_id=form_id,
            sheet=sheet,
            row_id=event_data["row_id"],
            direction=direction,
            user=user,
        )
        return result


class ProjectEventProcess:
    def __init__(self, db: AsyncSession):
        self.project_service: ProjectService = ProjectService(db)
        self.prl_service: RfProjectReportLineService = RfProjectReportLineService(db)
        self.pr_service: RfProjectReportService = RfProjectReportService(db)
        self.user_service: UserService = UserService(db)

    async def process(
        self, 
        event_data: dict, 
        project_id: int,
        user_id: int,
        year: int,
        report_type: str,
    ) -> int | bool | dict:
        curr_user: AppUser = await self.user_service.get(user_id)

        match event_data["event"]:
            case "cell_updated":
                return await self.__update_cell(
                    event_data=event_data["data"],
                    project_id=project_id,
                    year=year,
                    report_type=report_type,
                    user=curr_user,
                )
            case "row_added":
                return await self.__add_row(
                    event_data=event_data["data"],
                    project_id=project_id,
                    year=year,
                    report_type=report_type,
                    user=curr_user,
                )
            case "row_deleted":
                return await self.__del_row(
                    event_data=event_data["data"],
                    project_id=project_id,
                    year=year,
                    report_type=report_type,
                    user=curr_user,
                )
            case _:
                pass

    async def __update_cell(
            self, 
            event_data: dict, 
            project_id: int,
            year: int,
            report_type: str,
            user: AppUser,
    ) -> list[tuple]:
        """
        event_data: {
            "line_id": int,
            "column": str,
            "value": any,
        }
        """
        project = await self.project_service.get(project_id=project_id, user=user)
        if not project:
            return None

        return await self.project_service.upd_form3_cell(
            project_id=project_id,
            year=year,
            report_type=report_type,
            line_id=event_data["line_id"],
            column=event_data["column"],
            value=event_data["value"],
            user=user,
        )

    async def __add_row(
            self, 
            event_data: dict,
            project_id: int,
            year: int,
            report_type: str,
            user: AppUser,
    ) -> list[tuple]:
        """
            data = {
                expense_item_id: Optional[int] = None
            }
        """
        project_report = await self.pr_service.get(
            year=year,
            report_type=report_type,
            project_id=project_id,
        )
        if not project_report:
            return None
        
        ids = await self.prl_service.get_ids(project_report_id=project_report.id)
        result = await self.project_service.add_form3_line(
            project_id=project_id,
            year=year,
            report_type=report_type,
            expense_item_id=event_data["expense_item_id"],
            user=user,
        )

        final_result = {
            "data": result,
            "new_line_id": None,
        }
        for el in result:
            if el[3]["line_id"] and el[3]["line_id"] not in ids:
                final_result["new_line_id"] = el[3]["line_id"]
                return final_result
        return final_result

    async def __del_row(
            self, 
            event_data: dict, 
            project_id: int,
            year: int,
            report_type: str,
            user: AppUser,
    ) -> list[tuple]:
        """
            data = {
                "row_id": int,
            }
        """
        project = await self.project_service.get(project_id=project_id, user=user)
        if not project:
            return None

        return await self.project_service.del_form3_line(
            project_id=project_id,
            year=year,
            report_type=report_type,
            line_id=event_data["row_id"],
            user=user,
        )


manager = ConnectionManager()


def _convert_error(o):
    try:
        return asdict(o)
    except TypeError:
        return o


async def login(websocket: WebSocket, **kwargs) -> int:
    user_data = await manager.get_data(websocket=websocket)
    if user_data.get("event") != "user_login":
        return None

    async with get_db_session() as db:
        user = await get_user_by_token(token=user_data["data"].get("token"), db=db)
        if not user:
            return None
        manager.set_user(
            websocket=websocket, 
            user_id=user.id, 
            **kwargs,
        )
        return user.id


def resolve_cell_key(event_data: dict) -> dict | None:
    data = event_data.get("data")

    if not data or ("line_id" not in data and "line_id_code" not in data) or "column" not in data:
        return None
    return {
        "line_id": data.get("line_id_code") or data.get("line_id"),
        "column": data["column"],
    }


async def process_websocket(
    websocket: WebSocket,
    processor_cls,
    **kwargs,
):
    await manager.connect(
        websocket=websocket, 
        **kwargs,
    )

    
    kwargs_process = kwargs.copy()
    if "con_key" in kwargs_process:
        kwargs_process.pop("con_key")
    con_key = kwargs.get("con_key")

    try:
        user_id = await login(
            websocket=websocket,
            **kwargs,
        )
        if user_id is None:
            
            await manager.disconnect(
                websocket=websocket,
                reason="Ошибка авторизации",
                **kwargs,
            )
            return
    
        while True:
            data = loads(await websocket.receive_text())

            try:
                event_name = data.get("event")
                cell_key = resolve_cell_key(data)

                match event_name:
                    case "cell_edit_start":
                        if not manager.acquire_cell_lock(
                            con_key=con_key,
                            user_id=user_id,
                            cell_key=cell_key,
                            form_key=kwargs_process,
                        ):
                            data["error"] = "Ячейка уже редактируется другим пользователем"
                            await manager.send_back(
                                data, 
                                websocket,
                            )
                            continue
                        await manager.send_back(
                            data, 
                            websocket,
                        )
                        await manager.broadcast_to_other(
                            data=data,
                            # user_id=user_id,
                            websocket=websocket,
                            con_key=con_key,
                            **kwargs_process,
                        )
                        continue
                    case "cell_edit_end":
                        manager.release_cell_lock(
                            con_key=con_key,
                            user_id=user_id,
                            cell_key=cell_key,
                            form_key=kwargs_process,
                        )
                        await manager.send_back(
                            data, 
                            websocket,
                        )
                        await manager.broadcast_to_other(
                            data=data,
                            # user_id=user_id,
                            websocket=websocket,
                            con_key=con_key,
                            **kwargs_process,
                        )
                        continue
                    case "cell_updated":
                        if not manager.is_cell_available_for_edit(
                                con_key=con_key,
                                user_id=user_id,
                                cell_key=cell_key,
                                form_key=kwargs_process,
                        ):
                            data["error"] = "Ячейка уже редактируется другим пользователем"
                            await manager.send_back(
                                data, 
                                websocket,
                            )
                            continue
                        # if cell_key and not manager.acquire_cell_lock(
                        #     con_key=con_key,
                        #     user_id=user_id,
                        #     cell_key=cell_key,
                        #     form_key=kwargs_process,
                        # ):
                        #     data["error"] = "Ячейка уже редактируется другим пользователем"
                        #     await manager.send_back(
                        #         data, 
                        #         websocket,
                        #     )
                        #     continue
                    case "row_deleted":
                        if manager.is_row_locked_by_other(
                            con_key=con_key,
                            user_id=user_id,
                            line_id=data["data"]["row_id"]
                        ):
                            data["error"] = (
                                "Нельзя удалить строку: в этой строке есть редактируемые ячейки"
                            )
                            await manager.send_back(
                                data, 
                                websocket,
                            )
                            continue
                    case _:
                        pass

                async with get_db_session(user_id=user_id) as db:
                    processor = processor_cls(db)
                    t0 = time.perf_counter()
                    data_old = data.copy()
                    data["result"] = await processor.process(
                        event_data=data, 
                        user_id=user_id,
                        **kwargs_process
                    )
                    await db.commit()
                    db_ms = (time.perf_counter() - t0) * 1000

                    extra = {
                        "type": "db_time",
                        "time_ms": f"{db_ms:.2f}", 
                        "event": data.get("event"),
                    }
                    extra.update(data_old)
                    logger.info(
                        f"db_time: {extra}", 
                        extra=extra,
                    )
                
                if data.get("event") == "row_deleted":
                    manager.release_locks_for_row(
                        con_key=con_key,
                        user_id=user_id,
                        line_id=data["data"]["row_id"]
                    )
                
                if data.get("event") in [
                    "cell_updated",
                    "row_added",
                    "program_added",
                    "project_added",
                    "row_deleted",
                ]:
                    await manager.broadcast_to_all(
                        data=data,
                        websocket=websocket,
                        # user_id=user_id,
                        **kwargs
                    )
                else:
                    await manager.broadcast_to_other(
                        data=data,
                        # user_id=user_id,
                        websocket=websocket,
                        **kwargs,
                    )

            except BasicAppException as e:
                data["error"] = e.description or "Неизвестная ошибка"
                await manager.send_back(
                    data=data, 
                    websocket=websocket,
                )
            except IntegrityError as e:
                data["error"] = str(e)
                await manager.send_back(
                    data=data, 
                    websocket=websocket,
                )
            except SheetValidationError as e:
                data["error"] = e.message
                await manager.send_back(
                    data=data, 
                    websocket=websocket,
                )

            except Exception as e:
                tp = type(e)
                handlers = websocket.app.exception_handlers
                if tp in handlers:
                    data["error"] = loads((await handlers[tp](request=None, exc=e)).body)
                    await manager.send_back(
                        data=data, 
                        websocket=websocket,
                    )
                else:
                    raise e

    except WebSocketDisconnect:
        await manager.disconnect(
            websocket=websocket,
            **kwargs
        )

    except Exception as e:
        # Логируем ошибку, но не бросаем HTTPException — это WebSocket
        print(f"Error: {e}")
        await manager.disconnect(
            websocket=websocket,
            **kwargs,
        )


@router.websocket("/form/{form_id}/sheet/{sheet}")
async def websocket_form(
    websocket: WebSocket,
    form_id: int,
    sheet: str,
    direction: Optional[str] = None,
):
    await process_websocket(
        websocket=websocket,
        form_id=form_id,
        sheet=sheet,
        direction=direction,
        processor_cls=FormEventProcess,
        con_key=ConnectionKeyEnum.FORM,
    )


@router.websocket("/projects/{project_id}/report/{year}/{report_type}")
async def websocket_project(
    websocket: WebSocket,
    project_id: int,
    year: int,
    report_type: str,
):
    await process_websocket(
        websocket=websocket,
        project_id=project_id,
        year=year,
        report_type=report_type,
        processor_cls=ProjectEventProcess,
        con_key=ConnectionKeyEnum.PROJECT,
    )
