"""Build an Excel research workbook from company-specific operating KPIs.

https://vizualstocks.com/api/tutorials/company-kpi-workbook-python-excel
Python 3.10+. Install requests and openpyxl. Set VIZUALSTOCKS_API_KEY locally.
"""
import argparse
from datetime import datetime, timezone
import math
import os
from pathlib import Path
import re
from urllib.parse import urlsplit, urldefrag

import requests
from requests.adapters import HTTPAdapter
from urllib3.util.retry import Retry
from openpyxl import Workbook
from openpyxl.styles import Alignment, Font, PatternFill
from openpyxl.utils import get_column_letter

BASE_URL = "https://api.vizualstocks.com/v1"
DEFAULT_METRICS = ["customerCount", "netDollarRetention", "remainingPerformanceObligations"]
PERIOD_FIELDS = ["normalizedPeriod", "fiscalYear", "fiscalQuarter", "periodEnd"]


def make_session():
    key = os.environ.get("VIZUALSTOCKS_API_KEY", "").strip()
    if not key:
        raise ValueError("Set VIZUALSTOCKS_API_KEY in this terminal first.")
    session = requests.Session()
    session.headers.update({"X-API-Key": key, "Accept": "application/json"})
    retry = Retry(total=3, backoff_factor=1, status_forcelist=[429, 502, 503, 504],
                  allowed_methods=["GET"], respect_retry_after_header=True)
    session.mount("https://", HTTPAdapter(max_retries=retry))
    return session


def api_get(session, path, **params):
    response = session.get(BASE_URL + path, params=params, timeout=(10, 60), allow_redirects=False)
    if 300 <= response.status_code < 400:
        raise ValueError("Unexpected API redirect; check the documented API URL.")
    response.raise_for_status()
    return response.json()


def period_key(label, frequency):
    pattern = r"Q([1-4])'(\d{2})" if frequency == "quarterly" else r"(\d{4})"
    match = re.fullmatch(pattern, label or "")
    if not match:
        raise ValueError("Unexpected normalized period for the requested frequency.")
    return (2000 + int(match[2]), int(match[1])) if frequency == "quarterly" else (int(match[1]), 0)


def selected_history(payload, catalog, metric_ids, frequency):
    if payload.get("ticker") != catalog.get("ticker") or payload.get("reportingCurrency") != catalog.get("reportingCurrency"):
        raise ValueError("Catalog and history identity differ. Retry before combining them.")
    definitions = {item["metricId"]: item for item in catalog.get("metrics", [])}
    returned = {item["metricId"] for item in payload.get("metrics", [])}
    if not metric_ids or len(metric_ids) != len(set(metric_ids)):
        raise ValueError("Select distinct metric IDs.")
    for metric in metric_ids:
        if metric not in returned or metric not in definitions or frequency not in definitions[metric].get("frequencies", []):
            raise ValueError(f"Metric {metric} is unavailable for this ticker/frequency. Run --list-metrics.")
    rows, seen = [], set()
    for point in payload.get("data", []):
        label = point.get("normalizedPeriod")
        key = period_key(label, frequency)
        if key in seen:
            raise ValueError("Duplicate reporting period returned; inspect the response.")
        seen.add(key)
        for metric in metric_ids:
            value = point.get(metric)
            if value is not None and (isinstance(value, bool) or not isinstance(value, (int, float)) or not math.isfinite(value)):
                raise ValueError(f"Invalid numeric value for {metric}.")
        # A row with no selected observations is not useful history. Retain partial gaps.
        if any(point.get(metric) is not None for metric in metric_ids):
            rows.append(point)
    if not rows:
        raise ValueError("No observations returned for the selected metrics.")
    return sorted(rows, key=lambda row: period_key(row["normalizedPeriod"], frequency)), definitions


def text_cell(sheet, row, column, text):
    cell = sheet.cell(row, column, str(text))
    cell.data_type = "s"  # API descriptions are literal text, never Excel formulas.
    return cell


def source_url(value):
    """Only ordinary web links; never formulas, local files or embedded credentials."""
    if not isinstance(value, str):
        return None
    try:
        parsed = urlsplit(value)
        if parsed.scheme in ("http", "https") and parsed.hostname and not parsed.username and not parsed.password:
            return value
    except ValueError:
        pass
    return None


