from memory.session_memory import SessionMemory

session = SessionMemory()

# Detect portfolio from source content
portfolio = session.detect_portfolio(source_content)
# Returns: "workers_comp" with confidence 0.85

# Record metrics as they're written (prevents contradictions)
session.record_metric("model_performance", "R² = 0.72")
session.record_metric("sample_size", "450,000 policies")

# Track compliance issues for revision
session.record_compliance_issue(
    description="Tail development factors not documented",
    severity="HIGH",
    section="Loss Development Analysis"
)

# Writer agent queries: "What R² did I already write?"
metrics = session.get_recorded_metrics()
# Returns: {"model_performance": "R² = 0.72", "sample_size": "450,000 policies"}

from memory.cross_session_memory import CrossSessionMemory

cross_session = CrossSessionMemory(db_path="data/memory/cross_session.db")

# After each generation, record outcomes
cross_session.record_generation(
    session_id="session_abc123",
    portfolio="workers_comp",
    success=True,
    quality_score=8.2,
    iterations=2,
    compliance_issues=["Fee schedule documentation incomplete"],
    sections_revised=["Medical Cost Trends", "Loss Development Analysis"]
)

# Before next Workers' Comp generation, query historical patterns
common_issues = cross_session.get_common_issues("workers_comp", top_n=5)
# Returns:
# [
#     {"issue": "Tail development inadequate", "frequency": 0.73},
#     {"issue": "Fee schedule not addressed", "frequency": 0.65},
#     {"issue": "NCCI mapping incomplete", "frequency": 0.42},
#     {"issue": "Medical trend assumption missing", "frequency": 0.38},
#     {"issue": "Experience mod calculation unclear", "frequency": 0.31}
# ]

# Get quality baseline for comparison
baseline = cross_session.get_quality_baseline("workers_comp")
# Returns: {"avg_score": 7.8, "avg_iterations": 2.3, "total_generations": 12}

# Query successful patterns
patterns = cross_session.get_successful_patterns("workers_comp")
# Returns sections and approaches that consistently score 8.0+

from memory.user_memory import UserMemory

user = UserMemory(user_id="analyst_sarah")

# Set preferences
user.set_preference("tone", "technical")  # vs "executive" or "regulatory"
user.set_preference("detail_level", "high")
user.set_preference("include_code_snippets", True)
user.set_preference("max_section_length", 1500)  # words

# Add custom compliance rules
user.add_custom_rule("Always verify CECL alignment for loss reserving models")
user.add_custom_rule("Include state-specific fee schedule references")

# Exclude sections user doesn't need
user.set_preference("skip_sections", ["Business Context"])

# Writer agent retrieves preferences
instructions = user.get_writing_instructions()
# Returns formatted string:
# "TONE: technical
#  DETAIL: high
#  Include code snippets where relevant
#  Maximum section length: 1500 words
#  CUSTOM RULES:
#    - Always verify CECL alignment for loss reserving models
#    - Include state-specific fee schedule references"

from memory.memory_manager import MemoryManager

# Initialize with user context
memory = MemoryManager(user_id="analyst_sarah", db_path="data/memory/cross_session.db")

# Start session (detects portfolio, loads historical patterns)
session_id = memory.start_session(source_content)
print(f"Portfolio: {memory.session.detected_portfolio}")
print(f"Historical issues to watch: {memory.get_proactive_warnings()}")

# Get context for each agent phase
research_context = memory.get_research_context("Methodology")
# Returns: RAG focus areas + successful patterns from past docs

writing_context = memory.get_writing_context("Validation")
# Returns: Metrics already written + user tone preferences + section patterns

compliance_context = memory.get_compliance_context()
# Returns: Portfolio checkpoints + custom rules + common historical issues

editorial_context = memory.get_editorial_context()
# Returns: Quality baseline + user preferences + issues already flagged

# End session (persists learnings)
memory.end_session(final_score=8.2, success=True)
# Writes to cross_session.db for future generations

# From agents/graph_nodes.py

