# -*- coding: utf-8 -*-
"""
Legal-i: Argument Generator
============================
Causal argument generation for legal cases.

Features:
- IRAC structure generation (Issue, Rule, Application, Conclusion)
- Counter-argument anticipation
- Rebuttal generation
- Argument strength scoring
- Rhetorical strategy selection

TAU Signal Integration:
- τ (Tau): Argument flow / momentum
- φ (Phi): Balance between offense and defense
- ψ (Psi): Logical consistency
- ξ (Xi): Novelty / unusual arguments
- Ω (Omega): Overall argument quality
"""

import numpy as np
from dataclasses import dataclass, field
from typing import Dict, List, Optional, Any, Tuple
from enum import Enum, auto
import logging

logger = logging.getLogger(__name__)


class ArgumentType(Enum):
    """Types of legal arguments."""

    SUBSTANTIVE = "substantive"  # On the merits
    PROCEDURAL = "procedural"  # Procedural grounds
    FACTUAL = "factual"  # Based on facts
    LEGAL = "legal"  # Based on law
    POLICY = "policy"  # Policy-based
    EQUITABLE = "equitable"  # Equity arguments
    CONSTITUTIONAL = "constitutional"  # Constitutional grounds


class ArgumentStrength(Enum):
    """Strength levels for arguments."""

    STRONG = "strong"  # Very persuasive
    MODERATE = "moderate"  # Reasonably persuasive
    WEAK = "weak"  # Marginally persuasive
    SPECULATIVE = "speculative"  # Uncertain


class RhetoricalStrategy(Enum):
    """Rhetorical strategies."""

    LOGOS = "logos"  # Logic and reason
    PATHOS = "pathos"  # Emotional appeal
    ETHOS = "ethos"  # Credibility/authority
    KAIROS = "kairos"  # Timing/context


@dataclass
class IRACElement:
    """An element of IRAC analysis."""

    issue: str
    rule: str
    application: str
    conclusion: str

    # Supporting details
    citations: List[str] = field(default_factory=list)
    key_facts: List[str] = field(default_factory=list)

    def to_dict(self) -> Dict[str, Any]:
        return {
            "issue": self.issue,
            "rule": self.rule,
            "application": self.application,
            "conclusion": self.conclusion,
            "citations": self.citations,
            "key_facts": self.key_facts,
        }


@dataclass
class LegalArgument:
    """A structured legal argument."""

    id: str
    title: str
    argument_type: ArgumentType
    strength: ArgumentStrength

    # Content
    thesis: str
    supporting_points: List[str] = field(default_factory=list)
    evidence: List[str] = field(default_factory=list)
    citations: List[str] = field(default_factory=list)

    # IRAC structure
    irac: Optional[IRACElement] = None

    # Strategy
    rhetorical_strategy: RhetoricalStrategy = RhetoricalStrategy.LOGOS

    # Scoring
    persuasiveness: float = 0.5
    relevance: float = 0.5
    novelty: float = 0.5

    # Counter-arguments
    anticipated_rebuttals: List[str] = field(default_factory=list)
    rebuttal_responses: List[str] = field(default_factory=list)

    # TAU signals
    tau: float = 0.5
    psi: float = 0.7

    def compute_score(self) -> float:
        """Compute overall argument score."""
        strength_weights = {
            ArgumentStrength.STRONG: 1.0,
            ArgumentStrength.MODERATE: 0.7,
            ArgumentStrength.WEAK: 0.4,
            ArgumentStrength.SPECULATIVE: 0.2,
        }
        base = strength_weights.get(self.strength, 0.5)
        return (
            0.4 * base
            + 0.3 * self.persuasiveness
            + 0.2 * self.relevance
            + 0.1 * self.novelty
        )

    def to_dict(self) -> Dict[str, Any]:
        return {
            "id": self.id,
            "title": self.title,
            "type": self.argument_type.value,
            "strength": self.strength.value,
            "thesis": self.thesis,
            "supporting_points": self.supporting_points,
            "evidence": self.evidence,
            "citations": self.citations,
            "irac": self.irac.to_dict() if self.irac else None,
            "strategy": self.rhetorical_strategy.value,
            "score": self.compute_score(),
            "anticipated_rebuttals": self.anticipated_rebuttals,
            "signals": {"tau": self.tau, "psi": self.psi},
        }


