"""Export company segment history from the VizualStocks API to CSV.

Tutorial: https://vizualstocks.com/api/tutorials/revenue-segmentation-python
Requires Python 3.10+, requests and pandas. Set VIZUALSTOCKS_API_KEY locally.
"""

import argparse
import math
import os
from pathlib import Path
import re

import pandas as pd
import requests
from requests.adapters import HTTPAdapter
from urllib3.util.retry import Retry

BASE_URL = "https://api.vizualstocks.com/v1"
PERIOD_COLUMNS = ["normalizedPeriod", "fiscalYear", "fiscalQuarter", "periodEnd"]


def make_session():
    key = os.environ.get("VIZUALSTOCKS_API_KEY", "").strip()
    if not key:
        raise ValueError("Set VIZUALSTOCKS_API_KEY to your API key before running.")
    session = requests.Session()
    session.headers.update({"X-API-Key": key, "Accept": "application/json"})
    # Retry only read requests, with backoff and respect for Retry-After.
    retries = 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=retries))
    return session


def api_get(session, path, **params):
    response = session.get(
        BASE_URL + path, params=params, timeout=(10, 60), allow_redirects=False,
    )
    # Never forward the API key to an unexpected redirect destination.
    if 300 <= response.status_code < 400:
        raise ValueError("Unexpected API redirect. Check the documented API URL.")
    response.raise_for_status()
    return response.json()


def segment_frame(payload, group_id, selected_metrics=None, keep_empty=False):
    groups = {group["groupId"]: group for group in payload.get("segmentGroups", [])}
    if group_id not in groups:
        raise ValueError("Group not returned. Run --list-groups and check plan coverage.")
    metric_ids = selected_metrics or groups[group_id]["metricIds"]
    if not metric_ids or len(set(metric_ids)) != len(metric_ids):
        raise ValueError("Choose one or more distinct metric IDs.")
    if not set(metric_ids).issubset(groups[group_id]["metricIds"]):
        raise ValueError("A selected metric does not belong to this group.")
    if not payload.get("data"):
        raise ValueError("No history returned for this company and frequency.")
    if not payload.get("reportingCurrency"):
        raise ValueError("No reporting currency returned; verify units before exporting.")
    frame = pd.DataFrame(payload["data"])
    # Keep issuer fiscal identity alongside the API's normalized chart label.
    frame = frame.reindex(columns=PERIOD_COLUMNS + metric_ids)
    for metric_id in metric_ids:
        frame[metric_id] = pd.to_numeric(frame[metric_id], errors="raise").astype("Float64")
    # Missing values remain missing. Do not invent zeros or forward-fill history.
    if not keep_empty:
        frame = frame.dropna(subset=metric_ids, how="all")
    if frame.empty:
        raise ValueError("No observations for the selected metrics and frequency.")
    frame.insert(0, "ticker", payload["ticker"])
    frame.insert(1, "reportingCurrency", payload["reportingCurrency"])
    return frame.reset_index(drop=True)


