import json
from typing import TYPE_CHECKING, Any, Dict, List, Optional
from uuid import uuid4

from environs import Env
from sqlalchemy import create_engine
from sqlalchemy.exc import SQLAlchemyError
from sqlalchemy.orm import Session

from database.data_models import OpenAPITestCaseCollection, OpenAPITestCaseGeneration

if TYPE_CHECKING:
    from models.test_case import Case

env = Env()
env.read_env()

DATABASE_URL = env.str("DATABASE_URL")  # DATABASE_URL or TEST_DATABASE_URL


class OpenAPITestCaseGenerationDAO:
    def __init__(self):
        self.engine = create_engine(DATABASE_URL)
        self.session = Session(bind=self.engine)

    def insert(self, metadata: dict) -> Optional[str]:
        try:
            if self.get_by_trace_id(metadata["trace_id"]):
                return metadata["trace_id"]
            if metadata.get("trace_id"):
                trace_id = metadata["trace_id"]
            else:
                trace_id = str(uuid4())
            db_generation = OpenAPITestCaseGeneration(
                id=trace_id,
                run_id=None,
                trace_id=trace_id,
                service_name=metadata["service_name"],
                protocol=metadata.get("protocol", "OpenAPI"),
                spec_hash=metadata["spec_hash"],
                spec_commit_id=metadata["spec_commit_id"],
                spec_pr_url=metadata["spec_pr_url"],
                generator_commit_id=metadata["generator_commit_id"],
                generator_pr_url=metadata["generator_pr_url"],
                trigger_type=metadata["trigger_type"],
                enhanced_operations=metadata["enhanced_operations"],
            )
            self.session.add(db_generation)
            self.session.commit()
            return trace_id
        except SQLAlchemyError:
            self.session.rollback()
            return None

    def get_by_trace_id(self, trace_id: str) -> Optional[OpenAPITestCaseGeneration]:
        return (
            self.session.query(OpenAPITestCaseGeneration)
            .filter(OpenAPITestCaseGeneration.trace_id == trace_id)
            .first()
        )


class OpenAPITestCaseCollectionDAO:
    def __init__(self):
        self.engine = create_engine(DATABASE_URL)
        self.session = Session(bind=self.engine)

    def insert(self, case: "Case", generation_id: str, trace_id: str) -> Optional[str]:
        try:
            if self.get_by_id(case.id):
                return case.id
            db_case = OpenAPITestCaseCollection(
                id=case.id,
                generation_id=generation_id,
                trace_id=trace_id,
                operation_id=case.operation_id,
                method=case.method,
                path=case.path,
                environment=case.environment,
                request=case.request,
                type=case.type,
                category=case.category,
                affected_element=case.affected_element,
                element_type=case.element_type,
                generation_strategy=case.generation_strategy,
                expected_status=case.expected_status,
                description=case.description,
                # ignored=case.ignored,
            )
            self.session.add(db_case)
            self.session.commit()
            return case.id
        except SQLAlchemyError:
            self.session.rollback()
            return None

    def batch_insert(
        self, cases: List["Case"], generation_id: str, trace_id: str
    ) -> bool:
        try:
            db_cases = [
                OpenAPITestCaseCollection(
                    id=case.id,
                    generation_id=generation_id,
                    trace_id=trace_id,
                    operation_id=case.operation_id,
                    method=case.method,
                    path=case.path,
                    environment=case.environment,
                    request=case.request,
                    type=case.type,
                    category=case.category,
                    affected_element=case.affected_element,
                    element_type=case.element_type,
                    generation_strategy=case.generation_strategy,
                    expected_status=case.expected_status,
                    description=case.description,
                    # ignored=case.ignored,
                )
                for case in cases
            ]
            self.session.bulk_save_objects(db_cases)
            self.session.flush()
            self.session.commit()
            return True
        except SQLAlchemyError as e:
            print(f"Error: {str(e)}")
            self.session.rollback()
            return False

    def get_by_id(self, case_id: str) -> Optional[OpenAPITestCaseCollection]:
        return (
            self.session.query(OpenAPITestCaseCollection)
            .filter(OpenAPITestCaseCollection.id == case_id)
            .first()
        )

    def get_by_operation_id(self, operation_id: str) -> List[OpenAPITestCaseCollection]:
        return (
            self.session.query(OpenAPITestCaseCollection)
            .filter(OpenAPITestCaseCollection.operation_id == operation_id)
            .all()
        )

    def get_by_generation_id(
        self, generation_id: str
    ) -> List[OpenAPITestCaseCollection]:
        return (
            self.session.query(OpenAPITestCaseCollection)
            .filter(OpenAPITestCaseCollection.generation_id == generation_id)
            .all()
        )