import os
import pandas as pd
import plotly.graph_objects as go

from ui.state import MULTI_STATE
from pdf_processing.pdf_utils import process_pdf

from models.model_loader import get_model
from models.predictor import predict_with_model, predict
from models.pmfc_classifier import classify_quantitative_paragraphs

QQ_TASK = "Qualitative & Quantitative Sustainability Text Identification"


# =========================
# MULTI-FILE UPLOAD
# =========================
def load_multi_pdfs(files):
    """
    Accepts a list of uploaded PDF files (gr.File with file_count='multiple').
    Extracts paragraphs from every PDF and combines them into one DataFrame,
    tagged by Doc_name so reports can be told apart later.
    No paragraph-preview UI is shown for this page by design, but the
    combined CSV is made available for download.
    """
    if not files:
        return None, "No files uploaded."

    all_dfs = []
    filenames = []

    for file in files:
        path = file.name
        filenames.append(os.path.basename(path))

        tmp_csv = f"/tmp/{os.path.basename(path)}.csv"
        process_pdf(path, tmp_csv)

        df = pd.read_csv(tmp_csv)
        all_dfs.append(df)

    combined = pd.concat(all_dfs, ignore_index=True)

    MULTI_STATE["df"] = combined
    MULTI_STATE["filenames"] = filenames
    MULTI_STATE["coverage_df"] = None
    MULTI_STATE["commitment_df"] = None

    # Save the combined paragraph-level CSV for download
    combined_csv_path = "/tmp/combined_extracted.csv"
    combined.to_csv(combined_csv_path, index=False)

    summary = f"Uploaded {len(filenames)} report(s): " + ", ".join(filenames)

    return combined_csv_path, summary


# =========================
# COVERAGE SCORES (Step: SA model)
# =========================
def run_coverage_scores():
    """
    Runs the SA (Relevant/Irrelevant) model on all paragraphs across all
    uploaded reports (if not already run), then computes, per Doc_name:

        Coverage Score (%) = Relevant / (Relevant + Irrelevant) * 100

    Returns:
        result          -- ranked aggregate table (for the on-screen Dataframe)
        agg_csv_path     -- CSV of the ranked aggregate table
        detail_csv_path  -- CSV with one row per paragraph: Doc_name, Page_num,
                            Paragraph_num, Paragraph_content, SA_label
        status message
    """
    df = MULTI_STATE.get("df")

    if df is None or df.empty:
        return None, None, None, "Please upload reports first."

    # Run SA only if not already present
    if "SA_label" not in df.columns:
        bundle = get_model("SA")
        preds = predict_with_model(df["Paragraph_content"].tolist(), bundle)
        df["SA_label"] = preds
        MULTI_STATE["df"] = df

    # Group by Doc_name and count Relevant / Irrelevant
    grouped = df.groupby("Doc_name")["SA_label"].value_counts().unstack(fill_value=0)

    for col in ["Relevant", "Irrelevant"]:
        if col not in grouped.columns:
            grouped[col] = 0

    grouped["Coverage Score (%)"] = (
        grouped["Relevant"] / (grouped["Relevant"] + grouped["Irrelevant"]).replace(0, pd.NA) * 100
    ).round(2)

    result = grouped[["Relevant", "Irrelevant", "Coverage Score (%)"]].reset_index()
    result = result.sort_values("Coverage Score (%)", ascending=False).reset_index(drop=True)

    MULTI_STATE["coverage_df"] = result

    agg_csv_path = "/tmp/coverage_scores.csv"
    result.to_csv(agg_csv_path, index=False)

    # Per-paragraph detail CSV: Doc_name, Page_num, Paragraph_num, content, SA_label
    detail_cols = ["Doc_name", "Page_num", "Paragraph_num", "Paragraph_content", "SA_label"]
    detail_csv_path = "/tmp/coverage_paragraph_detail.csv"
    df[detail_cols].to_csv(detail_csv_path, index=False)

    return result, agg_csv_path, detail_csv_path, f"Coverage scores computed for {result['Doc_name'].nunique()} report(s)."


