import pandas as pd
import numpy as np
import math, json, re, ast
from schemas import ScenarioPlan, TaskSpec

class ScenarioEngine:
    @staticmethod
    def render_plan(plan: ScenarioPlan, results: dict) -> str:
        sections = ["# Scenario Output\n"]
        for t in plan.tasks:
            sections.append(ScenarioEngine._render_task(t, results))
        return "\n".join(sections)

    @staticmethod
    def _df(v):
        if isinstance(v, pd.DataFrame): return v
        if isinstance(v, list) and v and isinstance(v[0], dict): return pd.DataFrame(v)
        if isinstance(v, dict): return pd.DataFrame([v])
        return None

    @staticmethod
    def _render_task(t: TaskSpec, results: dict) -> str:
        out=[f"## {t.title}\n"]
        df=ScenarioEngine._df(results.get(t.data_key)) if t.data_key else None
        if df is None: return "\n".join(out+["_No data_"])
        if t.filter: df=df.query(t.filter)
        if t.group_by or t.agg: df=df.groupby(t.group_by).agg("first").reset_index()
        if t.sort_by: df=df.sort_values(by=t.sort_by, ascending=(t.sort_dir=="asc"))
        if t.top: df=df.head(t.top)
        if t.fields: df=df[t.fields]

        if t.format=="list":
            out += [f"- {row.to_dict()}" for _,row in df.iterrows()]
        else:
            out.append(df.to_markdown(index=False))
        return "\n".join(out)