"""Abbe MCP Server — exposes Abbe Core API as tools for AI coding assistants."""

from __future__ import annotations

from typing import Optional

from mcp.server.fastmcp import FastMCP

from abbe_mcp import client

mcp = FastMCP(
    "Abbe",
    instructions="AI workforce impact analysis — resolve job titles, score occupations, analyze portfolios",
)


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


def _confidence_label(conf: float) -> str:
    if conf >= 0.55:
        return "Very strong match"
    if conf >= 0.45:
        return "Strong match"
    if conf >= 0.35:
        return "Moderate match"
    return "Weak match"


def _format_pct(val: float) -> str:
    sign = "+" if val > 0 else ""
    return f"{sign}{val:.1f}%"


# ---------------------------------------------------------------------------
# Tools
# ---------------------------------------------------------------------------


@mcp.tool()
async def resolve_job_title(
    title: str,
    description: Optional[str] = None,
    top_k: int = 3,
) -> str:
    """Resolve a job title (and optional description) to a standardized O*NET occupation.

    Use this first when you have a job title and need to find its occupation code
    for use with other tools like get_ai_impact or get_skills.

    Args:
        title: Job title to resolve, e.g. 'Software Engineer', 'Head of Product'
        description: Short job description to improve match quality for ambiguous titles
        top_k: Number of candidate matches to return (default 3)
    """
    query: dict = {"title": title}
    if description:
        query["description"] = description

    data = await client.resolve(queries=[query], top_k=top_k)
    result = data["results"][0]
    matches = result["matches"]

    if not matches:
        return f'No matches found for "{title}". Try a different title or add a description.'

    lines = [f'Resolved "{title}" to:']
    for m in matches:
        label = _confidence_label(m["confidence"])
        star = " *" if m["rank"] == 1 else ""
        lines.append(
            f"  {m['rank']}. {m['title']} ({m['soc_code']}) "
            f"— confidence: {m['confidence']:.3f}{star} {label}"
        )

    best = matches[0]
    lines.append("")
    lines.append(
        f'Use soc_code "{best["soc_code"]}" with get_ai_impact or get_skills '
        f"for further analysis."
    )
    return "\n".join(lines)


@mcp.tool()
async def get_ai_impact(
    soc_codes: list[str],
    scenario: str = "baseline",
    horizon: str = "3y",
    use_case: Optional[str] = None,
) -> str:
    """Get AI impact scores for one or more occupations.

    Shows how AI is projected to affect demand for each occupation.
    Accepts O*NET codes (e.g. '15-1252.00') or SOC-6 codes (e.g. '15-1252').

    Args:
        soc_codes: O*NET or SOC occupation codes
        scenario: AI adoption scenario — 'baseline', 'aggressive', or 'conservative'
        horizon: Time horizon — '1y', '3y', or '5y'
        use_case: Capability filter — 'genai' for generative AI only (default: all capabilities)
    """
    data = await client.scores(
        soc_codes=soc_codes,
        scenarios=[scenario],
        horizons=[horizon],
        include=["confidence_interval"],
        use_case=use_case,
    )

    items = data["data"]
    if not items:
        return "No results returned. Check that the SOC codes are valid."

    uc_label = f", use_case={use_case}" if use_case else ""
    lines = [f"AI Impact Scores ({scenario} scenario, {horizon} horizon{uc_label}):"]
    lines.append("")

    for item in items:
        code = item["soc_code"]
        title = item.get("title") or "Unknown"
        coverage = item["coverage"]

        if coverage != "scored" or not item.get("scores"):
            lines.append(f"  {title} ({code}) — not scored ({coverage})")
            continue

        s = item["scores"]
        impact = _format_pct(s["impact_pct"])
        exposure = s["exposure"]
        tilt = s["tilt"]

        lines.append(f"  {title} ({code})")
        lines.append(f"    Impact: {impact}  |  Exposure: {exposure:.2f}  |  Tilt: {tilt}")

        ci = item.get("confidence_interval")
        if ci:
            lines.append(
                f"    90% CI: [{_format_pct(ci['impact_pct_p10'])} to {_format_pct(ci['impact_pct_p90'])}]"
            )
        lines.append("")

    lines.append(
        "Use get_skills with these SOC codes to see which skills drive each occupation."
    )
    return "\n".join(lines)


