┌─────────────────────────────────────────────────────────────────┐
│ User Interface                                                  │
│ ┌─────────────┐ ┌─────────────┐ ┌─────────────────────┐       │
│ │   Notion    │ │  BrainTrust │ │  FastAPI Swagger    │       │
│ │  Database   │ │ Chat Agent  │ │     (Debug)         │       │
│ └──────┬──────┘ └──────┬──────┘ └──────────┬──────────┘       │
└─────────┼──────────────────┼─────────────────────┼─────────────┘
          │                  │                     │
          ▼                  ▼                     ▼
┌─────────────────────────────────────────────────────────────────┐
│ FastAPI Backend                                                 │
│ ┌─────────────┐ ┌─────────────┐ ┌─────────────────────┐       │
│ │  /actions   │ │  /insights  │ │   /enrichment       │       │
│ │  endpoints  │ │  endpoints  │ │    endpoints        │       │
│ └─────────────┘ └─────────────┘ └─────────────────────┘       │
└─────────────────────────────┬───────────────────────────────────┘
                              │
          ┌───────────────────┼───────────────────┐
          ▼                   ▼                   ▼
┌─────────────────┐ ┌─────────────────┐ ┌─────────────────┐
│   Notion API    │ │   Firecrawl     │ │   Claude API    │
│ (Data Layer)    │ │  (Web Intel)    │ │   (Reasoning)   │
└─────────────────┘ └─────────────────┘ └─────────────────┘

from firecrawl import Firecrawl
from pydantic import BaseModel, Field
from typing import Optional, List

# Pydantic schemas for structured extraction
class CompanyInfo(BaseModel):
    """Schema for company information extraction."""
    employee_count: Optional[str] = Field(description="Employee count or range (e.g., '1000-5000')")
    industry: Optional[str] = Field(description="Primary business category")
    company_size: Optional[str] = Field(description="Size category: startup/SMB/enterprise")
    headquarters: Optional[str] = Field(description="City, State/Country")

class CompanyNews(BaseModel):
    """Schema for company news extraction."""
    has_layoffs: bool = Field(default=False, description="Recent layoff announcements")
    has_ma: bool = Field(default=False, description="M&A activity (acquiring or being acquired)")
    is_expanding: bool = Field(default=False, description="Growth/expansion signals")
    sentiment: Optional[str] = Field(description="Overall news sentiment: positive/neutral/negative")
    key_events: List[str] = Field(default_factory=list, description="Recent significant events")
    sources: List[str] = Field(default_factory=list, description="URLs of source articles")

def enrich_company(firecrawl: Firecrawl, company_name: str) -> dict:
    """
    Enrich a single company using Firecrawl agent.
    Returns structured data via Pydantic schema extraction.
    """
    # Get company info
    info_result = firecrawl.agent(
        prompt=f"Find current information about {company_name}: employee count, industry, company size, headquarters location.",
        schema=CompanyInfo,
        max_pages=3
    )
    
    # Get recent news and events
    news_result = firecrawl.agent(
        prompt=f"Find recent news about {company_name} from the last 6 months. Look for: layoffs, acquisitions, mergers, expansion, funding, leadership changes.",
        schema=CompanyNews,
        max_pages=5
    )
    
    return {
        "company_name": company_name,
        "employee_count": info_result.employee_count,
        "industry": info_result.industry,
        "company_size": info_result.company_size,
        "headquarters": info_result.headquarters,
        "has_layoffs": news_result.has_layoffs,
        "has_ma": news_result.has_ma,
        "is_expanding": news_result.is_expanding,
        "sentiment": news_result.sentiment,
        "key_events": news_result.key_events,
        "sources": news_result.sources,
    }

class RiskScorer:
    """
    Evidence-based risk scoring for customer churn prediction.
    Each signal must be backed by a specific event, not just sentiment.
    """
    
    # Critical signals - immediate action required
    CRITICAL_SIGNALS = {
        "budget_cuts": 3,              # Explicit cost reduction programs
        "spending_freeze": 3,          # Procurement freezes announced
        "acquired_by_competitor": 3,   # Being bought by someone who has competing product
        "bankruptcy_filing": 3,        # Chapter 11 or equivalent
    }
    
    # High-risk signals - schedule outreach this week
    HIGH_RISK_SIGNALS = {
        "major_layoffs": 2,            # >10% workforce reduction
        "it_department_cuts": 2,       # Specifically IT/tech layoffs
        "cto_cio_departure": 2,        # Key technical stakeholder leaving
        "negative_earnings": 2,        # Missed earnings + negative guidance
        "contract_consolidation": 2,   # Announced vendor consolidation initiative
    }
    
    # Watch signals - monitor closely
    WATCH_SIGNALS = {
        "minor_layoffs": 1,            # <10% workforce reduction
        "leadership_transition": 1,    # C-suite changes (not tech-specific)
        "restructuring": 1,            # Org changes without clear direction
        "mixed_earnings": 1,           # Beat/miss with uncertain outlook
    }
    
    # Positive signals - reduce risk score
    POSITIVE_SIGNALS = {
        "acquisition_mode": -2,        # Company acquiring others (growth)
        "expansion_announced": -2,     # New markets, new offices
        "funding_round": -2,           # Fresh capital injection
        "record_revenue": -1,          # Strong financial performance
        "hiring_spree": -1,            # Significant hiring (especially IT)
    }
    
    def calculate_risk(self, events: list[CompanyEvent]) -> RiskAssessment:
        """
        Calculate risk score from list of detected events.
        Returns risk level + supporting evidence.
        """
        score = 0
        evidence = []
        
        for event in events:
            signal_weight = self._get_signal_weight(event.event_type)
            if signal_weight != 0:
                score += signal_weight
                evidence.append({
                    "signal": event.event_type,
                    "weight": signal_weight,
                    "date": event.event_date,
                    "summary": event.event_summary,
                    "source": event.source_url
                })
        
        risk_level = self._score_to_level(score)
        
        return RiskAssessment(
            level=risk_level,
            score=score,
            evidence=evidence,
            recommended_action=self._get_recommended_action(risk_level)
        )
    
    def _score_to_level(self, score: int) -> str:
        if score >= 5:
            return "CRITICAL"
        elif score >= 3:
            return "HIGH"
        elif score >= 1:
            return "WATCH"
        elif score <= -2:
            return "OPPORTUNITY"
        else:
            return "STABLE"

