core/llm.py

from __future__ import annotations

import json
from abc import ABC, abstractmethod

from core.config import MODEL_BACKEND
from core.risks import HIGH_JARGON, NO_PRODUCT_VISUAL, PRICING_UNCLEAR
from core.schemas import (
    CTA,
    AggregatedReport,
    Archetype,
    ContentSection,
    HeuristicReport,
    Persona,
    RedesignSuggestion,
    UIAnalysis,
    VisualAssessment,
)


class ModelClient(ABC):
    @abstractmethod
    def vision_analyze(self, image_path: str | None, app_description: str) -> UIAnalysis: ...

    @abstractmethod
    def customize_personas(
        self,
        archetypes: list[Archetype],
        app_description: str,
        target_audience: str,
    ) -> list[Persona]: ...

    @abstractmethod
    def run_swarm(
        self,
        personas: list[Persona],
        ui: UIAnalysis,
        heuristics: HeuristicReport,
        flow_steps: list[str],
    ) -> list[tuple[str, str]]: ...

    @abstractmethod
    def generate_redesign(
        self,
        report: AggregatedReport,
        ui: UIAnalysis,
        heuristics: HeuristicReport,
    ) -> RedesignSuggestion: ...


class MockClient(ModelClient):
    def vision_analyze(self, image_path: str | None, app_description: str) -> UIAnalysis:
        text = app_description.lower()
        good = any(term in text for term in ["clear", "good", "simple", "trusted"])
        if good:
            return UIAnalysis(
                headline="Ship customer onboarding in minutes",
                subheadline=(
                    "Create guided product tours with templates, analytics, and a free trial."
                ),
                primary_cta=CTA(
                    text="Start free trial", position="above_fold", visibility="prominent"
                ),
                secondary_ctas=["Watch demo"],
                navigation=["Product", "Pricing", "Customers", "Docs"],
                content_sections=[
                    ContentSection(
                        type="product_visual", summary="Large dashboard screenshot above the fold"
                    ),
                    ContentSection(type="pricing", summary="Free trial and transparent paid tiers"),
                    ContentSection(
                        type="testimonials", summary="Customer logos and two short reviews"
                    ),
                ],
                visual_assessment=VisualAssessment(
                    clutter_level="low",
                    mobile_friendliness="good",
                    text_density="moderate",
                    contrast_issues="none",
                ),
                overall_first_impression="Clear value proposition with a visible free-trial CTA.",
            )
        return UIAnalysis(
            headline="Revolutionary agentic synergy for scalable workflow transformation",
            subheadline=(
                "Leverage our holistic ecosystem to optimize enterprise productivity instantly."
            ),
            primary_cta=CTA(text="Get Started", position="above_fold", visibility="subtle"),
            secondary_ctas=["Learn More", "Contact Sales"],
            navigation=["Platform", "Solutions", "Company"],
            content_sections=[
                ContentSection(type="hero", summary="Text-heavy hero with abstract gradient art"),
                ContentSection(type="features", summary="Dense feature claims without screenshots"),
            ],
            visual_assessment=VisualAssessment(
                clutter_level="high",
                mobile_friendliness="poor",
                text_density="dense",
                contrast_issues="minor",
            ),
            overall_first_impression=(
                "Ambitious but vague; the page does not quickly explain the product."
            ),
        )

    def customize_personas(
        self,
        archetypes: list[Archetype],
        app_description: str,
        target_audience: str,
    ) -> list[Persona]:
        personas: list[Persona] = []
        for idx, archetype in enumerate(archetypes, start=1):
            personas.append(
                Persona(
                    persona_id=f"P{idx:02d}",
                    archetype=archetype.archetype,
                    name=f"Persona {idx}",
                    backstory=f"Evaluating this product for {target_audience or 'their team'}.",
                    goal="Understand the product and decide whether to continue onboarding.",
                    device="mobile_chrome" if "mobile" in archetype.archetype else "desktop_chrome",
                    tolerance=archetype.key_trait,
                    constraints=[archetype.key_trait],
                    brutal_quote_style="direct, specific, first-person",
                )
            )
        return personas

    def run_swarm(
        self,
        personas: list[Persona],
        ui: UIAnalysis,
        heuristics: HeuristicReport,
        flow_steps: list[str],
    ) -> list[tuple[str, str]]:
        from agents.personas import (
            archetype_is_vulnerable,
            issues_for_archetype,
            quote_for_archetype,
        )

        total_steps = max(len(flow_steps), 3)
        raw_results: list[tuple[str, str]] = []
        severe = bool(
            set(heuristics.risk_flags) & {HIGH_JARGON, PRICING_UNCLEAR, NO_PRODUCT_VISUAL}
        )
        for persona in personas:
            fails = severe and archetype_is_vulnerable(persona.archetype, heuristics)
            final_step = 2 if fails else total_steps
            issues = issues_for_archetype(persona.archetype, heuristics)
            result = {
                "persona_id": persona.persona_id,
                "archetype": persona.archetype,
                "completed_task": not fails,
                "final_step_reached": final_step,
                "total_steps": total_steps,
                "step_log": [
                    {
                        "step": 1,
                        "action": "Read the landing page headline",
                        "understanding": f"Saw headline: {ui.headline}",
                        "confusion_level": 0.75 if "HIGH_JARGON" in heuristics.risk_flags else 0.2,
                        "would_continue": True,
                        "issues": [
                            issue
                            for issue in issues
                            if "jargon" in issue.lower() or "headline" in issue.lower()
                        ],
                    },
                    {
                        "step": 2,
                        "action": "Looked for the next onboarding action",
                        "understanding": f"Primary CTA says '{ui.primary_cta.text}'.",
                        "confusion_level": 0.85 if fails else 0.25,
                        "would_continue": not fails,
                        "issues": issues,
                    },
                ],
                "overall_confusion_score": 0.78 if fails else 0.24,
                "overall_trust_score": 0.34 if fails else 0.76,
                "brutal_quote": quote_for_archetype(persona.archetype, ui.primary_cta.text, fails),
                "top_issues": issues,
                "suggested_fix": (
                    "Use plain copy, show the product, and clarify pricing near the CTA."
                )
                if fails
                else "Keep the clear CTA and supporting proof visible above the fold.",
            }
            raw_results.append((persona.persona_id, json.dumps(result)))
        return raw_results

    def generate_redesign(
        self,
        report: AggregatedReport,
        ui: UIAnalysis,
        heuristics: HeuristicReport,
    ) -> RedesignSuggestion:
        return RedesignSuggestion.model_validate(
            {
                "priority_fixes": [
                    {
                        "issue": report.top_issues[0].issue
                        if report.top_issues
                        else "Clarify value proposition",
                        "severity": "high",
                        "affected_segments": [
                            segment.archetype for segment in report.worst_segments[:3]
                        ],
                        "current": ui.headline,
                        "proposed_alternatives": [
                            "Understand your onboarding failures before launch",
                            "See where new users get confused in 60 seconds",
                        ],
                        "reasoning": (
                            "The weakest segments need immediate clarity before "
                            "they will trust the CTA."
                        ),
                    }
                ],
                "onboarding_flow_fix": {
                    "current": [step.label for step in report.failure_funnel],
                    "proposed": [
                        "Show concrete example",
                        "Let user try",
                        "Ask for signup after value",
                    ],
                    "reasoning": "Move proof and value before commitment.",
                },
                "quick_wins": [
                    "Rewrite the headline in plain language under 8 words",
                    "Add pricing or 'Free trial' copy next to the primary CTA",
                    "Show a real product screenshot above the fold",
                    "Reduce competing CTAs to one primary and one secondary action",
                ],
            }
        )


class VLLMClient(MockClient):
    """Placeholder preserving the backend boundary until AMD/vLLM infra is available."""


class HFInferenceClient(MockClient):
    """Placeholder preserving the backend boundary until HF Inference is wired."""


def get_model_client() -> ModelClient:
    if MODEL_BACKEND == "vllm":
        return VLLMClient()
    if MODEL_BACKEND == "hf_api":
        return HFInferenceClient()
    return MockClient()