@mcp.tool()
async def analyze_workforce(
    workforce: list[dict],
    scenario: str = "baseline",
    horizon: str = "3y",
) -> str:
    """Analyze AI impact across a workforce portfolio.

    Provide a list of roles with headcount and optional salary to get aggregate
    workforce-level impact projections.

    Args:
        workforce: List of roles, each with 'soc_code' (str), 'headcount' (int), and optional 'salary' (float)
        scenario: AI adoption scenario — 'baseline', 'aggressive', or 'conservative'
        horizon: Time horizon — '1y', '3y', or '5y'
    """
    # Normalize: accept 'salary' as shorthand for 'annual_wage'
    api_workforce = []
    for role in workforce:
        row: dict = {
            "soc_code": role["soc_code"],
            "headcount": role["headcount"],
        }
        wage = role.get("salary") or role.get("annual_wage")
        if wage is not None:
            row["annual_wage"] = wage
        api_workforce.append(row)

    include = ["top_displaced", "top_augmented"]
    # Include wage-related enrichments if any role has a salary
    has_wages = any(r.get("annual_wage") for r in api_workforce)
    if not has_wages:
        include.append("by_group")

    data = await client.portfolio(
        workforce=api_workforce,
        scenario=scenario,
        horizon=horizon,
        include=include,
    )

    summary = data["data"]["summary"]
    lines = [f"Workforce Portfolio Analysis ({scenario}, {horizon}):"]
    lines.append("")
    lines.append(f"  Total headcount: {summary['total_headcount']:,}")
    lines.append(f"  Adjusted headcount: {summary['adjusted_headcount']:,.0f}")
    lines.append(f"  Headcount change: {_format_pct(summary['headcount_change_pct'])}")

    if summary.get("total_wage_bill"):
        lines.append(f"  Total wage bill: ${summary['total_wage_bill']:,.0f}")
    if summary.get("total_savings_usd"):
        lines.append(f"  Projected savings: ${summary['total_savings_usd']:,.0f}")
    if summary.get("savings_pct_of_wage_bill") is not None:
        lines.append(f"  Savings as % of wage bill: {summary['savings_pct_of_wage_bill']:.1f}%")

    lines.append(f"  Scored roles: {summary['scored_roles']} / {summary['scored_roles'] + summary['unscored_roles']}")
    lines.append("")

    displaced = data["data"].get("top_displaced")
    if displaced:
        lines.append("  Most displaced roles:")
        for d in displaced[:5]:
            lines.append(
                f"    - {d['title']} ({d['soc_code']}): "
                f"{_format_pct(d['impact_pct'])}, "
                f"{d['headcount_change']:+.0f} headcount"
            )
        lines.append("")

    augmented = data["data"].get("top_augmented")
    if augmented:
        lines.append("  Most augmented roles:")
        for a in augmented[:5]:
            lines.append(
                f"    - {a['title']} ({a['soc_code']}): {_format_pct(a['impact_pct'])}"
            )
        lines.append("")

    return "\n".join(lines)


@mcp.tool()
async def get_skills(
    soc_codes: list[str],
    top_n: int = 10,
    include_dimensions: bool = False,
) -> str:
    """Get the most important and distinctive skills for one or more occupations.

    Shows which skills define and differentiate each role. Use after resolve_job_title
    to understand what skills matter most for an occupation.

    Args:
        soc_codes: O*NET or SOC occupation codes
        top_n: Number of top skills to return per occupation (default 10)
        include_dimensions: Include extended metrics (posting count, salary premium, growth, etc.)
    """
    include = ["dimensions"] if include_dimensions else None
    data = await client.skills(soc_codes=soc_codes, top_n=top_n, include=include)

    skills_list = data["data"]
    meta = data["meta"]
    codes = meta.get("soc_codes", soc_codes)

    if not skills_list:
        return f"No skill data found for {', '.join(codes)}."

    lines = [f"Top {top_n} skills for {', '.join(codes)}:"]
    if meta.get("low_skill_coverage"):
        lines.append("  ⚠ Low skill coverage for one or more queried occupations")
    lines.append("")

    for i, sk in enumerate(skills_list, 1):
        distinctive = " [distinctive]" if sk.get("is_distinctive") else ""
        lines.append(
            f"  {i}. {sk['skill_subcategory']} ({sk['skill_category']})"
            f" — importance: {sk['importance']:.2f}{distinctive}"
        )
        # Show dimension fields when present
        if sk.get("posting_count") is not None:
            parts = []
            parts.append(f"posting_count: {sk['posting_count']:,}")
            if sk.get("demand_share") is not None:
                parts.append(f"demand_share: {sk['demand_share']:.1f}%")
            if sk.get("salary_premium") is not None:
                sal_label = f" ({sk['salary_label']})" if sk.get("salary_label") else ""
                parts.append(f"salary_premium: {sk['salary_premium']:+.1f}%{sal_label}")
            lines.append(f"     {' | '.join(parts)}")

            parts2 = []
            if sk.get("typical_experience_years") is not None:
                parts2.append(f"experience: {sk['typical_experience_years']:.1f} yrs")
            if sk.get("typical_education_level") is not None:
                parts2.append(f"education: {sk['typical_education_level']}")
            if parts2:
                lines.append(f"     {' | '.join(parts2)}")

            parts3 = []
            if sk.get("occupation_share") is not None:
                parts3.append(f"occupation_share: {sk['occupation_share']:.1f}%")
            if sk.get("demand_growth_pct") is not None:
                parts3.append(f"demand_growth: {sk['demand_growth_pct']:+.1f}%")
            if sk.get("ubiquity_pct") is not None:
                parts3.append(f"ubiquity: {sk['ubiquity_pct']:.1f}%")
            if sk.get("ecosystem_density_pct") is not None:
                parts3.append(f"ecosystem_density: {sk['ecosystem_density_pct']:.1f}%")
            if parts3:
                lines.append(f"     {' | '.join(parts3)}")

    lines.append("")
    if include_dimensions:
        lines.append(
            "Tip: Ask me to explain any of these dimensions if you want to understand what they measure."
        )
    lines.append(
        "Use get_ai_impact with these SOC codes to see projected employment changes."
    )
    return "\n".join(lines)


