main.py

import asyncio
import os
from typing import Annotated, Dict, List, Mapping, Optional, Union

import dotenv
import pydantic
from fastapi import FastAPI, File, HTTPException, UploadFile
from fastapi.middleware.cors import CORSMiddleware

from models import ExtendedRulesConfig
from rules import get_config_variables, get_rules_metadata, initialize_rules, validate_config, default_config
from utils import oas_parser
from utils.rule import RulesConfig, Violation, ViolationKey

app = FastAPI()
# disable cors

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

dotenv.load_dotenv()

assert (
    os.environ.get("MISTRAL_API_KEY") is not None
), "MISTRAL_API_KEY is not set in the environment variables."


@app.get("/")
async def read_root():
    return {"message": "See /docs for API documentation."}


# @app.get("/items/{item_id}")
# async def read_item(item_id: int, q: Union[str, None] = None):
#     """
#     Retrieve an item by its ID.
#     """
#     return {"item_id": item_id, "q": q}


# @app.post("/items/")
# async def create_item(item: dict):
#     """
#     Create an item.
#     """
#     return {"item": item}


# @app.get("/items/")
# async def read_items(q: Union[str, None] = None):
#     """
#     Retrieve a list of items.
#     """
#     return {"q": q}


# @app.delete("/items/{item_id}/delete")
# async def delete_item(item_id: int) -> None:
#     """
#     Delete an item by its ID.
#     """
#     return


# @app.get("/car/{car_id}")
# async def read_car(car_id: int, q: Union[str, None] = None):
#     """
#     Retrieve a car by its ID.
#     """
#     return {"car_id": car_id, "q": q}


# # violation /shops/products?shopId={shopId}&productId={productId}
# @app.get("/shops/products")
# async def read_shops_products(shopId: int, productId: int):
#     """
#     Retrieve products by shopId and productId.
#     """
#     return {"shopId": shopId, "productId": productId}


# # /employees/companies/{{companyId}}
# @app.get("/employees/companies/{companyId}")
# async def read_employees_companies(companyId: int):
#     """
#     Retrieve employees by companyId.
#     """
#     return {"companyId": companyId}


# # violation PUT for  collection create
# @app.put("/companies/{companyId}/employees")
# async def create_employee(companyId: int, employee: dict):
#     """
#     Create employee for a company.
#     """
#     return {"companyId": companyId, "employee": employee}


# # violation error_message in 2xx
# class BadSuccess(pydantic.BaseModel):
#     message: str
#     error: Optional[str] = None


# @app.get("/success")
# async def read_success() -> BadSuccess:
#     """
#     Retrieve a success message.
#     """
#     return BadSuccess(message="Success")


# Actual API endpoints


class InstanceStore:
    """
    A simple in-memory store for instances.
    """

    def __init__(self):
        self.last_instance_id = 0
        self.instances = {}
        self.instance_rules: Dict[int, RulesConfig] = {}

    def create_instance(self, spec: oas_parser.OpenAPISpecification) -> int:
        """
        Create a new instance and return its ID.
        """
        self.last_instance_id += 1
        self.instances[self.last_instance_id] = spec
        self.instance_rules[self.last_instance_id] = RulesConfig(
            config=default_config()
        )
        return self.last_instance_id

    def get_instance(self, instance_id: int) -> oas_parser.OpenAPISpecification:
        """
        Retrieve an instance by its ID.
        """
        return self.instances[instance_id]

    def get_rules(self, instance_id: int) -> RulesConfig:
        """
        Retrieve rules configured for an instance.
        """
        return self.instance_rules[instance_id]

    def upsert_rules(self, instance_id: int, rules: RulesConfig) -> None:
        """
        Upsert rules for an instance.
        """
        self.instance_rules[instance_id] = rules

    def update_instance(
        self, instance_id: int, spec: oas_parser.OpenAPISpecification
    ) -> None:
        """
        Update an existing instance.
        """
        if instance_id in self.instances:
            self.instances[instance_id] = spec
        else:
            raise ValueError(f"Instance with ID {instance_id} not found.")

    def delete_instance(self, instance_id: oas_parser.OpenAPISpecification) -> None:
        """
        Delete an instance by its ID.
        """
        if instance_id in self.instances:
            del self.instances[instance_id]
        else:
            raise ValueError(f"Instance with ID {instance_id} not found.")