# =========================
# COMMITMENT SCORES (Step: QQ model + PMFC classification, requires Coverage Scores first)
# =========================
def run_commitment_scores():
    """
    Requires Coverage Scores to have been run first (reuses its SA results).

    Pipeline:
      1. Run the Qualitative/Quantitative model on paragraphs marked
         'Relevant' by SA (if not already run).
      2. For paragraphs labeled 'Quantitative', retrieve the most similar
         Performance Metrics (PM) example and most similar Future
         Commitments (FC) example (independently, via all-MiniLM-L6-v2),
         then classify each as PM or FC via one-shot prompting with
         Qwen2.5-1.5B-Instruct.
      3. Compute, per Doc_name:

         Commitment Score (%) =
             (Qualitative*1 + Quantitative_PM*2 + Quantitative_FC*3)
             / (Total_Relevant * 3) * 100

         where Total_Relevant = Qualitative + Quantitative_PM + Quantitative_FC
         (i.e. equal to the Relevant count from Coverage Scores).

    Returns:
        result          -- ranked aggregate table (for the on-screen Dataframe)
        agg_csv_path     -- CSV of the ranked aggregate table
        detail_csv_path  -- CSV with one row per paragraph: Doc_name, Page_num,
                            Paragraph_num, SA_label, QQ label, retrieved PM
                            example + similarity, retrieved FC example +
                            similarity, and the final PMFC_Label
        status message
    """
    df = MULTI_STATE.get("df")

    if df is None or df.empty:
        return None, None, None, "Please upload reports first."

    if "SA_label" not in df.columns or MULTI_STATE.get("coverage_df") is None:
        return None, None, None, "Please run Coverage Scores first."

    # Run QQ model only if not already present
    if QQ_TASK not in df.columns:
        mask = df["SA_label"] == "Relevant"
        bundle = get_model(QQ_TASK)

        texts = df.loc[mask, "Paragraph_content"].tolist()
        if len(texts) > 0:
            preds = predict(texts, bundle)
            df.loc[mask, QQ_TASK] = preds

        df.loc[~mask, QQ_TASK] = "N/A"
        MULTI_STATE["df"] = df

    # Run PMFC classification on Quantitative paragraphs only if not already present
    if "PMFC_Label" not in df.columns:
        df = classify_quantitative_paragraphs(df, QQ_TASK)
        MULTI_STATE["df"] = df

    relevant_df = df[df["SA_label"] == "Relevant"].copy()

    commitment_scores = []
    for doc, group in relevant_df.groupby("Doc_name"):
        qualitative = group[group[QQ_TASK] == "Qualitative"].shape[0]
        pm = (group["PMFC_Label"] == "Answer: PM").sum()
        fc = (group["PMFC_Label"] == "Answer: FC").sum()

        total_relevant = qualitative + pm + fc
        commitment = (
            ((qualitative * 1) + (pm * 2) + (fc * 3))
            / (total_relevant * 3)
            * 100
        ) if total_relevant > 0 else 0

        commitment_scores.append({
            "Doc_name": doc,
            "Qualitative": qualitative,
            "Quantitative_PM": int(pm),
            "Quantitative_FC": int(fc),
            "Total_Relevant": total_relevant,
            "Commitment Score (%)": round(commitment, 2)
        })

    result = pd.DataFrame(commitment_scores)
    result = result.sort_values("Commitment Score (%)", ascending=False).reset_index(drop=True)

    MULTI_STATE["commitment_df"] = result

    agg_csv_path = "/tmp/commitment_scores.csv"
    result.to_csv(agg_csv_path, index=False)

    # Per-paragraph detail CSV
    detail_cols = [
        "Doc_name", "Page_num", "Paragraph_num", "SA_label", QQ_TASK,
        "EX_PM", "PM_Similarity", "EX_FC", "FC_Similarity", "PMFC_Label"
    ]
    detail_csv_path = "/tmp/commitment_paragraph_detail.csv"
    df[detail_cols].rename(columns={
        QQ_TASK: "QQ_label",
        "EX_PM": "Most_Similar_PM_Example",
        "EX_FC": "Most_Similar_FC_Example",
    }).to_csv(detail_csv_path, index=False)

    return result, agg_csv_path, detail_csv_path, f"Commitment scores computed for {result['Doc_name'].nunique()} report(s)."


