# middleware/guardrails.py
import time
from fastapi import Request
from starlette.middleware.base import BaseHTTPMiddleware
from starlette.responses import Response, JSONResponse
from presidio_analyzer import AnalyzerEngine
from presidio_anonymizer import AnonymizerEngine
from presidio_anonymizer.entities import OperatorConfig
from prometheus_client import Histogram, Counter
import logging

logger = logging.getLogger(__name__)

# ── Prometheus Metrics ──────────────────────────────────────────
# The histogram tracks the distribution of response times (not just the average)
REQUEST_DURATION = Histogram(
    "agent_request_duration_ms",
    "Agent response time in ms",
    buckets=[100, 500, 1000, 5000, 15000, 30000, 60000]
)
# The Counter increases monotonically — ideal for anomaly detection alerts
POLICY_VIOLATIONS = Counter(
    "agent_policy_violations_total", "Number of blocked responses",
    labelnames=["reason"]
)
PII_DETECTIONS = Counter(
    "agent_pii_detections_total", "Number of detected PII entities",
    labelnames=["entity_type"]
)

# ── Prohibited Phrases (Prompt Injection / Compliance) ──────────
FORBIDDEN_PHRASES = [
    "ignore previous instructions",
    "ignore all instructions",
    "you are now",    # classic prompt injection attack
    "disregard your",
]


class AgentGuardrailsMiddleware(BaseHTTPMiddleware):

    def __init__(self, app, languages: list[str] = None):
        super().__init__(app)
        # Initialize ONCE in the constructor — NLP models are resource-heavy (~200ms);
        # we want to avoid loading them per-request
        self.languages = languages or ['en', 'pl']
        self.analyzer = AnalyzerEngine()
        self.anonymizer = AnonymizerEngine()

    async def dispatch(self, request: Request, call_next):
        start_time = time.monotonic()

        # ── STEP 1: Call proper handler (agent) ───────────────────
        # call_next passes the request deeper into the middleware stack
        # until it reaches the intended endpoint. The agent performs all its logic:
        # querying the LLM, executing tools, and building the response.
        # We wait, and only THEN do we receive the response object.
        response = await call_next(request)

        # ── STEP 2: Collect streaming body ───────────────────────
        # Starlette streams responses chunk-by-chunk — there is
        # no single 'response.text'. We must assemble the body manually.
        # Note: this buffers the entire response in RAM.
        body = b''
        async for chunk in response.body_iterator:
            body += chunk
        text = body.decode('utf-8')

        # ── STEP 3: Compliance check ─────────────────────────────
        # Check if the agent returned any prohibited content.
        # We block the response here, before anything reaches the client.
        violation = self._check_policy(text)
        if violation:
            POLICY_VIOLATIONS.labels(reason=violation).inc()
            logger.warning(f'Policy violation [{violation}]')
            return JSONResponse(
                status_code=403,
                content={'error': 'Response blocked', 'reason': violation}
            )

        # ── STEP 4: PII masking with Presidio ────────────────────
        # Presidio uses NLP models (spaCy) instead of pure RegEx.
        # It understands context: '48123456789' as a phone number
        # in 'call me at...' but not when it's an Order ID.
        masked_text = self._mask_pii(text)

        # ── STEP 5: Metrics ──────────────────────────────────────
        # duration_ms is sent to the Prometheus Histogram, allowing us
        # to calculate p50/p95/p99 latency and trigger alerts on threshold violations.
        duration_ms = (time.monotonic() - start_time) * 1000
        REQUEST_DURATION.observe(duration_ms)
        logger.info(f'Agent: {duration_ms:.0f}ms, {len(text)} chars')

        # ── STEP 6: Return response ──────────────────────────────
        # Crucial: replicate the original status_code and headers.
        # Without this you lose Content-Type and bypass original 403/404 codes.
        return Response(
            content=masked_text.encode('utf-8'),
            status_code=response.status_code,   # keep original code
            headers=dict(response.headers),      # keep original headers
            media_type=response.media_type,
        )

    def _check_policy(self, text: str) -> str | None:
        """Returns the violation description or None if the text is compliant."""
        lower = text.lower()
        for phrase in FORBIDDEN_PHRASES:
            if phrase in lower:
                return f'prompt_injection: "{phrase}"'
        return None

    def _mask_pii(self, text: str) -> str:
        # Presidio operates in two stages:
        # 1. analyzer.analyze() — detects entities and returns their positions
        # 2. anonymizer.anonymize() — replaces them according to the defined strategy
        results = []
        for lang in self.languages:
            results.extend(self.analyzer.analyze(
                text=text, language=lang,
                entities=['PERSON', 'EMAIL_ADDRESS', 'PHONE_NUMBER',
                          'IBAN_CODE', 'CREDIT_CARD', 'IP_ADDRESS'],
                score_threshold=0.6,  # eliminates most false positives
            ))
        if not results:
            return text

        # Log the entity TYPE only — never the value itself (that would be a PII leak!)
        for r in results:
            PII_DETECTIONS.labels(entity_type=r.entity_type).inc()
            logger.info(f'PII: {r.entity_type} score={r.score:.2f}')

        anonymized = self.anonymizer.anonymize(
            text=text,
            analyzer_results=results,
            operators={
                "PERSON":        OperatorConfig("replace", {"new_value": "[PERSON]"}),
                "EMAIL_ADDRESS": OperatorConfig("replace", {"new_value": "[EMAIL]"}),
                "PHONE_NUMBER":  OperatorConfig("replace", {"new_value": "[PHONE]"}),
                "IBAN_CODE":     OperatorConfig("replace", {"new_value": "[IBAN]"}),
                "CREDIT_CARD":   OperatorConfig("replace", {"new_value": "[CARD]"}),
                "IP_ADDRESS":    OperatorConfig("replace", {"new_value": "[IP]"}),
            }
        )
        return anonymized.text