# main.py
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel

app = FastAPI(
    title="Sales Intelligence API",
    description="Portfolio intelligence and action management",
    version="1.0.0"
)

# ============ Actions Endpoints ============

@app.get("/actions/pending", response_model=list[ActionItem])
async def get_pending_actions(
    risk_level: str | None = None,
    limit: int = 20
):
    """
    Get pending action items, optionally filtered by risk level.
    Returns actions sorted by priority (CRITICAL first).
    """
    actions = await notion_service.get_actions(
        status="pending",
        risk_level=risk_level,
        limit=limit
    )
    return sorted(actions, key=lambda a: RISK_PRIORITY[a.risk_level])

@app.post("/actions/{action_id}/complete")
async def complete_action(
    action_id: str,
    outcome: ActionOutcome
):
    """
    Mark an action as complete with outcome tracking.
    """
    await notion_service.update_action(
        action_id=action_id,
        status="completed",
        outcome=outcome.result,
        notes=outcome.notes,
        completed_at=datetime.now()
    )
    return {"status": "completed", "action_id": action_id}

# ============ Insights Endpoints ============

@app.get("/insights/portfolio-summary", response_model=PortfolioSummary)
async def get_portfolio_summary():
    """
    Executive overview of portfolio health.
    """
    companies = await notion_service.get_all_companies()
    
    return PortfolioSummary(
        total_companies=len(companies),
        total_arr=sum(c.arr for c in companies if c.arr),
        risk_distribution={
            level: len([c for c in companies if c.risk_level == level])
            for level in ['CRITICAL', 'HIGH', 'WATCH', 'STABLE', 'OPPORTUNITY']
        },
        enrichment_coverage=len([c for c in companies if c.last_enriched]) / len(companies),
        pending_actions=await notion_service.count_pending_actions()
    )

@app.get("/insights/churn-risk", response_model=list[ChurnRiskAccount])
async def get_churn_risk_accounts(
    quarter: str | None = None,
    min_arr: float | None = None
):
    """
    Get accounts with elevated churn risk, filtered by renewal quarter and ARR.
    """
    companies = await notion_service.get_companies(
        risk_levels=['CRITICAL', 'HIGH'],
        renewal_quarter=quarter,
        min_arr=min_arr
    )
    
    return [
        ChurnRiskAccount(
            company_name=c.name,
            arr=c.arr,
            risk_level=c.risk_level,
            risk_evidence=c.risk_evidence,
            renewal_date=c.renewal_date,
            days_to_renewal=(c.renewal_date - date.today()).days if c.renewal_date else None,
            recommended_action=get_recommended_action(c.risk_level)
        )
        for c in companies
    ]

@app.get("/insights/expansion-opportunities", response_model=list[ExpansionOpportunity])
async def get_expansion_opportunities(
    min_employees: int | None = None,
    industry: str | None = None
):
    """
    Get accounts showing growth signals suitable for upsell.
    """
    companies = await notion_service.get_companies(
        risk_levels=['OPPORTUNITY'],
        min_employees=min_employees,
        industry=industry
    )
    
    return [
        ExpansionOpportunity(
            company_name=c.name,
            current_arr=c.arr,
            employee_count=c.employee_count,
            growth_signals=[e for e in c.events if e.event_type in GROWTH_EVENTS],
            expansion_potential=estimate_expansion_potential(c)
        )
        for c in companies
    ]

# ============ Company Endpoints ============

@app.get("/companies/{company_id}", response_model=CompanyDetail)
async def get_company_detail(company_id: str):
    """
    Get full detail for a single company including enrichment data.
    """
    company = await notion_service.get_company(company_id)
    if not company:
        raise HTTPException(status_code=404, detail="Company not found")
    return company

@app.post("/companies/{company_id}/enrich")
async def trigger_enrichment(company_id: str):
    """
    Manually trigger enrichment for a specific company.
    """
    company = await notion_service.get_company(company_id)
    if not company:
        raise HTTPException(status_code=404, detail="Company not found")
    
    # Invoke worker Lambda directly
    await trigger_worker(company_id, company.name, company.domain)
    
    return {"status": "enrichment_triggered", "company": company.name}