import os
import json
from pathlib import Path

import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
from dotenv import load_dotenv

load_dotenv(".env")

import sys
sys.path.insert(0, str(Path(".").resolve()))
from batch_api_client import BatchAPIClient, poll_until_complete, MetricsTracker

NOTEBOOK_DIR = Path(".").resolve()
CSV_PATH = NOTEBOOK_DIR / "global_all_caps.csv"
RESULTS_DIR = NOTEBOOK_DIR / "results"
RESULTS_DIR.mkdir(exist_ok=True)

api_key = os.getenv("BIGDATA_API_KEY")
if not api_key:
    raise ValueError("Set BIGDATA_API_KEY in .env")

jsonl_content = ("\n".join(json.dumps(line) for line in query_lines) + "\n").encode("utf-8")

client = BatchAPIClient(base_url="https://api.bigdata.com", api_key=api_key)
try:
    batch_id, presigned_url, _ = client.create_batch_job()
    print(f"Created batch: {batch_id}")
    client.upload_queries_file(presigned_url, jsonl_content)
    print(f"Uploaded {num_queries:,} queries.")
    metrics = MetricsTracker(batch_id, num_queries)
    status_data = poll_until_complete(client, batch_id, metrics, poll_interval=15)
    RESULTS_PATH = RESULTS_DIR / f"batch_{batch_id}_results.jsonl"
    data = status_data.get("data") if isinstance(status_data.get("data"), dict) else status_data
    if isinstance(data.get("body"), str):
        try:
            data = json.loads(data["body"])
        except json.JSONDecodeError:
            pass
    output_file = data.get("outputFile") or data.get("output_file") or {}
    output_url = output_file.get("url") if isinstance(output_file, dict) else None
    if not output_url:
        output_url = data.get("downloadUrl") or data.get("download_url")
    if output_url:
        results_bytes, _ = client.download_results(output_url)
        RESULTS_PATH.write_bytes(results_bytes)
        print(f"Results saved: {RESULTS_PATH}")
finally:
    client.close()

path_to_use = RESULTS_PATH  # or pick latest from RESULTS_DIR
universe_ids = set(df_universe["RP_ENTITY_ID"].astype(str))
long_rows = []
seen = set()

with open(path_to_use) as f:
    for line_idx, line in enumerate(f):
        rec = json.loads(line)
        q = query_lines[line_idx].get("query", {}) if line_idx < len(query_lines) else {}
        query_ids = {str(e) for e in ((q.get("filters") or {}).get("entity") or {}).get("any_of") or []}
        for doc in rec.get("response", {}).get("results") or rec.get("results") or []:
            doc_id = doc.get("id", "")
            for ch in doc.get("chunks") or []:
                rel, sent, cnum = ch.get("relevance"), ch.get("sentiment"), ch.get("cnum")
                if rel is None or sent is None:
                    continue
                for d in ch.get("detections") or []:
                    if d.get("type") != "entity" or not (eid := d.get("id")) or eid not in universe_ids or eid not in query_ids:
                        continue
                    key = (eid, doc_id, cnum)
                    if key not in seen:
                        seen.add(key)
                        long_rows.append({"entity_id": eid, "relevance": rel, "sentiment": sent})
                        break

df_long = pd.DataFrame(long_rows)
entity_score = df_long.groupby("entity_id").apply(lambda g: (g["sentiment"] * g["relevance"]).sum(), include_groups=False).rename("score")
entity_volume = df_long.groupby("entity_id").size().rename("volume")
entity_agg = pd.concat([entity_score, entity_volume], axis=1).reset_index().rename(columns={"entity_id": "RP_ENTITY_ID"})
df_metrics = entity_agg.merge(df_universe[["RP_ENTITY_ID", "COMPANY_NAME", "SECTOR", "COUNTRY"]], on="RP_ENTITY_ID", how="inner")