# main.py
from fastapi import FastAPI
from middleware.guardrails import AgentGuardrailsMiddleware

app = FastAPI()

# Middleware stack — order matters!
# Last in = first executed (LIFO)
app.add_middleware(AgentGuardrailsMiddleware)
app.add_middleware(CostTrackingMiddleware)
app.add_middleware(AuthMiddleware)  # <-- this executes first

# dependencies/sandboxed_db.py
from fastapi import Depends
from sqlalchemy import create_engine, text
from sqlalchemy.orm import Session

# Separate connection pool — read-only
# DB account with SELECT-only privileges; no INSERT/UPDATE/DELETE
READONLY_DB_URL = 'postgresql://agent_ro:pass@db-replica/prod'

readonly_engine = create_engine(
    READONLY_DB_URL,
    pool_size=5,           # Small pool — prevents the agent from exhausting connections
    max_overflow=2,        # Maximum of 7 total concurrent connections
    pool_timeout=10,       # Timeout — ensures the agent doesn't wait indefinitely
    connect_args={'options': '-c default_transaction_read_only=on'}
)

class SandboxedDatabase:
    """Database with enforced read-only access at the connection level."""

    def __init__(self, session: Session):
        self._session = session
        self._query_count = 0
        self._max_queries = 10  # Rate limit per request

    def execute_query(self, sql: str, params: dict = None):
        if self._query_count >= self._max_queries:
            raise PermissionError('Query limit exceeded for this agent session')

        # Block dangerous operations at the application level
        forbidden = ['DROP', 'DELETE', 'TRUNCATE', 'ALTER', 'GRANT']
        for keyword in forbidden:
            if keyword in sql.upper():
                raise PermissionError(f'Forbidden SQL keyword: {keyword}')

        self._query_count += 1
        return self._session.execute(text(sql), params or {})


def get_sandboxed_db() -> SandboxedDatabase:
    with Session(readonly_engine) as session:
        yield SandboxedDatabase(session)


@app.post('/agent/query')
async def agent_query(
    request: AgentRequest,
    db: SandboxedDatabase = Depends(get_sandboxed_db),  # Injection
):
    # The agent has access ONLY to db.execute_query()
    # No access to the engine, session, or other databases
    result = db.execute_query(request.sql)
    return {'data': result.fetchall()}

