"""Token economics analysis for AI-to-AI communication patterns.

Analyzes how current tokenizers encode common multi-agent message patterns
and quantifies token waste across categories.
"""

import json
import tiktoken
from dataclasses import dataclass, field
from typing import Optional


@dataclass
class TokenAnalysis:
    text: str
    token_count: int
    tokens: list[int]
    token_strings: list[str]
    category: str = ""
    waste_breakdown: dict[str, int] = field(default_factory=dict)


@dataclass
class CompressionResult:
    scenario: str
    natural_language: TokenAnalysis
    json_format: Optional[TokenAnalysis] = None
    condensa_format: Optional[TokenAnalysis] = None

    @property
    def nl_to_json_ratio(self) -> Optional[float]:
        if self.json_format:
            return self.json_format.token_count / self.natural_language.token_count
        return None

    @property
    def nl_to_condensa_ratio(self) -> Optional[float]:
        if self.condensa_format:
            return self.condensa_format.token_count / self.natural_language.token_count
        return None


class TokenEconomicsAnalyzer:
    """Analyzes token usage patterns in AI-to-AI communication."""

    def __init__(self, encoding_name: str = "cl100k_base"):
        self.enc = tiktoken.get_encoding(encoding_name)

    def analyze(self, text: str, category: str = "") -> TokenAnalysis:
        """Tokenize text and return detailed analysis."""
        tokens = self.enc.encode(text)
        token_strings = [self.enc.decode([t]) for t in tokens]
        return TokenAnalysis(
            text=text,
            token_count=len(tokens),
            tokens=tokens,
            token_strings=token_strings,
            category=category,
        )

    def count_tokens(self, text: str) -> int:
        return len(self.enc.encode(text))

    def classify_waste(self, analysis: TokenAnalysis) -> dict[str, int]:
        """Classify tokens into waste categories."""
        waste = {
            "structural_overhead": 0,  # JSON braces, quotes, colons
            "politeness_filler": 0,    # "please", "could you", "thank you"
            "redundant_context": 0,    # Re-stated information
            "ambiguity_hedging": 0,    # "maybe", "possibly", "I think"
            "format_compliance": 0,    # Whitespace, newlines for readability
            "semantic_content": 0,     # Actual meaningful tokens
        }

        filler_words = {
            "please", "could", "would", "you", "kindly", "thank",
            "thanks", "hi", "hello", "hey", "sure", "okay", "ok",
            "great", "nice", "perfect", "awesome", "certainly",
            "absolutely", "of course", "no problem", "happy to",
        }
        hedge_words = {
            "maybe", "perhaps", "possibly", "might", "probably",
            "I think", "it seems", "apparently", "likely", "unlikely",
            "roughly", "approximately", "around", "about",
        }
        structural_chars = set('{}[]":,')

        text_lower = analysis.text.lower()
        for ts in analysis.token_strings:
            stripped = ts.strip().lower()
            if all(c in structural_chars or c.isspace() for c in ts) and ts.strip():
                waste["structural_overhead"] += 1
            elif stripped in filler_words:
                waste["politeness_filler"] += 1
            elif stripped in hedge_words:
                waste["ambiguity_hedging"] += 1
            elif ts.isspace() or ts == "\n":
                waste["format_compliance"] += 1
            else:
                waste["semantic_content"] += 1

        analysis.waste_breakdown = waste
        return waste

    def compare_formats(
        self,
        scenario: str,
        natural_lang: str,
        json_format: str = "",
        condensa_format: str = "",
    ) -> CompressionResult:
        """Compare token counts across formats for the same semantic content."""
        nl_analysis = self.analyze(natural_lang, "natural_language")
        self.classify_waste(nl_analysis)

        result = CompressionResult(
            scenario=scenario,
            natural_language=nl_analysis,
        )

        if json_format:
            json_analysis = self.analyze(json_format, "json")
            self.classify_waste(json_analysis)
            result.json_format = json_analysis

        if condensa_format:
            condensa_analysis = self.analyze(condensa_format, "condensa")
            self.classify_waste(condensa_analysis)
            result.condensa_format = condensa_analysis

        return result

    def format_report(self, result: CompressionResult) -> str:
        """Generate a human-readable comparison report."""
        lines = [f"### Scenario: {result.scenario}", ""]

        def _fmt(label: str, analysis: TokenAnalysis) -> list[str]:
            out = [f"**{label}**: {analysis.token_count} tokens"]
            out.append(f"  Text: `{analysis.text[:120]}{'...' if len(analysis.text) > 120 else ''}`")
            if analysis.waste_breakdown:
                out.append(f"  Waste: {analysis.waste_breakdown}")
            return out

        lines.extend(_fmt("Natural Language", result.natural_language))
        if result.json_format:
            lines.extend(_fmt("JSON", result.json_format))
            ratio = result.nl_to_json_ratio
            lines.append(f"  JSON/NL ratio: {ratio:.2f}x")
        if result.condensa_format:
            lines.extend(_fmt("Condensa", result.condensa_format))
            ratio = result.nl_to_condensa_ratio
            lines.append(f"  Condensa/NL ratio: {ratio:.2f}x ({(1-ratio)*100:.1f}% reduction)")

        return "\n".join(lines)


