from __future__ import annotations
import base64, csv, io, json, re
from html.parser import HTMLParser
from pathlib import Path
from .common import fetch_bytes, cache_payload, normalized_source, clean_text, sha256_bytes

class TextHTMLParser(HTMLParser):
    def __init__(self):
        super().__init__(); self.current_heading="Document"; self.blocks=[]; self.buf=[]; self.tag=""
    def handle_starttag(self,tag,attrs):
        self.tag=tag
        if tag in ('h1','h2','h3','h4','p','li','td','th') and self.buf: self.flush()
    def handle_endtag(self,tag):
        if tag in ('h1','h2','h3','h4'):
            text=clean_text(' '.join(self.buf)); self.buf=[]
            if text: self.current_heading=text
        elif tag in ('p','li','tr'):
            self.flush()
    def handle_data(self,data): self.buf.append(data)
    def flush(self):
        text=clean_text(' '.join(self.buf)); self.buf=[]
        if text: self.blocks.append((self.current_heading,text))

def fetch_url(url: str, cache_dir):
    body,meta=fetch_bytes(url)
    cm=cache_payload(cache_dir,"url",url,url,body,meta)
    ctype=meta.get('content-type','').lower()
    warnings=[]; sections=[]
    if 'html' in ctype or body[:100].lower().find(b'<html')>=0:
        p=TextHTMLParser(); p.feed(body.decode('utf-8',errors='replace')); p.flush()
        sections=[{"section":h,"subsection":"","locator":f"HTML block {i+1}","text":t} for i,(h,t) in enumerate(p.blocks)]
    elif 'pdf' in ctype or body.startswith(b'%PDF'):
        sections,warnings=parse_pdf_bytes(body)
    else:
        txt=body.decode('utf-8',errors='replace')
        sections=[{"section":"Document","subsection":"","locator":"text body","text":clean_text(txt)}]
    return normalized_source(connector="url",title=url.split('/')[-1] or url,url=url,source_type="web_or_document",sections=sections,raw_meta=cm,warnings=warnings)

def parse_pdf_bytes(body: bytes):
    try:
        from pypdf import PdfReader
    except Exception:
        return [], ["PDF received but pypdf is not installed. Run INSTALL_RESEARCH_DEPS.cmd, then retry."]
    reader=PdfReader(io.BytesIO(body)); sections=[]
    for i,page in enumerate(reader.pages):
        text=clean_text(page.extract_text() or "")
        sections.append({"section":f"Page {i+1}","subsection":"","locator":f"page {i+1}","text":text})
    return sections, []

def parse_upload(filename: str, data_b64: str):
    body=base64.b64decode(data_b64); lower=filename.lower(); sections=[]; warnings=[]
    if lower.endswith('.pdf'): sections,warnings=parse_pdf_bytes(body)
    elif lower.endswith(('.csv','.tsv')):
        delim='\t' if lower.endswith('.tsv') else ','
        rows=list(csv.reader(io.StringIO(body.decode('utf-8',errors='replace')),delimiter=delim))
        if rows:
            header=rows[0]
            for i,row in enumerate(rows[1:]): sections.append({"section":"Table","subsection":"","locator":f"row {i+2}","text":" | ".join(row),"data":dict(zip(header,row))})
    elif lower.endswith(('.html','.htm')):
        p=TextHTMLParser(); p.feed(body.decode('utf-8',errors='replace')); p.flush()
        sections=[{"section":h,"subsection":"","locator":f"HTML block {i+1}","text":t} for i,(h,t) in enumerate(p.blocks)]
    elif lower.endswith(('.png','.jpg','.jpeg','.webp')):
        warnings=["Image/chart intake is preserved but automated OCR/chart digitization is not enabled in this local build. Add a manual source record or use the current PhUCKIT OCR path."]
    else:
        sections=[{"section":"Uploaded document","subsection":"","locator":"file body","text":clean_text(body.decode('utf-8',errors='replace'))}]
    return normalized_source(connector="upload",title=filename,url="",source_type="uploaded_document",sections=sections,raw_meta={"sha256":sha256_bytes(body),"bytes":len(body)},warnings=warnings)