def add_source_table(workbook, payload, rows, metric_ids, frequency):
    if not isinstance(payload.get("sources"), list):
        raise ValueError("Request history with includeSources='true' before building a source-linked workbook.")
    sources = {}
    for item in payload["sources"]:
        key = (item.get("normalizedPeriod"), item.get("metricId"))
        if key[1] not in metric_ids:
            continue
        if key in sources:
            raise ValueError("Duplicate source for a metric/period; inspect the response before linking it.")
        sources[key] = item.get("source") or {}
    sheet = workbook.create_sheet("Source-linked data")
    text_cell(sheet, 2, 1, f"{payload['ticker']} source-linked KPI observations")
    text_cell(sheet, 3, 1, f"{frequency.title()} | Currency: {payload.get('reportingCurrency') or 'Not returned'} | Values in API base units")
    text_cell(sheet, 4, 1, "Links match this metric and normalized period only. Not returned means no source reference was supplied for that value.")
    headers = ["Normalized period", "Fiscal year", "Fiscal quarter", "Period end", "Metric ID", "Value", "Source coverage", "Original source", "Archived PDF", "PDF page", "Source image", "Location"]
    for column, title in enumerate(headers, 1): text_cell(sheet, 6, column, title)
    for history_row, point in enumerate(rows, 7):
        for metric in metric_ids:
            if point.get(metric) is None:
                continue  # No observation to cite; the wide history retains this gap.
            row = sheet.max_row + 1
            source = sources.get((point["normalizedPeriod"], metric), {})
            values = [point.get(field) for field in PERIOD_FIELDS] + [metric, point[metric]]
            for column, value in enumerate(values, 1):
                if column == 4 and value:
                    sheet.cell(row, column, datetime.strptime(value, "%Y-%m-%d")).number_format = "yyyy-mm-dd"
                elif isinstance(value, str): text_cell(sheet, row, column, value)
                elif value is not None: sheet.cell(row, column, value)
            # Reuse the exact wide-table format, including literal percentage points.
            sheet.cell(row, 6).number_format = workbook["KPI history"].cell(history_row, 5 + metric_ids.index(metric)).number_format
            text_cell(sheet, row, 7, "Returned" if source else "Not returned")
            page = source.get("pageNumber")
            if isinstance(page, int) and not isinstance(page, bool) and page > 0:
                sheet.cell(row, 10, page)
            for column, field, label in [(8, "originalSourceUrl", "Original source"), (9, "archivedPdfUrl", "Open PDF"), (11, "imageUrl", "Source image")]:
                target = source_url(source.get(field))
                if target:
                    if field == "archivedPdfUrl" and sheet.cell(row, 10).value:
                        target = urldefrag(target)[0] + f"#page={page}"
                    text_cell(sheet, row, column, label).hyperlink = target
            text_cell(sheet, row, 12, source.get("locationDescription") or "")
    for column, width in enumerate([21, 14, 15, 17, 38, 24, 20, 21, 19, 12, 20, 55], 1):
        sheet.column_dimensions[get_column_letter(column)].width = width
    sheet.freeze_panes = "F7"
    sheet.auto_filter.ref = f"A6:L{sheet.max_row}"
    sheet.row_dimensions[6].height = 34
    sheet.print_title_rows = "1:6"
    return sheet


