from galileo_observe import ObserveWorkflows
import galileo_protect as gp
from pydantic import BaseModel
import promptquality as pq
from promptquality import CustomizedScorerName, Models
from dotenv import load_dotenv
import os
from datetime import datetime
from typing import Optional
load_dotenv()

class GalileoPlatformConfig(BaseModel):
    """Base configuration for Galileo platform."""
    evaluate_project_name: str
    observe_project_name: str
    protect_project_name: str
    protect_stage_name: str


class GalileoPlatform:
    """Implementation of Galileo Features"""

    def __init__(self, config: GalileoPlatformConfig):
        self.config = config
        pq.login(api_key=os.getenv("GALILEO_API_KEY"))
        self.evaluate_run = self.create_evaluate_run()
        self.observe_logger = ObserveWorkflows(project_name=config.observe_project_name)
        self.protect_stage_id = self.get_protect_stage()

    def create_evaluate_run(self):
        """Create a Galileo Evaluate run."""
        scorers = [
            pq.Scorers.context_adherence_luna,
            pq.Scorers.chunk_attribution_utilization_luna,
            pq.Scorers.completeness_luna
        ]
        evaluate_run = pq.EvaluateRun(
            project_name=self.config.evaluate_project_name,
            scorers=scorers,
        )
        return evaluate_run

    def get_protect_stage(self):
        """Get or create a Galileo Protect stage."""
        try:
            protect_project = gp.get_project(
                project_name=self.config.protect_project_name
            )
        except Exception as _:
            protect_project = gp.create_project(name=self.config.protect_project_name)

        protect_project_id = protect_project.id

        try:
            protect_stage = gp.get_stage(
                project_id=protect_project_id, stage_name=self.config.protect_stage_name
            )
        except Exception as _:
            protect_stage = gp.create_stage(
                project_id=protect_project_id,
                name=self.config.protect_stage_name,
            )

        return protect_stage.id

    def run_protect(self, prompt: str, output: str, workflow: Optional[ObserveWorkflows] = None) -> dict:
        """Run Galileo Protect on input and output."""
        response = gp.invoke(
            payload=gp.Payload(input=prompt, output=output),
            prioritized_rulesets=[
                gp.Ruleset(
                    rules=[
                        gp.Rule(
                            metric=gp.RuleMetrics.context_adherence_luna,
                            operator=gp.RuleOperator.lte,
                            target_value=0.01,
                        ),
                    ],
                    action=gp.OverrideAction(
                        choices=["Sorry, the input is hallucinatory."]
                    ),
                ),
                gp.Ruleset(
                    rules=[
                        gp.Rule(
                            metric=gp.RuleMetrics.pii,
                            operator=gp.RuleOperator.any,
                            target_value=["email", "phone_number", "name"],
                        )
                    ],
                    action=gp.OverrideAction(
                        choices=["Sorry, the output contains PII."]
                    ),
                ),
                # gp.Ruleset(
                #     rules=[
                #         gp.Rule(
                #             metric="deutsche_bank_company_pii_0",
                #             operator=gp.RuleOperator.gte,
                #             target_value=0.1,
                #         )
                #     ],
                #     action=gp.OverrideAction(
                #         choices=["Sorry, the output contains PII."]
                #     ),
                # )
            ],
            stage_id=self.protect_stage_id,
        )

        if workflow:
            workflow.add_protect(
                payload=gp.Payload(input=prompt, output=output),
                response=response,
            )

        return dict(response)