@mcp.tool()
async def explain_skill_dimensions(
    fields: Optional[list[str]] = None,
) -> str:
    """Explain what the skill dimension fields mean.

    Use this when a user asks "what does salary_premium mean?" or wants to
    understand any of the metrics returned by get_skills with include_dimensions=True.

    Args:
        fields: Specific field names to explain (e.g. ['salary_premium', 'ubiquity_pct']).
                If omitted, all 13 fields are explained.
    """
    descriptions: dict[str, str] = {
        "importance": (
            "How important this skill is to the occupation, balancing how commonly "
            "it appears in job postings and how distinctive it is to this role versus "
            "others. 0 to 1, higher means more important."
        ),
        "distinctiveness": (
            "Whether this skill is a specialization of this occupation relative to "
            "the broader labor market. -1 to +1. Positive means the occupation "
            "relies on this skill more than average; negative means less than average."
        ),
        "is_distinctive": (
            "True if this skill is a revealed specialization of the occupation."
        ),
        "posting_count": (
            "Number of job postings for this occupation that mentioned this skill "
            "in the last 18 months."
        ),
        "demand_share": (
            "Percentage of the occupation's job postings that mention this skill. "
            "A value of 74.3 means about 3 in 4 postings require it."
        ),
        "salary_premium": (
            "Percentage difference between the average salary for postings with "
            "this skill and the overall average salary for the same occupation. "
            "+9.1 means postings requiring this skill pay about 9% above the "
            "occupation's average. Calculated within the occupation, not across "
            "the labor market."
        ),
        "salary_label": (
            'Salary category based on the premium: "high premium" (>10% above '
            'average), "moderate premium" (3-10%), "average" (within 3%), or '
            '"below average" (>3% below).'
        ),
        "typical_experience_years": (
            "Average minimum years of experience requested in postings that "
            "mention this skill. Reflects the occupation's typical seniority "
            "level for this skill."
        ),
        "typical_education_level": (
            "Most common minimum education level requested in postings that "
            "mention this skill."
        ),
        "occupation_share": (
            "This occupation's share of total market demand for the skill, as a "
            "percentage. A value of 52.5 means this occupation accounts for over "
            "half of all postings requiring the skill."
        ),
        "demand_growth_pct": (
            "Percentage change in demand for this skill across the entire labor "
            "market, weighted toward recent trends. Positive means growing, "
            "negative means declining."
        ),
        "ubiquity_pct": (
            "Percentage of scored occupations that require this skill. Higher "
            "means more widely demanded across the labor market."
        ),
        "ecosystem_density_pct": (
            "How densely this skill connects to other specialized skills across "
            "occupations, expressed as a percentile (0-100). Higher means the "
            "skill sits in a tightly connected cluster of related capabilities. "
            "Lower means the skill is more structurally isolated."
        ),
    }

    # Determine which fields to explain
    if fields:
        selected = fields
    else:
        selected = list(descriptions.keys())

    lines = ["Skill Dimension Definitions:"]
    lines.append("")

    for name in selected:
        if name in descriptions:
            lines.append(f"  {name}")
            lines.append(f"    {descriptions[name]}")
            lines.append("")
        else:
            lines.append(f"  {name}")
            lines.append(
                f"    Unknown field. Available fields: {', '.join(descriptions.keys())}"
            )
            lines.append("")

    return "\n".join(lines).rstrip()