INSTANCE_STORE_MEMORY = InstanceStore()


class InstanceResponse(pydantic.BaseModel):
    instance_id: int

# /instances
@app.post("/instances", status_code=201)
async def create_instance(
    file: Annotated[
        UploadFile,
        File(
            description="The OpenAPI spec file to evaluate.",
            media_type="application/json",
        ),
    ]
) -> InstanceResponse:
    """
    Create a new instance.
    """
    parser = oas_parser.OASParser()
    parser.load_spec(file.file.read().decode("utf-8"))
    parser.is_valid()
    spec = parser.to_spec()

    # Store the instance
    instance_id = INSTANCE_STORE_MEMORY.create_instance(spec)

    return InstanceResponse(instance_id=instance_id)

@app.get("/instances", response_model=List[InstanceResponse])
async def list_instances() -> List[InstanceResponse]:
    """
    List all instances.
    """
    return [InstanceResponse(instance_id=instance_id) for instance_id in INSTANCE_STORE_MEMORY.instances.keys()]

@app.put("/instances/{instance_id}")
async def upsert_instance(
    instance_id: int,
    file: Annotated[
        UploadFile,
        File(
            description="The OpenAPI spec file to evaluate.",
            media_type="application/json",
        ),
    ],
) -> InstanceResponse:
    """
    Upsert an instance by its ID.
    """
    if instance_id not in INSTANCE_STORE_MEMORY.instances:
        raise ValueError(f"Instance with ID {instance_id} not found.")

    parser = oas_parser.OASParser()
    parser.load_spec(file.file.read().decode("utf-8"))
    parser.is_valid()
    spec = parser.to_spec()

    # Update the instance
    INSTANCE_STORE_MEMORY.update_instance(instance_id, spec)

    return InstanceResponse(instance_id=instance_id)


class ViolationKVSeralized(pydantic.BaseModel):
    key: ViolationKey
    value: List[Violation]


class ViolationsSerialized(pydantic.BaseModel):
    violations: List[ViolationKVSeralized]

    @classmethod
    def from_mapping(
        cls, violations: Mapping[ViolationKey, List[Violation]]
    ) -> "ViolationsSerialized":
        return cls(
            violations=[
                ViolationKVSeralized(key=k, value=v) for k, v in violations.items()
            ]
        )

    def to_mapping(self) -> Mapping[ViolationKey, List[Violation]]:
        return {v.key: v.value for v in self.violations}


class EvaluationResult(pydantic.BaseModel):
    violations: ViolationsSerialized
    tree: oas_parser.PathNode
    structures: Dict[str, oas_parser.NodeParameterSchema]


@app.get("/instances/{instance_id}")
async def read_instance(instance_id: int) -> EvaluationResult:
    """
    Retrieve an instance by its ID, returns spec + identified violations.
    """
    if instance_id not in INSTANCE_STORE_MEMORY.instances:
        raise ValueError(f"Instance with ID {instance_id} not found.")

    spec = INSTANCE_STORE_MEMORY.get_instance(instance_id)
    config = INSTANCE_STORE_MEMORY.get_rules(instance_id)

    rules = await initialize_rules(spec, config=config)

    await asyncio.gather(*[rule.evaluate() for rule in rules if config.enabled.get(rule.rule_id, True)])

    # Collect violations
    violations: Dict[ViolationKey, List[Violation]] = {}
    for rule in rules:
        v = rule.get_violations()
        filters = list(config.ignore.get(rule.rule_id, [])) + config.ignore_all
        # append to respective key
        for k, v in v.items():
            # TODO: Set theory confusing, I might have flipped the subset logic
            if any(k.equal(f) for f in filters):
                continue
            if k not in violations:
                violations[k] = []
            violations[k].extend(v)

    return EvaluationResult(
        violations=ViolationsSerialized.from_mapping(violations),
        tree=spec.tree,
        structures=spec.get_schemas(),
    )