def research_phase(state: Dict[str, Any]) -> Dict[str, Any]:
    memory = state.get("memory_manager")
    
    for section in sections:
        # Memory provides historical context
        context = memory.get_research_context(section)
        
        # Includes:
        # - "Past workers_comp docs prioritized NCCI classification"
        # - "Successful Methodology sections averaged 1,200 words"
        # - "User prefers technical detail with code snippets"
        
        findings = research_agent.research_topic(
            topic=f"{section} {model_type}",
            additional_context=context  # Memory-enhanced
        )

# From agents/graph_state.py

from typing import TypedDict, List, Dict, Optional

class DocumentState(TypedDict, total=False):
    # === Input ===
    document_title: str
    source_content: str
    user_id: str
    
    # === Portfolio Detection ===
    detected_portfolio: str      # personal_auto, homeowners, workers_comp, commercial_auto
    detected_model_type: str     # frequency, severity
    portfolio_confidence: float  # 0.0 to 1.0
    detection_keywords: List[str]
    
    # === Dynamic Configuration ===
    required_sections: List[str]
    quality_threshold: float     # 7.0 for personal_auto, 7.5 for others
    max_iterations: int          # 3 standard, 4 for workers_comp
    compliance_strictness: str   # standard, elevated, strict
    custom_instructions: str     # Built from portfolio config
    
    # === Generation State ===
    current_document: str
    sections_written: List[Dict]
    research_contexts: Dict[str, str]
    
    # === Quality Control ===
    compliance_passed: bool
    editorial_passed: bool
    quality_score: float
    compliance_issues: List[Dict]
    editorial_issues: List[Dict]
    
    # === Iteration Control ===
    current_iteration: int
    revision_history: List[Dict]
    
    # === Output ===
    final_document: Optional[str]
    generation_successful: bool

# From agents/graph_nodes.py

def detect_portfolio(state: Dict[str, Any]) -> Dict[str, Any]:
    """Analyze source content to determine insurance portfolio."""
    
    source_content = state.get("source_content", "").lower()
    
    PORTFOLIO_KEYWORDS = {
        "personal_auto": [
            "driver age", "vehicle type", "collision", "territory",
            "prior claims", "comprehensive", "liability", "pip"
        ],
        "homeowners": [
            "dwelling", "protection class", "cat ", "hurricane",
            "roof", "coverage a", "demand surge", "coastal"
        ],
        "workers_comp": [
            "ncci", "payroll", "medical cost", "indemnity",
            "classification code", "experience mod", "fee schedule",
            "tail development", "monopolistic state"
        ],
        "commercial_auto": [
            "fleet", "gvw", "social inflation", "nuclear verdict",
            "driver turnover", "trucking", "mcs-90", "jurisdiction"
        ]
    }
    
    # Count keyword matches
    match_counts = {}
    for portfolio, keywords in PORTFOLIO_KEYWORDS.items():
        matches = [kw for kw in keywords if kw in source_content]
        match_counts[portfolio] = len(matches)
    
    best_portfolio = max(match_counts, key=match_counts.get)
    confidence = min(1.0, match_counts[best_portfolio] / 10.0)
    
    # Require minimum confidence
    if match_counts[best_portfolio] < 3:
        best_portfolio = "personal_auto"
        confidence = 0.3
    
    return {
        "detected_portfolio": best_portfolio,
        "portfolio_confidence": confidence,
        "detection_keywords": [kw for kw in PORTFOLIO_KEYWORDS[best_portfolio] 
                              if kw in source_content][:10]
    }


def configure_for_portfolio(state: Dict[str, Any]) -> Dict[str, Any]:
    """Load portfolio-specific configuration."""
    
    portfolio = state.get("detected_portfolio")
    config = PORTFOLIO_REGISTRY.get_config(portfolio)
    
    # Build custom instructions
    instructions = [
        f"PORTFOLIO: {config.display_name}",
        f"EXPOSURE BASIS: {config.exposure_basis}",
        "",
        "FOCUS AREAS:",
        *[f"  - {area}" for area in config.research_focus_areas],
        "",
        "COMMON PITFALLS TO AVOID:",
        *[f"  - {pitfall}" for pitfall in config.common_pitfalls],
        "",
        "COMPLIANCE CHECKPOINTS:",
        *[f"  - {check}" for check in config.compliance_checkpoints]
    ]
    
    return {
        "required_sections": config.required_sections + config.portfolio_specific_sections,
        "quality_threshold": config.min_quality_score,
        "max_iterations": config.max_iterations,
        "compliance_strictness": config.compliance_strictness,
        "custom_instructions": "\n".join(instructions)
    }