# ── Representative multi-agent conversation samples ──────────────────────────

SAMPLE_CONVERSATIONS = [
    # 1. Simple task delegation
    {
        "name": "task_delegation_simple",
        "natural_language": "Hey Agent B, could you please search the web for the latest news about SpaceX Starship and give me a summary of the top 3 results? Format them as bullet points please.",
        "json": '{"action": "search_and_summarize", "query": "SpaceX Starship latest news", "result_count": 3, "format": "bullet_points", "source": "web"}',
    },
    # 2. Status update
    {
        "name": "status_update",
        "natural_language": "Hi, just wanted to let you know that I've completed the data extraction step. I processed 15,432 records successfully, with 23 failures due to malformed dates. The output is stored in /tmp/extracted_data.parquet. Ready for the next step whenever you are.",
        "json": '{"type": "status_update", "step": "data_extraction", "status": "completed", "records_processed": 15432, "failures": 23, "failure_reason": "malformed_dates", "output_path": "/tmp/extracted_data.parquet", "ready_for_next": true}',
    },
    # 3. Error handling
    {
        "name": "error_with_retry",
        "natural_language": "I encountered an error while trying to call the payment API. The server returned a 503 Service Unavailable error. I'm going to wait for 30 seconds and then retry. This will be my second attempt out of a maximum of 3 retries. If all retries fail, I'll fall back to the cached data from the last successful call.",
        "json": '{"type": "error_report", "service": "payment_api", "error_code": 503, "error_type": "service_unavailable", "retry": {"attempt": 2, "max_attempts": 3, "delay_seconds": 30}, "fallback": {"strategy": "cached_data", "source": "last_successful_call"}}',
    },
    # 4. Result passing
    {
        "name": "result_passing",
        "natural_language": "Here are the results of the sentiment analysis you requested. The overall sentiment is positive with a confidence score of 0.87. Breaking it down by section: the introduction is neutral (0.52), the main body is strongly positive (0.94), and the conclusion is mildly positive (0.71). I also detected 3 named entities: SpaceX, Elon Musk, and Starlink.",
        "json": '{"type": "result", "task": "sentiment_analysis", "overall": {"sentiment": "positive", "confidence": 0.87}, "sections": [{"name": "introduction", "sentiment": "neutral", "score": 0.52}, {"name": "main_body", "sentiment": "strongly_positive", "score": 0.94}, {"name": "conclusion", "sentiment": "mildly_positive", "score": 0.71}], "entities": ["SpaceX", "Elon Musk", "Starlink"]}',
    },
    # 5. Context update
    {
        "name": "context_update",
        "natural_language": "I need to update the shared context. The user has changed their preference from English to Spanish for all output. Also, the budget constraint has been increased from $500 to $750. Everything else remains the same as before.",
        "json": '{"type": "context_update", "changes": [{"field": "output_language", "old": "English", "new": "Spanish"}, {"field": "budget_constraint", "old": 500, "new": 750}], "unchanged": "all_other_fields"}',
    },
    # 6. Complex workflow delegation
    {
        "name": "complex_workflow",
        "natural_language": "I need you to handle a multi-step workflow. First, fetch the user's purchase history from the database for the last 90 days. Then, run a clustering algorithm to identify their top 3 product categories. After that, query our recommendation engine for 5 products in each of those categories, excluding items they've already purchased. Finally, rank the results by predicted click-through rate and return the top 10 overall. If the database query takes longer than 10 seconds, use the cached profile instead.",
        "json": '{"type": "workflow", "steps": [{"id": 1, "action": "db_query", "table": "purchase_history", "filter": {"days": 90}, "timeout": 10, "fallback": "cached_profile"}, {"id": 2, "action": "cluster", "input": "step_1", "algorithm": "top_categories", "k": 3}, {"id": 3, "action": "recommend", "input": "step_2", "per_category": 5, "exclude": "purchased_items"}, {"id": 4, "action": "rank", "input": "step_3", "by": "predicted_ctr", "limit": 10}]}',
    },
    # 7. Negotiation / capability check
    {
        "name": "capability_negotiation",
        "natural_language": "Before I send you this task, I need to check: can you handle image processing? Specifically, I need OCR on a scanned PDF document that's about 50 pages long. The output should be structured JSON with page numbers. If you can't do OCR, can you at least extract the text layer if one exists?",
        "json": '{"type": "capability_check", "required": [{"capability": "ocr", "input": "scanned_pdf", "pages": 50, "output": "structured_json_with_pages"}], "fallback": [{"capability": "text_extraction", "condition": "no_ocr", "input": "pdf_text_layer"}]}',
    },
    # 8. Partial result with continuation
    {
        "name": "partial_result",
        "natural_language": "I've completed processing the first 500 items out of 2000 total. Here's what I have so far: 340 items classified as Category A, 112 as Category B, and 48 as Category C. My confidence is high for A and B but moderate for C due to ambiguous features. I'm continuing to process the remaining 1500 items. Estimated completion: 3 more iterations. Should I continue with the same parameters or would you like to adjust the classification thresholds?",
        "json": '{"type": "partial_result", "progress": {"completed": 500, "total": 2000, "remaining": 1500, "eta_iterations": 3}, "results_so_far": {"category_a": 340, "category_b": 112, "category_c": 48}, "confidence": {"a": "high", "b": "high", "c": "moderate", "c_reason": "ambiguous_features"}, "continuation": {"status": "in_progress", "query": "adjust_thresholds?"}}',
    },
    # 9. Correction / revision request
    {
        "name": "correction_request",
        "natural_language": "The previous response was partially incorrect. The data for Q3 was accurate, but Q4 numbers were wrong. The correct Q4 revenue is $2.3M, not $1.8M. Please regenerate the summary chart with this correction. Keep everything else the same. Also, add a footnote explaining the discrepancy was due to a late invoice adjustment.",
        "json": '{"type": "correction", "target": "previous_response", "accurate": ["Q3"], "corrections": [{"field": "Q4_revenue", "was": 1800000, "should_be": 2300000}], "action": "regenerate_summary_chart", "preserve": "everything_else", "additions": [{"type": "footnote", "content": "discrepancy due to late invoice adjustment"}]}',
    },
    # 10. Parallel task dispatch
    {
        "name": "parallel_dispatch",
        "natural_language": "I need three things done simultaneously. Task 1: Translate this document from English to French. Task 2: Generate a summary of no more than 200 words. Task 3: Extract all dates, monetary amounts, and person names as structured data. All three tasks should use the same input document. When all are complete, merge the results into a single response with sections for each task.",
        "json": '{"type": "parallel_dispatch", "input": "current_document", "tasks": [{"id": 1, "action": "translate", "from": "en", "to": "fr"}, {"id": 2, "action": "summarize", "max_words": 200}, {"id": 3, "action": "extract_entities", "types": ["dates", "monetary_amounts", "person_names"]}], "on_complete": {"action": "merge", "format": "sectioned_response"}}',
    },
    # 11. Conditional branching
    {
        "name": "conditional_branching",
        "natural_language": "Analyze the uploaded image. If it contains text, perform OCR and return the extracted text. If it contains a chart or graph, describe the data trends and extract any visible data points. If it's a photograph of a real scene, describe what's in the image and identify any notable objects or people. In all cases, also return the image dimensions and dominant colors.",
        "json": '{"type": "conditional", "input": "uploaded_image", "branches": [{"condition": "contains_text", "action": "ocr", "output": "extracted_text"}, {"condition": "contains_chart", "actions": ["describe_trends", "extract_data_points"]}, {"condition": "is_photograph", "actions": ["describe_scene", "identify_objects_people"]}], "always": ["return_dimensions", "return_dominant_colors"]}',
    },
    # 12. Authentication/permission context
    {
        "name": "auth_context",
        "natural_language": "You're operating under the following constraints for this session: the user has read-only access to the production database, read-write access to the staging database, and no access to the admin panel. Their API rate limit is 100 requests per minute. The session token expires in 45 minutes. Please factor these constraints into any actions you take.",
        "json": '{"type": "session_context", "permissions": {"production_db": "read", "staging_db": "read_write", "admin_panel": "none"}, "rate_limit": {"requests": 100, "per": "minute"}, "session": {"expires_in_minutes": 45}}',
    },
    # 13. Aggregation request
    {
        "name": "aggregation",
        "natural_language": "I've received results from 5 different agents. Agent 1 says the answer is Paris with 90% confidence. Agent 2 says Paris with 85% confidence. Agent 3 says Lyon with 60% confidence. Agent 4 says Paris with 95% confidence. Agent 5 abstained due to insufficient context. Please aggregate these results using weighted voting, determine the consensus answer, calculate the overall confidence, and flag any significant dissent.",
        "json": '{"type": "aggregate", "method": "weighted_vote", "inputs": [{"agent": 1, "answer": "Paris", "confidence": 0.90}, {"agent": 2, "answer": "Paris", "confidence": 0.85}, {"agent": 3, "answer": "Lyon", "confidence": 0.60}, {"agent": 4, "answer": "Paris", "confidence": 0.95}, {"agent": 5, "answer": null, "reason": "insufficient_context"}], "output": ["consensus_answer", "overall_confidence", "dissent_flags"]}',
    },
    # 14. Schema/format negotiation
    {
        "name": "format_negotiation",
        "natural_language": "For the next set of results you send me, please use the following format: each result should have a title, a relevance score between 0 and 1, a snippet of no more than 100 characters, and a source URL. Return them as a flat array, not nested. Maximum 20 results. If you can't provide a source URL for an item, omit that item entirely rather than returning a null URL.",
        "json": '{"type": "format_spec", "schema": {"type": "array", "max_items": 20, "item": {"title": "string", "relevance": {"type": "float", "min": 0, "max": 1}, "snippet": {"type": "string", "max_chars": 100}, "source_url": {"type": "string", "required": true}}}, "rules": [{"if": "no_source_url", "then": "omit_item"}]}',
    },
    # 15. Memory/state management
    {
        "name": "state_management",
        "natural_language": "Please store the following in our shared conversation state: the user's name is Alice Chen, they are working on Project Falcon, their preferred output format is markdown, and they have approved a budget of $10,000 for API calls. Reference ID for this session is SESS-2024-0847. I may refer to these later, so keep them accessible.",
        "json": '{"type": "state_store", "entries": [{"key": "user_name", "value": "Alice Chen"}, {"key": "project", "value": "Project Falcon"}, {"key": "output_format", "value": "markdown"}, {"key": "api_budget", "value": 10000, "unit": "USD"}, {"key": "session_id", "value": "SESS-2024-0847"}], "persistence": "session"}',
    },
    # 16. Cancellation/abort
    {
        "name": "cancellation",
        "natural_language": "Cancel the current task immediately. The user has changed their mind and no longer needs the report generated. Discard any partial results. Do not save anything to the database. Acknowledge the cancellation and confirm that no side effects occurred.",
        "json": '{"type": "cancel", "target": "current_task", "discard_partial": true, "prevent_side_effects": ["database_writes"], "require_confirmation": true}',
    },
    # 17. Priority/urgency signaling
    {
        "name": "priority_signal",
        "natural_language": "URGENT: Drop whatever you're currently working on. We have a production incident. The authentication service is returning 500 errors for approximately 30% of login requests. I need you to immediately check the last 5 minutes of auth service logs, identify the error pattern, and suggest a mitigation. This takes priority over everything else in your queue.",
        "json": '{"type": "priority_override", "urgency": "critical", "context": "production_incident", "service": "auth", "symptom": {"error_code": 500, "affected_percentage": 30, "operation": "login"}, "actions": [{"action": "check_logs", "service": "auth", "timeframe_minutes": 5}, {"action": "identify_pattern"}, {"action": "suggest_mitigation"}], "queue_behavior": "preempt_all"}',
    },
    # 18. Feedback loop
    {
        "name": "feedback_loop",
        "natural_language": "The summary you generated was good but needs some adjustments. The tone should be more formal—remove any casual language. The length is about right at 3 paragraphs. Add more specific numbers from the source data where possible. The structure is fine, don't reorganize. Please regenerate with these adjustments and show me a diff of what changed.",
        "json": '{"type": "feedback", "target": "previous_summary", "adjustments": [{"aspect": "tone", "change": "more_formal", "detail": "remove_casual_language"}, {"aspect": "length", "change": "keep", "current": "3_paragraphs"}, {"aspect": "content", "change": "add_specific_numbers", "source": "original_data"}, {"aspect": "structure", "change": "keep"}], "action": "regenerate", "show_diff": true}',
    },
    # 19. Tool/function call result
    {
        "name": "tool_result",
        "natural_language": "I called the weather API and here are the results: New York City current temperature is 72°F, humidity 65%, wind speed 12mph from the northwest. The forecast for tomorrow shows a high of 78°F and a low of 61°F with a 40% chance of afternoon showers. The API responded in 230ms and I used 1 of our 1000 daily quota.",
        "json": '{"type": "tool_result", "tool": "weather_api", "location": "New York City", "current": {"temp_f": 72, "humidity_pct": 65, "wind": {"speed_mph": 12, "direction": "NW"}}, "forecast": {"tomorrow": {"high_f": 78, "low_f": 61, "precipitation": {"chance_pct": 40, "type": "showers", "timing": "afternoon"}}}, "meta": {"response_ms": 230, "quota": {"used": 1, "total": 1000, "period": "daily"}}}',
    },
    # 20. End-of-chain summary
    {
        "name": "chain_summary",
        "natural_language": "All tasks in the pipeline are now complete. Here's the final summary: we processed 50,000 customer records through 4 stages. Stage 1 (validation) passed 48,200 records. Stage 2 (enrichment) added demographic data to 47,100 records. Stage 3 (scoring) assigned risk scores to all enriched records. Stage 4 (output) generated the final report and uploaded it to S3 at s3://reports/customer-risk-2024-03.csv. Total execution time: 4 minutes 23 seconds. No critical errors. 1,800 records were excluded due to missing required fields. The pipeline run ID is PIPE-20240315-001 for your reference.",
        "json": '{"type": "pipeline_complete", "run_id": "PIPE-20240315-001", "input_records": 50000, "stages": [{"name": "validation", "passed": 48200}, {"name": "enrichment", "enriched": 47100, "data_added": "demographics"}, {"name": "scoring", "scored": 47100, "type": "risk_score"}, {"name": "output", "format": "csv", "location": "s3://reports/customer-risk-2024-03.csv"}], "excluded": {"count": 1800, "reason": "missing_required_fields"}, "execution_time_seconds": 263, "critical_errors": 0}',
    },
]