def revenue_mix_frame(payload, group_id, metric_ids, frequency="quarterly"):
    """Shares of an explicitly selected, non-overlapping revenue basket (not growth rates).

    The caller checks reporting definitions and comparability. Group membership alone
    does not establish non-overlap. Missing/negative/zero-total rows get no shares.
    """
    if frequency not in ("quarterly", "annual") or not metric_ids or len(metric_ids) < 2:
        raise ValueError("Select at least two non-overlapping revenue metrics and quarterly or annual frequency.")
    for point in payload.get("data", []):
        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 revenue for {metric}.")
    frame = segment_frame(payload, group_id, metric_ids, keep_empty=True)
    def ordinal(label):
        match = re.fullmatch(r"Q([1-4])'(\d{2})" if frequency == "quarterly" else r"(\d{4})", str(label))
        if not match:
            raise ValueError("Unexpected normalized period for the selected frequency.")
        return (2000 + int(match[2])) * 4 + int(match[1]) - 1 if frequency == "quarterly" else int(match[1])
    frame["_period"] = frame.normalizedPeriod.map(ordinal)
    if frame._period.duplicated().any():
        raise ValueError("Duplicate period returned; inspect the response before calculating mix.")
    frame = frame.sort_values("_period").reset_index(drop=True)
    amounts = frame[metric_ids]
    total = amounts.sum(axis=1, min_count=len(metric_ids))
    complete = amounts.notna().all(axis=1)
    negative = amounts.lt(0).any(axis=1)
    finite_total = total.map(lambda value: pd.notna(value) and math.isfinite(value))
    valid = complete & ~negative & total.gt(0).fillna(False) & finite_total
    frame["mixStatus"] = "ok"
    frame.loc[~complete, "mixStatus"] = "missing_revenue"
    frame.loc[complete & negative, "mixStatus"] = "negative_revenue"
    frame.loc[complete & ~negative & ~valid, "mixStatus"] = "nonpositive_or_nonfinite_total"
    frame["selectedRevenueTotal"] = total.where(valid)
    adjacent = frame._period.diff().eq(1) & valid & valid.shift(1, fill_value=False)
    frame["comparisonPeriod"] = frame.normalizedPeriod.shift(1).where(adjacent)
    for metric in metric_ids:
        share = amounts[metric].div(frame.selectedRevenueTotal).mul(100)
        frame[metric + "SharePercent"] = share
        frame[metric + "ChangePp"] = share.diff().where(adjacent)
    return frame.drop(columns="_period")


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--ticker", default="MSFT")
    parser.add_argument("--frequency", choices=["quarterly", "annual"], default="quarterly")
    parser.add_argument("--list-companies", action="store_true", help="List accessible segment coverage")
    parser.add_argument("--list-groups", action="store_true", help="Show group IDs and metric definitions")
    parser.add_argument("--group", help="Exact group ID returned by --list-groups")
    parser.add_argument("--metrics", nargs="+", help="Optional subset of metric IDs in the group")
    parser.add_argument("--revenue-mix", action="store_true", help="Calculate shares and percentage-point changes for explicitly selected, non-overlapping revenue metrics")
    parser.add_argument("--output", default="segments.csv", help="New CSV file; existing files are not overwritten")
    args = parser.parse_args()
    ticker = args.ticker.strip().upper()
    if not ticker or any(char not in "ABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789.-" for char in ticker):
        parser.error("Use a ticker containing letters, digits, periods or hyphens.")
    if not (args.list_companies or args.list_groups or args.group):
        parser.error("Use --list-companies, --list-groups or --group.")
    if args.revenue_mix and (not args.group or not args.metrics):
        parser.error("Revenue mix requires --group and explicit --metrics. Choose non-overlapping revenue series.")

    try:
        with make_session() as session:
            if args.list_companies:
                companies = api_get(session, "/companies", hasSegments="true")
                for company in companies:
                    print(f'{company["ticker"]}: {company["name"]}')
                return
            payload = api_get(
                session, f"/companies/{ticker}/metrics", type="segment",
                frequency=args.frequency, status="all", order="asc",
            )
            if args.list_groups:
                definitions = {metric["metricId"]: metric for metric in payload.get("metrics", [])}
                for group in payload.get("segmentGroups", []):
                    print(f'{group["groupId"]}: {group["title"]}')
                    for metric_id in group["metricIds"]:
                        metric = definitions.get(metric_id, {})
                        print(f'  {metric_id}: {metric.get("name", "")} [{metric.get("status", "")}]')
                        if metric.get("description"):
                            print(f'    {metric["description"]}')
                if not payload.get("segmentGroups"):
                    print("No groups returned. Check company, frequency and your plan's coverage.")
                return
            frame = (revenue_mix_frame(payload, args.group, args.metrics, args.frequency)
                     if args.revenue_mix else segment_frame(payload, args.group, args.metrics))
            # Exclusive creation avoids replacing an analyst's existing export.
            with Path(args.output).open("x", encoding="utf-8-sig", newline="") as output:
                frame.to_csv(output, index=False, na_rep="")
            print(frame.tail(3).to_string(index=False))
            print(f"Saved {len(frame)} periods to {args.output}; values are not scaled to millions.")
    except requests.HTTPError as error:
        status = error.response.status_code
        hints = {401: "Check your API key.", 403: "Check API plan access and quota.",
                 404: "Check the ticker and its coverage.", 429: "Wait for your rate limit to reset."}
        parser.exit(1, f"HTTP {status}. {hints.get(status, 'Check the API status and try again.')}\n")
    except requests.RequestException:
        parser.exit(1, "Request failed after retries. Check connectivity and retry later.\n")
    except (ValueError, KeyError, OSError) as error:
        parser.exit(1, f"{error}\n")


if __name__ == "__main__":
    main()