class AgentToolkit:
    def __init__(self, allowed_tools: list[str]):
        self._tools = {
            'search_products': self._search_products,   # OK
            'get_order_status': self._get_order_status, # OK
            # 'send_email': ...    — NOT registered = unavailable
            # 'execute_code': ... — NOT registered = unavailable
        }
        self._available = {k: v for k, v in self._tools.items()
                          if k in allowed_tools}

    def call(self, tool_name: str, **kwargs):
        if tool_name not in self._available:
            raise PermissionError(f'Tool not available: {tool_name}')
        return self._available[tool_name](**kwargs)

# observability/tracing.py
from opentelemetry import trace
from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.sdk.trace.export import BatchSpanProcessor
from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import OTLPSpanExporter
from opentelemetry.instrumentation.fastapi import FastAPIInstrumentor
from opentelemetry.instrumentation.httpx import HTTPXClientInstrumentor

def setup_telemetry(app):
    provider = TracerProvider(
        resource=Resource.create({
            'service.name': 'agent-service',
            'service.version': '2.1.0',
            'deployment.environment': 'production',
        })
    )

    exporter = OTLPSpanExporter(endpoint='http://otel-collector:4317')
    provider.add_span_processor(BatchSpanProcessor(exporter))
    trace.set_tracer_provider(provider)

    # Auto-instrumentation — all FastAPI requests are automatically traced
    FastAPIInstrumentor.instrument_app(app)

    # HTTP client instrumentation — LLM API calls are also tracked
    HTTPXClientInstrumentor().instrument()

# agent/reasoning.py
from opentelemetry import trace

tracer = trace.get_tracer('agent.reasoning')

async def execute_agent(user_query: str) -> str:
    with tracer.start_as_current_span('agent.execute') as agent_span:
        agent_span.set_attribute('agent.query', user_query)
        agent_span.set_attribute('agent.model', 'claude-3-5-sonnet')

        for step in range(MAX_STEPS):
            with tracer.start_as_current_span(f'cot.step_{step}') as step_span:
                step_span.set_attribute('cot.step_number', step)

                with tracer.start_as_current_span('llm.call') as llm_span:
                    llm_span.set_attribute('llm.prompt_tokens', len(prompt))
                    response = await call_llm(prompt)
                    llm_span.set_attribute('llm.completion_tokens',
                                          response.usage.completion_tokens)
                    llm_span.set_attribute('llm.cost_usd',
                                          calculate_cost(response.usage))

                if response.tool_call:
                    with tracer.start_as_current_span('tool.call') as tool_span:
                        tool_span.set_attribute('tool.name', response.tool_call.name)
                        result = await execute_tool(response.tool_call)
                        tool_span.set_attribute('tool.success', True)

                if response.is_final:
                    break

        return response.content

agent.execute                          [0ms ────────────────────── 45,230ms]
  cot.step_0                           [12ms ─────── 8,400ms]
    llm.call                           [15ms ─── 7,800ms]   tokens: 412/623
    tool.call: search_products         [8,410ms ─ 8,890ms]  success: true
  cot.step_1                           [8,900ms ─────── 28,100ms]
    llm.call                           [8,900ms ─ 27,600ms] tokens: 1024/892
    tool.call: get_order_status        [27,610ms ─ 28,080ms]
  cot.step_2 (final)                   [28,100ms ─── 45,200ms]
    llm.call                           [28,100ms ─ 45,100ms] tokens: 2048/312

# Synchronous — blocks the thread
def handle_request_sync(user_query: str) -> str:
    response = requests.post(LLM_API, json={...})  # <-- BLOCKS 8 seconds
    return response.json()['content']              # Thread sits idle

# Asynchronous — releases the thread while waiting
async def handle_request_async(user_query: str) -> str:
    async with httpx.AsyncClient() as client:
        response = await client.post(LLM_API, json={...})
        # await = 'wait for the result, but release the event loop'
        # During the 8-second wait, the event loop processes other requests
    return response.json()['content']

@app.post('/agent/query')
async def agent_endpoint(request: AgentRequest):
    result = await handle_request_async(request.query)
    return {'response': result}

# WRONG — sync operation blocks the event loop
async def bad_agent_step():
    data = heavy_json_parse(large_response)  # Blocks everyone!
    return data

# CORRECT — delegate CPU-bound work to a thread pool
async def good_agent_step():
    data = await asyncio.to_thread(heavy_json_parse, large_response)
    return data