@dataclass
class ArgumentSet:
    """A complete set of arguments for a case."""

    case_id: str
    side: str  # "plaintiff" or "defendant"

    main_arguments: List[LegalArgument] = field(default_factory=list)
    alternative_arguments: List[LegalArgument] = field(default_factory=list)
    rebuttal_arguments: List[LegalArgument] = field(default_factory=list)

    # TAU signals
    tau: float = 0.5
    phi: float = 0.618
    psi: float = 0.7
    xi: float = 0.9

    def compute_omega(self) -> float:
        """Compute overall argument set quality."""
        all_args = self.main_arguments + self.alternative_arguments
        if not all_args:
            return 0.5

        avg_score = np.mean([a.compute_score() for a in all_args])
        coverage = min(len(all_args) / 5, 1.0)  # Aim for 5 arguments

        return 0.4 * avg_score + 0.3 * self.psi + 0.2 * coverage + 0.1 * self.phi

    def to_dict(self) -> Dict[str, Any]:
        return {
            "case_id": self.case_id,
            "side": self.side,
            "main_arguments": [a.to_dict() for a in self.main_arguments],
            "alternative_arguments": [a.to_dict() for a in self.alternative_arguments],
            "rebuttal_arguments": [a.to_dict() for a in self.rebuttal_arguments],
            "signals": {
                "tau": self.tau,
                "phi": self.phi,
                "psi": self.psi,
                "xi": self.xi,
                "omega": self.compute_omega(),
            },
        }


