from fastapi.exceptions import HTTPException
from .base_storage import BaseStorage
from core.models.office.db import JdeOfficeLocalDB
from core.models.office.requests import UpdateOfficeRequest


class OfficeStorage(BaseStorage):

    def get_by_id(self, idx: int) -> JdeOfficeLocalDB:
        row = self.get_by('office', code=idx)
        if row is None:
            raise HTTPException(status_code=404, detail='Office not found')
        office = JdeOfficeLocalDB(code=str(row[0]), features=row[1], contact_person_id=row[2], person_count=row[3],
                                  rating=row[4])
        return office

    def list_all_offices(self):
        rows = self.list_by('office')
        return [
            JdeOfficeLocalDB(code=str(row[0]),
                             features=row[1],
                             contact_person_id=row[2],
                             person_count=row[3],
                             rating=row[4])
            for row in rows]

    def update_office(self, data: UpdateOfficeRequest):
        try:
            self.get_by_id(data.code)
            sql = f"update office set features = %s, contact_person_id = %s, person_count = %s, rating = %s where code = %s"
            self.db.execute(sql, (data.features, data.contact_person_id, data.person_count, data.rating, data.code))
        except HTTPException as _:
            sql = f"insert into office (code, features, contact_person_id, person_count, rating) values (%s, %s, %s, %s, %s)"
            self.db.execute(sql, (data.code, data.features, data.contact_person_id, data.person_count, data.rating))