def revision_phase(state: Dict[str, Any]) -> Dict[str, Any]:
    """Revise sections based on compliance and editorial feedback."""
    
    compliance_issues = state.get("compliance_issues", [])
    editorial_issues = state.get("editorial_issues", [])
    
    # Build revision instructions from feedback
    revision_instructions = []
    
    for issue in compliance_issues:
        if issue.get("severity") in ["CRITICAL", "HIGH"]:
            revision_instructions.append(f"COMPLIANCE: {issue['description']}")
    
    for issue in editorial_issues:
        if issue.get("priority") in ["CRITICAL", "HIGH"]:
            revision_instructions.append(f"EDITORIAL: {issue['description']}")
    
    # Revise each section with WriterAgent
    writer = WriterAgent()
    revised_sections = []
    
    for section in state.get("sections_written", []):
        revised = writer.revise_section(
            original_content=section,
            revision_instructions="\n".join(revision_instructions),
            source_content=state.get("source_content")
        )
        revised_sections.append(revised)
    
    return {
        "sections_written": revised_sections,
        "current_iteration": state.get("current_iteration", 0) + 1,
        "revision_history": state.get("revision_history", []) + [{
            "iteration": state.get("current_iteration"),
            "issues_addressed": len(revision_instructions)
        }]
    }

# From agents/graph_nodes.py

def route_after_compliance(state: Dict[str, Any]) -> str:
    """Decide next step after compliance check."""
    
    compliance_passed = state.get("compliance_passed", False)
    current_iteration = state.get("current_iteration", 0)
    max_iterations = state.get("max_iterations", 3)
    
    if compliance_passed:
        return "editorial"
    elif current_iteration < max_iterations:
        return "revision"
    else:
        return "complete"  # Accept what we have


def route_after_editorial(state: Dict[str, Any]) -> str:
    """Decide next step after editorial review."""
    
    editorial_passed = state.get("editorial_passed", False)
    current_iteration = state.get("current_iteration", 0)
    max_iterations = state.get("max_iterations", 3)
    
    if editorial_passed:
        return "complete"
    elif current_iteration < max_iterations:
        return "revision"
    else:
        return "complete"


def route_after_revision(state: Dict[str, Any]) -> str:
    """After revision, always loop back to compliance."""
    return "compliance"

# From agents/langgraph_orchestrator.py

from langgraph.graph import StateGraph, END

def create_workflow():
    """Assemble the documentation generation graph."""
    
    workflow = StateGraph(DocumentState)
    
    # Add nodes
    workflow.add_node("detect_portfolio", detect_portfolio)
    workflow.add_node("configure", configure_for_portfolio)
    workflow.add_node("research", research_phase)
    workflow.add_node("write", writing_phase)
    workflow.add_node("compliance", compliance_phase)
    workflow.add_node("editorial", editorial_phase)
    workflow.add_node("revision", revision_phase)
    workflow.add_node("complete", complete_workflow)
    
    # Linear edges
    workflow.add_edge("detect_portfolio", "configure")
    workflow.add_edge("configure", "research")
    workflow.add_edge("research", "write")
    workflow.add_edge("write", "compliance")
    
    # Conditional edges (this is where cycles are defined)
    workflow.add_conditional_edges(
        "compliance",
        route_after_compliance,
        {
            "editorial": "editorial",
            "revision": "revision",
            "complete": "complete"
        }
    )
    
    workflow.add_conditional_edges(
        "editorial",
        route_after_editorial,
        {
            "complete": "complete",
            "revision": "revision"
        }
    )
    
    workflow.add_conditional_edges(
        "revision",
        route_after_revision,
        {
            "compliance": "compliance"  # Loop back
        }
    )
    
    # Entry and exit
    workflow.set_entry_point("detect_portfolio")
    workflow.add_edge("complete", END)
    
    return workflow.compile()