class ArgumentGenerator:
    """
    Legal argument generation engine.

    Features:
    - IRAC-based argument structuring
    - Multi-strategy argument generation
    - Counter-argument anticipation
    - Argument optimization
    """

    PHI = 0.618033988749895

    # Argument templates by type
    ARGUMENT_TEMPLATES = {
        ArgumentType.SUBSTANTIVE: {
            "structure": ["establish_elements", "apply_facts", "distinguish_precedent"],
            "typical_strength": ArgumentStrength.STRONG,
        },
        ArgumentType.PROCEDURAL: {
            "structure": ["identify_rule", "show_violation", "demonstrate_prejudice"],
            "typical_strength": ArgumentStrength.MODERATE,
        },
        ArgumentType.POLICY: {
            "structure": ["identify_policy", "show_alignment", "predict_consequences"],
            "typical_strength": ArgumentStrength.MODERATE,
        },
        ArgumentType.EQUITABLE: {
            "structure": ["establish_hardship", "show_balance", "appeal_fairness"],
            "typical_strength": ArgumentStrength.WEAK,
        },
    }

    def __init__(self):
        self.argument_sets: Dict[str, ArgumentSet] = {}
        self.generated_arguments: Dict[str, LegalArgument] = {}

        logger.info("Initialized ArgumentGenerator")

    def create_argument_set(
        self,
        case_id: str,
        side: str,
    ) -> ArgumentSet:
        """Create a new argument set for a case."""
        arg_set = ArgumentSet(
            case_id=case_id,
            side=side,
        )

        key = f"{case_id}_{side}"
        self.argument_sets[key] = arg_set
        return arg_set

    def generate_argument(
        self,
        case_id: str,
        side: str,
        argument_type: ArgumentType,
        thesis: str,
        facts: List[str] = None,
        rules: List[str] = None,
        citations: List[str] = None,
        strategy: RhetoricalStrategy = RhetoricalStrategy.LOGOS,
    ) -> LegalArgument:
        """Generate a structured legal argument."""
        arg_id = f"arg_{len(self.generated_arguments)}"

        # Generate IRAC structure
        irac = self._generate_irac(thesis, facts or [], rules or [])

        # Determine strength based on support
        strength = self._assess_strength(facts, rules, citations)

        # Generate supporting points
        supporting = self._generate_supporting_points(
            argument_type, thesis, facts or []
        )

        # Anticipate counter-arguments
        rebuttals = self._anticipate_rebuttals(argument_type, thesis)

        argument = LegalArgument(
            id=arg_id,
            title=f"{argument_type.value.title()} Argument",
            argument_type=argument_type,
            strength=strength,
            thesis=thesis,
            supporting_points=supporting,
            evidence=facts or [],
            citations=citations or [],
            irac=irac,
            rhetorical_strategy=strategy,
            anticipated_rebuttals=rebuttals,
        )

        # Compute scores
        argument.persuasiveness = self._score_persuasiveness(argument)
        argument.relevance = 0.7  # Would analyze against case issues
        argument.novelty = self._score_novelty(argument)

        # Update TAU signals
        argument.tau = 0.5 + argument.compute_score() * 0.5
        argument.psi = argument.persuasiveness

        self.generated_arguments[arg_id] = argument

        # Add to argument set
        key = f"{case_id}_{side}"
        if key in self.argument_sets:
            self.argument_sets[key].main_arguments.append(argument)

        return argument

    def _generate_irac(
        self,
        thesis: str,
        facts: List[str],
        rules: List[str],
    ) -> IRACElement:
        """Generate IRAC structure from thesis and facts."""
        # Issue: convert thesis to question form
        issue = f"Whether {thesis.lower().rstrip('.')}"
        if not issue.endswith("?"):
            issue += "?"

        # Rule: combine provided rules
        if rules:
            rule = " ".join(rules)
        else:
            rule = "The applicable legal standard requires analysis of the relevant factors."

        # Application: apply facts to rule
        if facts:
            application = f"Here, the facts demonstrate that {facts[0].lower()}. "
            if len(facts) > 1:
                application += f"Additionally, {facts[1].lower()}."
        else:
            application = (
                "Applying the legal standard to these facts supports the conclusion."
            )

        # Conclusion: restate thesis as conclusion
        conclusion = f"Therefore, {thesis.lower()}"
        if not conclusion.endswith("."):
            conclusion += "."

        return IRACElement(
            issue=issue,
            rule=rule,
            application=application,
            conclusion=conclusion,
            key_facts=facts[:3] if facts else [],
        )

    def _assess_strength(
        self,
        facts: List[str],
        rules: List[str],
        citations: List[str],
    ) -> ArgumentStrength:
        """Assess argument strength based on support."""
        score = 0

        # Facts add strength
        score += min(len(facts or []) * 0.1, 0.3)

        # Rules add strength
        score += min(len(rules or []) * 0.15, 0.3)

        # Citations add significant strength
        score += min(len(citations or []) * 0.2, 0.4)

        if score >= 0.7:
            return ArgumentStrength.STRONG
        elif score >= 0.4:
            return ArgumentStrength.MODERATE
        elif score >= 0.2:
            return ArgumentStrength.WEAK
        else:
            return ArgumentStrength.SPECULATIVE

    def _generate_supporting_points(
        self,
        arg_type: ArgumentType,
        thesis: str,
        facts: List[str],
    ) -> List[str]:
        """Generate supporting points for an argument."""
        points = []

        template = self.ARGUMENT_TEMPLATES.get(arg_type, {})
        structure = template.get("structure", ["point_1", "point_2", "point_3"])

        for i, step in enumerate(structure):
            if step == "establish_elements":
                points.append(
                    "The elements required by law are satisfied in this case."
                )
            elif step == "apply_facts":
                if facts:
                    points.append(f"The facts establish that {facts[0].lower()}")
                else:
                    points.append("The factual record supports this position.")
            elif step == "distinguish_precedent":
                points.append("This case is distinguishable from contrary authority.")
            elif step == "identify_rule":
                points.append("The applicable procedural rule is clear and mandatory.")
            elif step == "show_violation":
                points.append(
                    "There has been a clear violation of the applicable standard."
                )
            elif step == "identify_policy":
                points.append(
                    "The relevant policy objectives support this interpretation."
                )
            elif step == "establish_hardship":
                points.append("The hardship to our client is significant and concrete.")
            else:
                points.append(f"Supporting point {i + 1} for the thesis.")

        return points

    def _anticipate_rebuttals(
        self,
        arg_type: ArgumentType,
        thesis: str,
    ) -> List[str]:
        """Anticipate potential counter-arguments."""
        rebuttals = []

        if arg_type == ArgumentType.SUBSTANTIVE:
            rebuttals.append("Opposing party may argue the elements are not satisfied.")
            rebuttals.append("Contrary precedent may be cited.")

        elif arg_type == ArgumentType.PROCEDURAL:
            rebuttals.append("Opposing party may argue waiver or forfeiture.")
            rebuttals.append("Harmless error doctrine may be invoked.")

        elif arg_type == ArgumentType.POLICY:
            rebuttals.append("Counter-policy arguments may be raised.")
            rebuttals.append("Slippery slope concerns may be presented.")

        elif arg_type == ArgumentType.EQUITABLE:
            rebuttals.append("Clean hands doctrine may be invoked.")
            rebuttals.append("Legal remedy adequacy may be argued.")

        else:
            rebuttals.append("General opposition to the thesis is expected.")

        return rebuttals

    def _score_persuasiveness(self, argument: LegalArgument) -> float:
        """Score argument persuasiveness."""
        score = 0.5

        # IRAC completeness
        if argument.irac:
            score += 0.2

        # Evidence support
        score += min(len(argument.evidence) * 0.05, 0.15)

        # Citation support
        score += min(len(argument.citations) * 0.05, 0.15)

        return min(score, 1.0)

    def _score_novelty(self, argument: LegalArgument) -> float:
        """Score argument novelty."""
        # Would compare against precedent database
        # For now, base on argument type
        novelty_map = {
            ArgumentType.POLICY: 0.7,
            ArgumentType.EQUITABLE: 0.6,
            ArgumentType.CONSTITUTIONAL: 0.5,
            ArgumentType.SUBSTANTIVE: 0.3,
            ArgumentType.PROCEDURAL: 0.2,
        }
        return novelty_map.get(argument.argument_type, 0.5)

    def generate_rebuttal(
        self,
        original_argument: LegalArgument,
        rebuttal_points: List[str] = None,
    ) -> LegalArgument:
        """Generate a rebuttal to an argument."""
        arg_id = f"rebuttal_{original_argument.id}"

        thesis = (
            f"The opposing argument regarding {original_argument.title.lower()} fails"
        )

        rebuttal = LegalArgument(
            id=arg_id,
            title=f"Rebuttal: {original_argument.title}",
            argument_type=original_argument.argument_type,
            strength=ArgumentStrength.MODERATE,
            thesis=thesis,
            supporting_points=rebuttal_points
            or original_argument.anticipated_rebuttals,
            rhetorical_strategy=RhetoricalStrategy.LOGOS,
        )

        rebuttal.persuasiveness = 0.6
        rebuttal.relevance = 0.9
        rebuttal.novelty = 0.4

        self.generated_arguments[arg_id] = rebuttal
        return rebuttal

    def optimize_argument_order(
        self,
        case_id: str,
        side: str,
    ) -> List[LegalArgument]:
        """Optimize the order of arguments for maximum persuasiveness."""
        key = f"{case_id}_{side}"
        if key not in self.argument_sets:
            return []

        arg_set = self.argument_sets[key]
        all_args = arg_set.main_arguments + arg_set.alternative_arguments

        # Sort by: strongest first, then by type
        # (Some prefer strongest in middle - "primacy-recency")
        sorted_args = sorted(all_args, key=lambda a: a.compute_score(), reverse=True)

        # Apply primacy-recency: best at start and end
        if len(sorted_args) >= 3:
            reordered = []
            for i, arg in enumerate(sorted_args):
                if i % 2 == 0:
                    reordered.insert(0, arg)
                else:
                    reordered.append(arg)
            return reordered

        return sorted_args

    def generate_brief_outline(
        self,
        case_id: str,
        side: str,
    ) -> Dict[str, Any]:
        """Generate an outline for a legal brief."""
        key = f"{case_id}_{side}"
        if key not in self.argument_sets:
            return {"error": "Argument set not found"}

        arg_set = self.argument_sets[key]
        optimized = self.optimize_argument_order(case_id, side)

        outline = {
            "case_id": case_id,
            "side": side,
            "sections": [
                {
                    "name": "Statement of Issues",
                    "content": [a.irac.issue for a in optimized if a.irac],
                },
                {
                    "name": "Statement of Facts",
                    "content": list(
                        set(fact for a in optimized for fact in a.evidence)
                    ),
                },
                {
                    "name": "Argument",
                    "subsections": [
                        {
                            "title": a.title,
                            "thesis": a.thesis,
                            "points": a.supporting_points,
                            "citations": a.citations,
                        }
                        for a in optimized
                    ],
                },
                {
                    "name": "Conclusion",
                    "content": [a.irac.conclusion for a in optimized if a.irac],
                },
            ],
            "quality_score": arg_set.compute_omega(),
        }

        return outline

    def get_stats(self) -> Dict[str, Any]:
        """Get generator statistics."""
        return {
            "argument_sets": len(self.argument_sets),
            "total_arguments": len(self.generated_arguments),
            "argument_types": list(
                set(a.argument_type.value for a in self.generated_arguments.values())
            ),
        }

    def to_dict(self) -> Dict[str, Any]:
        """Serialize generator state."""
        return {
            "argument_sets": {k: v.to_dict() for k, v in self.argument_sets.items()},
            "stats": self.get_stats(),
        }


# Factory functions
def create_argument_generator() -> ArgumentGenerator:
    """Create a new argument generator."""
    return ArgumentGenerator()