@app.get("/instances/{instance_id}/rules")
async def read_instance_rules(instance_id: int) -> ExtendedRulesConfig:
    """
    Retrieve rules configured for an instance.
    """
    if instance_id not in INSTANCE_STORE_MEMORY.instances:
        raise ValueError(f"Instance with ID {instance_id} not found.")

    rules = INSTANCE_STORE_MEMORY.get_rules(instance_id)
    return ExtendedRulesConfig.from_config(
        config=rules,
        available_config=await get_config_variables(),
        rules_metadata=await get_rules_metadata(),
    )


@app.put("/instances/{instance_id}/rules")
async def upsert_instance_rules(
    instance_id: int,
    rules: RulesConfig,
) -> RulesConfig:
    """
    Upsert rules for an instance.
    """
    if instance_id not in INSTANCE_STORE_MEMORY.instances:
        raise ValueError(f"Instance with ID {instance_id} not found.")

    # Validate the rules
    for rule_id, rule_config in rules.config.items():
        await validate_config(rule_id, rule_config)
    # Upsert the rules
    INSTANCE_STORE_MEMORY.upsert_rules(instance_id, rules)
    return INSTANCE_STORE_MEMORY.get_rules(instance_id)


@app.get("/instances/{instance_id}/rules/{rule_id}")
async def read_instance_rule(instance_id: int, rule_id: int) -> Dict[str, str]:
    """
    Retrieve a specific rule configured for an instance.
    """
    if instance_id not in INSTANCE_STORE_MEMORY.instances:
        raise ValueError(f"Instance with ID {instance_id} not found.")

    rules = INSTANCE_STORE_MEMORY.get_rules(instance_id)

    return rules.config.get(rule_id, {})


@app.put("/instances/{instance_id}/rules/{rule_id}", status_code=204)
async def configure_instance_rule(
    instance_id: int,
    rule_id: int,
    rule_config: Dict[str, str],
) -> None:
    """
    Configure a specific rule for an instance.
    """
    if instance_id not in INSTANCE_STORE_MEMORY.instances:
        raise ValueError(f"Instance with ID {instance_id} not found.")

    # Validate the rule configuration
    await validate_config(rule_id, rule_config)

    # Upsert the rule configuration
    rules = INSTANCE_STORE_MEMORY.get_rules(instance_id)
    rules.config[rule_id] = rule_config
    INSTANCE_STORE_MEMORY.upsert_rules(instance_id, rules)


@app.delete("/instances/{instance_id}/rules/{rule_id}/active", status_code=204)
async def deactivate_instance_rule(instance_id: int, rule_id: int) -> None:
    """
    Deactivate a specific rule configured for an instance.
    """
    if instance_id not in INSTANCE_STORE_MEMORY.instances:
        raise ValueError(f"Instance with ID {instance_id} not found.")

    rules = INSTANCE_STORE_MEMORY.get_rules(instance_id)
    rules.disable_rule(rule_id)


@app.put("/instances/{instance_id}/rules/{rule_id}/active", status_code=204)
async def activate_instance_rule(instance_id: int, rule_id: int) -> None:
    """
    Activate a specific rule configured for an instance.
    """
    if instance_id not in INSTANCE_STORE_MEMORY.instances:
        raise ValueError(f"Instance with ID {instance_id} not found.")

    rules = INSTANCE_STORE_MEMORY.get_rules(instance_id)
    rules.enable_rule(rule_id)


@app.get("/instances/{instance_id}/rules/{rule_id}/ignores")
async def read_instance_ignores(instance_id: int, rule_id: int) -> List[ViolationKey]:
    """
    Retrieve ignores for a specific rule configured for an instance.
    """
    if instance_id not in INSTANCE_STORE_MEMORY.instances:
        raise ValueError(f"Instance with ID {instance_id} not found.")

    rules = INSTANCE_STORE_MEMORY.get_rules(instance_id)
    return rules.ignore.get(rule_id, [])

@app.post("/instances/{instance_id}/rules/{rule_id}/ignores", status_code=204)
async def add_instance_ignore(
    instance_id: int,
    rule_id: int,
    ignore: ViolationKey,
) -> None:
    """
    Add an ignore for a specific rule configured for an instance.
    """
    if instance_id not in INSTANCE_STORE_MEMORY.instances:
        raise ValueError(f"Instance with ID {instance_id} not found.")

    rules = INSTANCE_STORE_MEMORY.get_rules(instance_id)
    rules.ignore.setdefault(rule_id, []).append(ignore)
    INSTANCE_STORE_MEMORY.upsert_rules(instance_id, rules)