def build_workbook(payload, catalog, metric_ids, frequency, include_sources=False):
    rows, definitions = selected_history(payload, catalog, metric_ids, frequency)
    workbook = Workbook()
    history = workbook.active
    history.title = "KPI history"
    guide = workbook.create_sheet("Metric definitions")
    ticker, currency = payload["ticker"], payload.get("reportingCurrency") or "Not returned"
    text_cell(history, 2, 1, f"{ticker} company KPI history")
    text_cell(history, 3, 1, f"{frequency.title()} observations | Reporting currency: {currency}")
    text_cell(history, 4, 1, "Values retain API base units. Blank means unavailable. Metric definitions are on the second sheet.")
    headers = PERIOD_FIELDS + metric_ids
    for column, name in enumerate(headers, 1): text_cell(history, 6, column, name)
    for row_index, point in enumerate(rows, 7):
        for column, field in enumerate(headers, 1):
            value = point.get(field)
            cell = history.cell(row_index, column)
            if field == "periodEnd" and value:
                cell.value = datetime.strptime(value, "%Y-%m-%d")
                cell.number_format = "yyyy-mm-dd"
            elif isinstance(value, str):
                text_cell(history, row_index, column, value)
            elif value is not None:
                cell.value = value
            if field in metric_ids:
                # API percentages are points (157 means 157%), not Excel fractions.
                format_name = definitions[field].get("valueFormat")
                cell.number_format = ('0.0"%"' if format_name == "percentage" else '#,##0.00' if format_name == "ratio"
                                      else '#,##0' if value is None or float(value).is_integer() else '#,##0.00######')
    history.freeze_panes = "E7"
    history.auto_filter.ref = f"A6:{get_column_letter(len(headers))}{history.max_row}"
    history.print_title_rows = "1:6"
    for column in range(1, len(headers) + 1):
        history.column_dimensions[get_column_letter(column)].width = 20 if column <= 4 else max(32, len(headers[column - 1]) + 2)
    history.row_dimensions[6].height = 34
    text_cell(guide, 2, 1, f"{ticker} metric definitions")
    text_cell(guide, 3, 1, "Retrieved UTC: " + datetime.now(timezone.utc).isoformat(timespec="seconds"))
    text_cell(guide, 4, 1, f"History: {BASE_URL}/companies/{ticker}/metrics")
    text_cell(guide, 5, 1, f"Definitions: {BASE_URL}/companies/{ticker}/metrics/catalog")
    definition_headers = ["Metric ID", "Name", "Definition", "Format", "Status", "First available", "Last available", "Frequencies"]
    for column, name in enumerate(definition_headers, 1): text_cell(guide, 6, column, name)
    for row_index, metric in enumerate(metric_ids, 7):
        item = definitions[metric]
        values = [metric, item.get("name", ""), item.get("description", ""), item.get("valueFormat", ""), item.get("status", ""),
                  item.get("firstReportedPeriod", ""), item.get("lastReportedPeriod", ""), ", ".join(item.get("frequencies", []))]
        for column, value in enumerate(values, 1): text_cell(guide, row_index, column, value or "")
        guide.row_dimensions[row_index].height = max(38, 16 * math.ceil(len(item.get("description") or "") / 65) + 12)
    for column, width in enumerate([32, 32, 65, 16, 16, 20, 20, 24], 1): guide.column_dimensions[get_column_letter(column)].width = width
    guide.freeze_panes = "C7"
    guide.auto_filter.ref = f"A6:H{6 + len(metric_ids)}"
    guide.row_dimensions[6].height = 30
    text_cell(guide, guide.max_row + 2, 1, "Catalog boundaries describe available coverage, not proof of a complete or unchanged reporting basis.")
    text_cell(guide, guide.max_row + 1, 1, "Fiscal identity and normalized chart labels can differ. Use fiscal year/quarter and period end when joining to statements.")
    source_sheet = add_source_table(workbook, payload, rows, metric_ids, frequency) if include_sources else None
    for sheet in workbook:
        sheet.sheet_view.showGridLines = False
        sheet.sheet_properties.pageSetUpPr.fitToPage = True
        sheet.page_setup.orientation = "landscape"
        sheet.page_setup.paperSize = sheet.PAPERSIZE_A3
        sheet.page_setup.fitToWidth = 1
        sheet.page_setup.fitToHeight = 0
        for row in sheet:
            for cell in row:
                cell.font = Font(name="Arial", size=11, color="0563C1" if cell.hyperlink else "24344B", underline="single" if cell.hyperlink else None)
                cell.alignment = Alignment(vertical="center", horizontal="right" if cell.data_type == "n" else "left", indent=1 if cell.row >= 7 else 0)
        sheet["A2"].font = Font(name="Arial", size=16, bold=True, color="172B4D")
        for cell in sheet[6]:
            cell.fill = PatternFill("solid", fgColor="203B60")
            cell.font = Font(name="Arial", size=11, bold=True, color="FFFFFF")
            cell.alignment = Alignment(horizontal="center", vertical="center", wrap_text=True)
        data_count = sheet.max_row - 6 if sheet is source_sheet else len(rows) if sheet is history else len(metric_ids)
        for row_index in range(7, 7 + data_count):
            if sheet is history: sheet.row_dimensions[row_index].height = 23
            if sheet is source_sheet: sheet.row_dimensions[row_index].height = 42
            for cell in sheet[row_index]:
                if row_index % 2 == 0: cell.fill = PatternFill("solid", fgColor="F0F4F9")
                if sheet is guide: cell.alignment = Alignment(vertical="top", wrap_text=True)
                if sheet is source_sheet: cell.alignment = Alignment(vertical="center", wrap_text=True, indent=1, horizontal="right" if cell.data_type == "n" else "left")
    return workbook, len(rows)


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--ticker", default="PLTR")
    parser.add_argument("--frequency", choices=["quarterly", "annual"], default="quarterly")
    parser.add_argument("--list-metrics", action="store_true")
    parser.add_argument("--metrics", nargs="+", default=DEFAULT_METRICS)
    parser.add_argument("--include-sources", action="store_true", help="Add an observation-level table with returned document and page links")
    parser.add_argument("--output", default="pltr-kpi-research.xlsx")
    args = parser.parse_args()
    try:
        ticker = args.ticker.strip().upper()
        if not re.fullmatch(r"[A-Z0-9][A-Z0-9.-]{0,19}", ticker): raise ValueError("Use a valid stock ticker.")
        output = Path(args.output)
        if not args.list_metrics and (output.suffix.lower() != ".xlsx" or output.exists()):
            raise ValueError("Choose a new .xlsx output filename. Existing files are never overwritten.")
        with make_session() as session:
            catalog = api_get(session, f"/companies/{ticker}/metrics/catalog", type="companyKpi", status="all", frequency=args.frequency)
            if args.list_metrics:
                for metric in catalog.get("metrics", []): print(f"{metric['metricId']}: {metric['name']} ({metric.get('valueFormat')}, {metric.get('status')})")
                return
            params = {"includeSources": "true"} if args.include_sources else {}
            payload = api_get(session, f"/companies/{ticker}/metrics", type="companyKpi", status="all", frequency=args.frequency, order="asc", **params)
        workbook, count = build_workbook(payload, catalog, args.metrics, args.frequency, include_sources=args.include_sources)
        with output.open("xb") as file: workbook.save(file)
        print(f"Saved {count} periods and {len(args.metrics)} metrics to {output}.")
    except requests.HTTPError as error:
        parser.exit(1, f"API HTTP {error.response.status_code}. Check your key, access and quota.\n")
    except requests.RequestException:
        parser.exit(1, "API connection failed after retries. Try again later.\n")
    except (ValueError, OSError, KeyError, TypeError) as error:
        parser.exit(1, f"{error}\n")


if __name__ == "__main__": main()