def run_full_analysis():
    """Run complete token economics analysis on all samples."""
    analyzer = TokenEconomicsAnalyzer()

    print("=" * 80)
    print("CONDENSA — Token Economics Audit")
    print("Tokenizer: cl100k_base (GPT-4 / Claude compatible BPE)")
    print("=" * 80)
    print()

    total_nl_tokens = 0
    total_json_tokens = 0
    waste_totals = {
        "structural_overhead": 0,
        "politeness_filler": 0,
        "redundant_context": 0,
        "ambiguity_hedging": 0,
        "format_compliance": 0,
        "semantic_content": 0,
    }

    results = []

    for sample in SAMPLE_CONVERSATIONS:
        result = analyzer.compare_formats(
            scenario=sample["name"],
            natural_lang=sample["natural_language"],
            json_format=sample.get("json", ""),
        )
        results.append(result)

        total_nl_tokens += result.natural_language.token_count
        if result.json_format:
            total_json_tokens += result.json_format.token_count

        for k, v in result.natural_language.waste_breakdown.items():
            waste_totals[k] += v

        print(analyzer.format_report(result))
        print()

    # Summary
    print("=" * 80)
    print("AGGREGATE RESULTS")
    print("=" * 80)
    print(f"\nTotal Natural Language tokens: {total_nl_tokens}")
    print(f"Total JSON tokens:            {total_json_tokens}")
    print(f"JSON/NL ratio:                {total_json_tokens/total_nl_tokens:.2f}x")
    print(f"\nNatural Language Waste Breakdown (token counts):")
    for k, v in waste_totals.items():
        pct = v / total_nl_tokens * 100
        print(f"  {k:25s}: {v:5d} ({pct:5.1f}%)")

    # Estimate ideal encoding
    semantic = waste_totals["semantic_content"]
    print(f"\nSemantic content tokens: {semantic} ({semantic/total_nl_tokens*100:.1f}% of NL)")
    print(f"Theoretical minimum (semantic only): {semantic} tokens")
    print(f"Theoretical max compression: {(1 - semantic/total_nl_tokens)*100:.1f}% reduction")

    return results


if __name__ == "__main__":
    run_full_analysis()
