import anthropic
import psycopg2
import json
from typing import Optional

client = anthropic.Anthropic()

def embed_text(text: str) -> list[float]:
    """Generate embeddings using a lightweight approach via Claude."""
    # In production: use a dedicated embedding model like text-embedding-3-small
    # or voyage-3 for cost efficiency. Claude isn't primarily an embedding model.
    # This is a placeholder showing the integration pattern.
    response = client.messages.create(
        model="claude-haiku-4-5-20251001",
        max_tokens=10,
        messages=[{"role": "user", "content": f"Embed: {text[:100]}"}]
    )
    # Real implementation: call your embedding API here
    return []  

def store_document_chunks(
    conn: psycopg2.extensions.connection,
    document_id: str,
    chunks: list[str],
    metadata: dict
) -> int:
    """Store document chunks with embeddings in pgvector."""
    stored = 0
    with conn.cursor() as cur:
        for i, chunk in enumerate(chunks):
            embedding = embed_text(chunk)

            cur.execute(
                """INSERT INTO document_chunks 
                   (document_id, chunk_index, content, embedding, metadata)
                   VALUES (%s, %s, %s, %s::vector, %s)
                   ON CONFLICT (document_id, chunk_index) DO UPDATE
                   SET content = EXCLUDED.content,
                       embedding = EXCLUDED.embedding""",
                (document_id, i, chunk, embedding, json.dumps(metadata))
            )
            stored += 1

    conn.commit()
    return stored

def semantic_search(
    conn: psycopg2.extensions.connection,
    query: str,
    limit: int = 5,
    metadata_filter: Optional[dict] = None
) -> list[dict]:
    """Search document chunks by semantic similarity."""
    query_embedding = embed_text(query)

    filter_clause = ""
    filter_params = []
    if metadata_filter:
        conditions = [f"metadata->>{repr(k)} = %s" for k in metadata_filter]
        filter_clause = "WHERE " + " AND ".join(conditions)
        filter_params = list(metadata_filter.values())

    with conn.cursor() as cur:
        cur.execute(
            f"""SELECT document_id, chunk_index, content, metadata,
                       1 - (embedding <=> %s::vector) AS similarity
                FROM document_chunks
                {filter_clause}
                ORDER BY embedding <=> %s::vector
                LIMIT %s""",
            [query_embedding] + filter_params + [query_embedding, limit]
        )

        return [
            {
                "document_id": row[0],
                "chunk_index": row[1],
                "content": row[2],
                "metadata": row[3],
                "similarity": float(row[4])
            }
            for row in cur.fetchall()
        ]

-- Instead of adding AI fields directly to the parent table:
CREATE TABLE customer_ai_attributes (
    customer_id UUID REFERENCES customers(id),
    attribute_key VARCHAR(100) NOT NULL,
    attribute_value TEXT,
    confidence FLOAT,
    model_version VARCHAR(50),
    generated_at TIMESTAMPTZ DEFAULT NOW(),
    expires_at TIMESTAMPTZ,  -- AI outputs can go stale
    PRIMARY KEY (customer_id, attribute_key)
);

-- This allows you to:
-- 1. Update AI attributes independently from the customer record
-- 2. Track which model version produced each attribute
-- 3. Expire stale AI outputs and regenerate them
-- 4. Roll back to previous AI-generated values if a model update regresses

import hashlib
import json
import redis
from anthropic import Anthropic

class CachedAnthropicClient:
    """
    Wrapper around Anthropic client with Redis caching.
    Appropriate for deterministic or near-deterministic use cases.
    """

    def __init__(self, cache_ttl_seconds: int = 3600):
        self.client = Anthropic()
        self.cache = redis.Redis()
        self.ttl = cache_ttl_seconds

    def cached_complete(self, model: str, messages: list, system: str = "", max_tokens: int = 1024, temperature: float = 0) -> str:
        """
        Complete with caching. Only cache when temperature=0 (deterministic).
        """
        if temperature > 0:
            # Don't cache non-deterministic outputs
            return self._complete(model, messages, system, max_tokens, temperature)

        cache_key = self._make_cache_key(model, messages, system, max_tokens)

        cached = self.cache.get(cache_key)
        if cached:
            return json.loads(cached)

        result = self._complete(model, messages, system, max_tokens, temperature)
        self.cache.setex(cache_key, self.ttl, json.dumps(result))
        return result

    def _complete(self, model, messages, system, max_tokens, temperature) -> str:
        kwargs = {"model": model, "max_tokens": max_tokens, "messages": messages}
        if system:
            kwargs["system"] = system
        response = self.client.messages.create(**kwargs)
        return response.content[0].text

    def _make_cache_key(self, model: str, messages: list, system: str, max_tokens: int) -> str:
        payload = json.dumps({"model": model, "messages": messages, "system": system, "max_tokens": max_tokens}, sort_keys=True)
        return f"llm_cache:{hashlib.sha256(payload.encode()).hexdigest()}"

import time
from anthropic import Anthropic
from dataclasses import dataclass

@dataclass
class LLMCallMetrics:
    feature: str
    model: str
    input_tokens: int
    output_tokens: int
    latency_ms: int
    cached: bool = False

class InstrumentedAnthropicClient:
    """Anthropic client with cost and latency tracking per feature."""

    COST_PER_MILLION = {
        "claude-sonnet-4-20250514": {"input": 3.0, "output": 15.0},
        "claude-haiku-4-5-20251001": {"input": 0.25, "output": 1.25},
    }

    def __init__(self, metrics_emitter):
        self.client = Anthropic()
        self.metrics = metrics_emitter  # Your metrics system (Datadog, Prometheus, etc.)

    def complete(self, feature: str, model: str, messages: list, **kwargs) -> str:
        start = time.time()

        response = self.client.messages.create(
            model=model, messages=messages, **kwargs
        )

        latency_ms = int((time.time() - start) * 1000)

        m = LLMCallMetrics(
            feature=feature,
            model=model,
            input_tokens=response.usage.input_tokens,
            output_tokens=response.usage.output_tokens,
            latency_ms=latency_ms
        )

        # Emit metrics tagged by feature
        self.metrics.histogram("llm.latency_ms", latency_ms, tags=[f"feature:{feature}", f"model:{model}"])
        self.metrics.increment("llm.input_tokens", m.input_tokens, tags=[f"feature:{feature}"])
        self.metrics.increment("llm.output_tokens", m.output_tokens, tags=[f"feature:{feature}"])

        cost = self._calculate_cost(model, m.input_tokens, m.output_tokens)
        self.metrics.gauge("llm.cost_usd", cost, tags=[f"feature:{feature}"])

        return response.content[0].text

    def _calculate_cost(self, model: str, input_tokens: int, output_tokens: int) -> float:
        rates = self.COST_PER_MILLION.get(model, {"input": 3.0, "output": 15.0})
        return (input_tokens * rates["input"] + output_tokens * rates["output"]) / 1_000_000