from agents.langgraph_orchestrator import LangGraphOrchestrator

# Initialize
orchestrator = LangGraphOrchestrator(user_id="analyst_sarah")

# Generate documentation
document, state = orchestrator.generate_documentation(
    document_title="2024 WC Frequency Model",
    document_type="model_doc",
    source_content=ppt_content
)

# Inspect results
print(f"Portfolio: {state['detected_portfolio']}")           # workers_comp
print(f"Confidence: {state['portfolio_confidence']:.0%}")    # 85%
print(f"Quality Score: {state['quality_score']:.1f}/10")     # 8.2/10
print(f"Iterations: {state['current_iteration']}")           # 2
print(f"Sections: {len(state['required_sections'])}")        # 10
print(f"Success: {state['generation_successful']}")          # True

orchestrator = LangGraphOrchestrator(enable_checkpointing=True)  # Default: ON

document, state = orchestrator.generate_documentation(
    document_title="2024 WC Frequency Model",
    source_content=ppt_content
)

# Get execution history
history = orchestrator.get_execution_history()
for entry in history:
    print(f"Step {entry.step}: {entry.node}")
    print(f"  Iteration: {entry.iteration}")
    print(f"  Compliance Passed: {entry.compliance_passed}")

audit_log = orchestrator.get_audit_log()

print(f"Thread ID: {audit_log.thread_id}")
print(f"Portfolio: {audit_log.detected_portfolio}")
print(f"Execution Path: {audit_log.get_path_summary()}")
print(f"Total Iterations: {audit_log.total_iterations}")
print(f"Final Quality: {audit_log.final_quality_score}/10")
print(f"Success: {audit_log.generation_successful}")

# Get decision points (where routing happened)
for decision in audit_log.get_decision_points():
    print(f"{decision['node']}: {'PASSED' if decision['passed'] else 'FAILED'} -> {decision['routed_to']}")

summary = orchestrator.get_execution_summary()

# Ready for Streamlit
st.subheader("Execution Audit Trail")
st.write(f"**Portfolio Detected:** {summary['portfolio_detected']}")
st.write(f"**Execution Path:** {summary['execution_path']}")
st.write(f"**Quality Score:** {summary['final_quality_score']}")

st.subheader("Decision Points")
for decision in summary['decision_points']:
    status = "✅ PASSED" if decision['passed'] else "❌ FAILED"
    st.write(f"**{decision['node']}**: {status} -> routed to {decision['routed_to']}")

orchestrator.print_execution_trace()

PORTFOLIO_KEYWORDS = {
    "personal_auto": [
        "driver age", "vehicle type", "collision", "bodily injury", "territory",
        "prior claims", "comprehensive", "liability", "pip", "uninsured motorist",
        "deductible", "annual mileage", "good driver", "multi-car", "garaging"
    ],
    "homeowners": [
        "dwelling", "protection class", "cat ", "hurricane", "hail", "roof",
        "construction type", "coverage a", "coverage b", "coverage c",
        "replacement cost", "wind", "tornado", "wildfire", "flood zone",
        "coastal", "air ", "rms", "corelogic", "demand surge"
    ],
    "workers_comp": [
        "injury", "ncci", "payroll", "medical cost", "indemnity", "classification code",
        "experience mod", "body part", "nature of injury", "cause of injury",
        "temporary disability", "permanent disability", "medical only", "lost time",
        "fee schedule", "monopolistic state", "tail development"
    ],
    "commercial_auto": [
        "fleet", "vehicle class", "radius", "gvw", "for-hire", "driver turnover",
        "commercial vehicle", "trucking", "light truck", "medium truck", "heavy truck",
        "tractor", "trailer", "hired auto", "non-owned auto", "mcs-90",
        "social inflation", "nuclear verdict", "jurisdiction"
    ]
}