@mcp.tool()
async def list_occupations(
    search: Optional[str] = None,
    soc_prefix: Optional[str] = None,
) -> str:
    """Browse available occupations in the Abbe taxonomy.

    Use to discover what occupations are available or search by keyword.
    Returns up to 20 results. Use resolve_job_title for fuzzy matching.

    Args:
        search: Search term to filter occupations by title
        soc_prefix: Filter by SOC code prefix, e.g. '15' for computer occupations
    """
    data = await client.occupations(
        search=search,
        soc_prefix=soc_prefix,
        limit=20,
    )

    items = data["data"]
    total = data["meta"]["total"]

    if not items:
        hint = f' matching "{search}"' if search else ""
        return f"No occupations found{hint}. Try a broader search term."

    lines = [f"Occupations ({total} total, showing {len(items)}):"]
    lines.append("")

    for item in items:
        exposure_str = f"exposure: {item['exposure']:.2f}" if item.get("exposure") is not None else "unscored"
        lines.append(f"  {item['soc_code']}  {item['title']} — {exposure_str}")

    if total > len(items):
        lines.append(f"\n  ... and {total - len(items)} more. Refine your search to narrow results.")

    lines.append("")
    lines.append(
        "Use get_ai_impact with a soc_code to see full impact analysis."
    )
    return "\n".join(lines)


@mcp.tool()
async def get_occupation_tasks(
    soc_code: str,
    include_reasoning: bool = False,
) -> str:
    """Get task-level AI impact labels for a specific occupation.

    Shows how AI affects each individual task within an occupation, including
    automation and augmentation scores. Use after get_ai_impact to understand
    why an occupation scored the way it did.

    Args:
        soc_code: O*NET occupation code (e.g. '23-1011.00' for Lawyers)
        include_reasoning: Include LLM reasoning for each label (default false)
    """
    include = ["reasoning"] if include_reasoning else None
    data = await client.tasks(soc_code=soc_code, include=include)

    title = data.get("title", "Unknown")
    tasks = data.get("tasks", [])
    meta = data.get("meta", {})

    if not tasks:
        return f"No tasks found for {soc_code}. Check that the SOC code is valid."

    task_count = meta.get("task_count", len(tasks))
    impacted_count = meta.get("impacted_count", 0)

    lines = [f"Tasks for {title} ({soc_code}):"]
    lines.append(f"  {task_count} tasks, {impacted_count} impacted by AI")
    lines.append("")

    for t in tasks:
        tag = "impacted" if t.get("is_impacted") else "not_impacted"
        z = t.get("z_score", 0)
        auto = t.get("automation_score", 0)
        aug = t.get("augmentation_score", 0)
        stmt = t.get("statement", "")
        lines.append(
            f"  [{tag}] z={z:+.2f}  auto={auto:.1f}  aug={aug:.1f}  {stmt}"
        )
        if include_reasoning and t.get("impact_reasoning"):
            lines.append(f"    Reasoning: {t['impact_reasoning']}")

    lines.append("")
    lines.append(
        "Use resolve_job_title first if you have a job title instead of a SOC code."
    )
    return "\n".join(lines)


@mcp.tool()
async def get_api_info() -> str:
    """Get Abbe API metadata including model version, score cache version, and coverage statistics."""
    data = await client.meta()

    lines = ["Abbe API Info:"]
    lines.append("")
    lines.append(f"  Model version: {data['model_version']}")
    lines.append(f"  Scores version: {data.get('scores_version', 'N/A')}")
    lines.append(f"  Computed at: {data['computed_at']}")

    cov = data.get("coverage", {})
    if cov:
        lines.append(f"  Occupations: {cov.get('total_occupations', '?')} total, {cov.get('scored', '?')} scored")

    scenarios = data.get("scenarios", [])
    if scenarios:
        names = [s["name"] for s in scenarios]
        lines.append(f"  Scenarios: {', '.join(names)}")

    horizons = data.get("horizons", [])
    if horizons:
        names = [h["name"] for h in horizons]
        lines.append(f"  Horizons: {', '.join(names)}")

    if data.get("wage_data_vintage"):
        lines.append(f"  Wage data: BLS OEWS {data['wage_data_vintage']}")

    return "\n".join(lines)


# ---------------------------------------------------------------------------
# Entry point
# ---------------------------------------------------------------------------


def main():
    """Entry point for the abbe-mcp command."""
    mcp.run(transport="stdio")


if __name__ == "__main__":
    main()