@app.delete("/instances/{instance_id}/rules/{rule_id}/ignores", status_code=204)
async def remove_instance_ignore(
    instance_id: int,
    rule_id: int,
    ignore: ViolationKey | None = None,
) -> None:
    """
    Remove an ignore for a specific rule configured for an instance.

    Args:
        ignore (ViolationKey | None): The ignore to remove. If None, remove all ignores for the rule.
    """
    if instance_id not in INSTANCE_STORE_MEMORY.instances:
        raise ValueError(f"Instance with ID {instance_id} not found.")

    rules = INSTANCE_STORE_MEMORY.get_rules(instance_id)
    if ignore is None:
        rules.ignore.pop(rule_id, None)
    else:
        rules.ignore[rule_id] = [
            i for i in rules.ignore.get(rule_id, []) if not i.equal(ignore)
        ]
    INSTANCE_STORE_MEMORY.upsert_rules(instance_id, rules)

@app.post("/instances/{instance_id}/rules/ignores", status_code=204)
async def add_instance_ignore_all(
    instance_id: int,
    ignore: ViolationKey,
) -> None:
    """
    Add an ignore for all rules configured for an instance.
    """
    if instance_id not in INSTANCE_STORE_MEMORY.instances:
        raise ValueError(f"Instance with ID {instance_id} not found.")

    rules = INSTANCE_STORE_MEMORY.get_rules(instance_id)
    rules.ignore_all.append(ignore)
    INSTANCE_STORE_MEMORY.upsert_rules(instance_id, rules)

@app.delete("/instances/{instance_id}/rules/ignores", status_code=204)
async def remove_instance_ignore_all(
    instance_id: int,
    ignore: ViolationKey | None = None,
) -> None:
    """
    Remove an ignore for all rules configured for an instance.
    Args:
        ignore (ViolationKey | None): The ignore to remove. If None, remove all ignores for all rules.
    """

    if instance_id not in INSTANCE_STORE_MEMORY.instances:
        raise ValueError(f"Instance with ID {instance_id} not found.")

    rules = INSTANCE_STORE_MEMORY.get_rules(instance_id)
    if ignore is None:
        rules.ignore_all = []
    else:
        rules.ignore_all = [
            i for i in rules.ignore_all if not i.equal(ignore)
        ]
    INSTANCE_STORE_MEMORY.upsert_rules(instance_id, rules)


@app.post("/evaluate", include_in_schema=False)
async def evaluate_oas(
    file: Annotated[
        UploadFile,
        File(
            description="The OpenAPI spec file to evaluate.",
            media_type="application/json",
        ),
    ]
):
    """
    Evaluate an OpenAPI spec file using the Mistral API.
    """

    parser = oas_parser.OASParser()
    parser.load_spec(file.file.read().decode("utf-8"))
    parser.is_valid()
    spec = parser.to_spec()
    # manulally traverse the tree to find where we cant do a model_dump

    # return spec.tree.model_dump(exclude_defaults=True)
    # return spec.tree.model_dump()
    # return str(spec.tree)

    # Create rules

    config = RulesConfig(
        config={
            75: {
                "component_structure_names": "PascalCase",
                "uri_arg_pattern": "snake_case",
                "json_schema_key": "snake_case",
            }
        }
    )

    rules = await initialize_rules(spec, config=config)

    await asyncio.gather(*[rule.evaluate() for rule in rules])

    # Collect violations
    violations: Dict[ViolationKey, List[Violation]] = {}
    for rule in rules:
        v = rule.get_violations()
        # append to respective key
        for k, v in v.items():
            if k not in violations:
                violations[k] = []
            violations[k].extend(v)

    # Return violations as JSON

    return None

    # return Response(
    #     content=spec.dump(),
    #     media_type="application/yaml",
    # )


@app.exception_handler(ValueError)
async def value_error_handler(request, exc: ValueError):
    raise HTTPException(
        status_code=400,
        detail=str(exc),
    )