# =========================
# SHARED SCORING HELPERS (used by both the live pipeline above and the
# revised-file recalculation below, so the math only lives in one place)
# =========================
def _compute_coverage_table(df):
    """Given a df with a SA_label column, returns the ranked Coverage
    Scores table. Rows with missing/blank SA_label are excluded from
    both the Relevant and Irrelevant counts (NaN-safe)."""
    valid = df[df["SA_label"].isin(["Relevant", "Irrelevant"])]

    grouped = valid.groupby("Doc_name")["SA_label"].value_counts().unstack(fill_value=0)
    for col in ["Relevant", "Irrelevant"]:
        if col not in grouped.columns:
            grouped[col] = 0

    grouped["Coverage Score (%)"] = (
        grouped["Relevant"] / (grouped["Relevant"] + grouped["Irrelevant"]).replace(0, pd.NA) * 100
    ).round(2)

    result = grouped[["Relevant", "Irrelevant", "Coverage Score (%)"]].reset_index()
    return result.sort_values("Coverage Score (%)", ascending=False).reset_index(drop=True)


def _compute_commitment_table(df, qq_column="QQ_label", pmfc_column="PMFC_Label"):
    """Given a df with SA_label, a QQ label column, and a PMFC_Label
    column, returns the ranked Commitment Scores table using the formula:

        (Qualitative*1 + Quantitative_PM*2 + Quantitative_FC*3)
        / (Total_Relevant * 3) * 100

    Rows with missing/blank labels are excluded from all counts (NaN-safe),
    so a partially-revised file doesn't silently miscount."""
    relevant_df = df[df["SA_label"] == "Relevant"].copy()

    commitment_scores = []
    for doc, group in relevant_df.groupby("Doc_name"):
        qualitative = group[group[qq_column] == "Qualitative"].shape[0]
        pm = (group[pmfc_column] == "Answer: PM").sum()
        fc = (group[pmfc_column] == "Answer: FC").sum()

        total_relevant = qualitative + pm + fc
        commitment = (
            ((qualitative * 1) + (pm * 2) + (fc * 3)) / (total_relevant * 3) * 100
        ) if total_relevant > 0 else 0

        commitment_scores.append({
            "Doc_name": doc,
            "Qualitative": qualitative,
            "Quantitative_PM": int(pm),
            "Quantitative_FC": int(fc),
            "Total_Relevant": total_relevant,
            "Commitment Score (%)": round(commitment, 2)
        })

    result = pd.DataFrame(commitment_scores)
    if result.empty:
        return result
    return result.sort_values("Commitment Score (%)", ascending=False).reset_index(drop=True)


QUADRANT_LABELS = {
    "top_left": "Symbolic Reporting (SDG Washing Risk)",
    "top_right": "Substantive Reporting (Best Practice)",
    "bottom_left": "Minimal Engagement",
    "bottom_right": "Focused Reporting (Strategic)",
}


def _build_positioning_figure(coverage_df, commitment_df, title_suffix=""):
    """
    Shared figure-building logic for the Positioning quadrant scatter
    chart. Used by both run_positioning_chart() (live results) and
    run_revised_positioning_chart() (user-revised results), so the
    visual design only lives in one place.
    """
    merged = pd.merge(
        coverage_df[["Doc_name", "Coverage Score (%)"]],
        commitment_df[["Doc_name", "Commitment Score (%)"]],
        on="Doc_name",
        how="inner"
    )

    if merged.empty:
        return None, "No matching companies found between Coverage and Commitment results."

    x = merged["Commitment Score (%)"]
    y = merged["Coverage Score (%)"]

    avg_x = x.mean()
    avg_y = y.mean()

    fig = go.Figure()

    fig.add_trace(go.Scatter(
        x=x,
        y=y,
        mode="markers+text",
        text=merged["Doc_name"],
        textposition="top center",
        marker=dict(size=12, color="#006c4b"),
        name="Companies"
    ))

    fig.add_vline(x=avg_x, line_dash="dash", line_color="gray")
    fig.add_hline(y=avg_y, line_dash="dash", line_color="gray")

    x_min, x_max = min(0, x.min() - 5), max(100, x.max() + 5)
    y_min, y_max = min(0, y.min() - 5), max(100, y.max() + 5)

    fig.add_annotation(
        x=(x_min + avg_x) / 2, y=y_max, showarrow=False,
        text=QUADRANT_LABELS["top_left"],
        font=dict(size=12, color="#8a4b00"), align="center"
    )
    fig.add_annotation(
        x=(avg_x + x_max) / 2, y=y_max, showarrow=False,
        text=QUADRANT_LABELS["top_right"],
        font=dict(size=12, color="#006c4b"), align="center"
    )
    fig.add_annotation(
        x=(x_min + avg_x) / 2, y=y_min, showarrow=False,
        text=QUADRANT_LABELS["bottom_left"],
        font=dict(size=12, color="#8a0000"), align="center"
    )
    fig.add_annotation(
        x=(avg_x + x_max) / 2, y=y_min, showarrow=False,
        text=QUADRANT_LABELS["bottom_right"],
        font=dict(size=12, color="#1158A6"), align="center"
    )

    fig.update_layout(
        title=f"Positioning: Coverage vs Commitment{title_suffix}",
        xaxis_title="Commitment Score (%)",
        yaxis_title="Coverage Score (%)",
        xaxis=dict(range=[x_min, x_max]),
        yaxis=dict(range=[y_min, y_max]),
        showlegend=False
    )

    return fig, f"Positioning chart generated for {len(merged)} report(s)."


