from anthropic import Anthropic

client = Anthropic()

# Один большой промпт — модель делает всё сразу
bad_prompt = """
Возьми этот отзыв клиента, определи тональность, извлеки ключевые проблемы,
сгенерируй ответ службы поддержки на русском, переведи его на английский,
оцени вероятность что клиент уйдёт и дай рекомендации команде продукта.
Отзыв: {review}
"""

# Проблемы: сложно контролировать каждый шаг, ошибка на одном шаге ломает всё,
# нельзя кэшировать промежуточные результаты, дорого если нужен только один шаг

from dataclasses import dataclass
from typing import Any, Callable


@dataclass
class PipelineStep:
    name: str
    fn: Callable
    model: str = "claude-haiku-4-5"  # дешевле для простых шагов


class LLMPipeline:
    def __init__(self):
        self.steps: list[PipelineStep] = []
        self.results: dict[str, Any] = {}

    def add_step(self, name: str, fn: Callable, model: str = "claude-haiku-4-5"):
        self.steps.append(PipelineStep(name=name, fn=fn, model=model))
        return self  # для chaining

    def run(self, initial_input: Any) -> dict:
        current_input = initial_input

        for step in self.steps:
            print(f"Шаг: {step.name}...")
            result = step.fn(current_input, self.results)
            self.results[step.name] = result
            current_input = result  # вывод идёт в следующий шаг

        return self.results

def analyze_sentiment(review: str, _: dict) -> dict:
    response = client.messages.create(
        model="claude-haiku-4-5",  # простая задача — дешёвая модель
        max_tokens=128,
        system="Отвечай только JSON. Без пояснений.",
        messages=[{
            "role": "user",
            "content": f"""Определи тональность отзыва.
JSON: {{"sentiment": "positive|negative|neutral", "score": 1-10, "is_urgent": true|false}}

Отзыв: {review}"""
        }]
    )
    import json, re
    text = re.sub(r'```(?:json)?|```', '', response.content[0].text).strip()
    return json.loads(text)


def extract_issues(review: str, context: dict) -> list[str]:
    sentiment = context.get("analyze_sentiment", {})

    # Пропускаем если позитивный отзыв
    if sentiment.get("sentiment") == "positive":
        return []

    response = client.messages.create(
        model="claude-haiku-4-5",
        max_tokens=256,
        messages=[{
            "role": "user",
            "content": f"""Извлеки конкретные проблемы из отзыва. 
Отвечай JSON-массивом строк: ["проблема 1", "проблема 2"]

Отзыв: {review}"""
        }]
    )
    import json, re
    text = re.sub(r'```(?:json)?|```', '', response.content[0].text).strip()
    return json.loads(text)


def generate_response(review: str, context: dict) -> str:
    sentiment = context.get("analyze_sentiment", {})
    issues = context.get("extract_issues", [])

    issues_text = "\n".join(f"- {issue}" for issue in issues) if issues else "нет"

    response = client.messages.create(
        model="claude-sonnet-4-5",  # важный шаг — более качественная модель
        max_tokens=512,
        system="""Пишешь ответы службы поддержки.
Тон: дружелюбный, конкретный, без шаблонных фраз.
Не обещай то что не можешь выполнить.""",
        messages=[{
            "role": "user",
            "content": f"""Напиши ответ на отзыв.
Тональность отзыва: {sentiment.get('sentiment')}
Выявленные проблемы: {issues_text}

Отзыв: {review}"""
        }]
    )
    return response.content[0].text


def churn_risk(review: str, context: dict) -> dict:
    sentiment = context.get("analyze_sentiment", {})

    response = client.messages.create(
        model="claude-haiku-4-5",
        max_tokens=128,
        messages=[{
            "role": "user",
            "content": f"""Оцени риск ухода клиента.
JSON: {{"churn_risk": "high|medium|low", "reason": "краткое объяснение"}}

Тональность: {sentiment}
Отзыв: {review}"""
        }]
    )
    import json, re
    text = re.sub(r'```(?:json)?|```', '', response.content[0].text).strip()
    return json.loads(text)


# Собираем pipeline
pipeline = LLMPipeline()
pipeline\
    .add_step("analyze_sentiment", analyze_sentiment)\
    .add_step("extract_issues", extract_issues)\
    .add_step("generate_response", generate_response, model="claude-sonnet-4-5")\
    .add_step("churn_risk", churn_risk)

review = "Приложение постоянно вылетает при оплате. Уже третий раз теряю корзину. Очень расстроен."
results = pipeline.run(review)

print(f"Тональность: {results['analyze_sentiment']}")
print(f"Проблемы: {results['extract_issues']}")
print(f"Риск ухода: {results['churn_risk']}")
print(f"\nОтвет:\n{results['generate_response']}")

class ConditionalPipeline(LLMPipeline):

    def add_conditional_step(
        self,
        name: str,
        condition: Callable[[dict], bool],
        fn: Callable,
        skip_value: Any = None
    ):
        def conditional_fn(input_data, context):
            if condition(context):
                return fn(input_data, context)
            print(f"  Шаг {name} пропущен по условию")
            return skip_value

        self.steps.append(PipelineStep(name=name, fn=conditional_fn))
        return self


pipeline = ConditionalPipeline()
pipeline\
    .add_step("analyze_sentiment", analyze_sentiment)\
    .add_conditional_step(
        name="extract_issues",
        # Извлекаем проблемы только если отзыв негативный
        condition=lambda ctx: ctx.get("analyze_sentiment", {}).get("sentiment") != "positive",
        fn=extract_issues,
        skip_value=[]
    )\
    .add_step("generate_response", generate_response)