def run_positioning_chart():
    """
    Builds the Positioning quadrant scatter chart from the LIVE,
    model-generated Coverage/Commitment results.
    Requires both Coverage Scores and Commitment Scores to have been run.
    """
    coverage_df = MULTI_STATE.get("coverage_df")
    commitment_df = MULTI_STATE.get("commitment_df")

    if coverage_df is None:
        return None, "Please run Coverage Scores first."
    if commitment_df is None:
        return None, "Please run Commitment Scores first."

    return _build_positioning_figure(coverage_df, commitment_df)


# =========================
# RECALCULATE FROM REVISED FILE
# =========================
def recalculate_from_revised_file(file):
    """
    Lets the user upload a manually-revised per-paragraph CSV (matching
    the detail-CSV format already downloadable from Steps 2/3: Doc_name,
    Page_num, Paragraph_num, SA_label, QQ_label, PMFC_Label, etc.) and
    recomputes Coverage Scores, Commitment Scores, and the Positioning
    chart from it -- entirely independent of the live model-generated
    results, so the original tables/chart stay untouched for comparison.

    Any row with a missing/invalid SA_label, QQ_label, or PMFC_Label is
    excluded from the relevant counts rather than crashing or defaulting,
    so a partially-edited file degrades gracefully.
    """
    if file is None:
        return None, None, "Please upload a revised CSV file first."

    try:
        revised_df = pd.read_csv(file.name)
    except Exception as e:
        return None, None, f"Could not read the uploaded file: {e}"

    required_cols = ["Doc_name", "SA_label"]
    missing = [c for c in required_cols if c not in revised_df.columns]
    if missing:
        return None, None, f"Uploaded file is missing required column(s): {missing}"

    qq_column = "QQ_label" if "QQ_label" in revised_df.columns else (
        "Qualitative & Quantitative Sustainability Text Identification"
        if "Qualitative & Quantitative Sustainability Text Identification" in revised_df.columns
        else None
    )

    revised_coverage = _compute_coverage_table(revised_df)

    revised_commitment = pd.DataFrame()
    if qq_column and "PMFC_Label" in revised_df.columns:
        revised_commitment = _compute_commitment_table(revised_df, qq_column=qq_column, pmfc_column="PMFC_Label")

    MULTI_STATE["revised_coverage_df"] = revised_coverage
    MULTI_STATE["revised_commitment_df"] = revised_commitment if not revised_commitment.empty else None

    status_parts = [f"Revised Coverage Scores computed for {revised_coverage['Doc_name'].nunique()} report(s)."]
    if revised_commitment.empty:
        status_parts.append("Revised file is missing QQ/PMFC label columns -- Commitment Scores and Positioning skipped.")
    else:
        status_parts.append(f"Revised Commitment Scores computed for {revised_commitment['Doc_name'].nunique()} report(s).")

    return revised_coverage, revised_commitment, " ".join(status_parts)


def run_revised_positioning_chart():
    """
    Same quadrant-chart logic as run_positioning_chart(), but built from
    the REVISED Coverage/Commitment tables (set by
    recalculate_from_revised_file), so it can be compared side-by-side
    against the original Positioning chart.
    """
    coverage_df = MULTI_STATE.get("revised_coverage_df")
    commitment_df = MULTI_STATE.get("revised_commitment_df")

    if coverage_df is None:
        return None, "Please upload and recalculate a revised file first."
    if commitment_df is None:
        return None, "Revised file has no Commitment Scores (missing QQ/PMFC columns) -- cannot plot Positioning."

    return _build_positioning_figure(coverage_df, commitment_df, title_suffix=" (Revised)")


# =========================
# POSITIONING QUADRANT CHART (requires both Coverage and Commitment scores)
# =========================