"""Resolution-specific DASH manifests. Video/audio bytes stay on the original CDN."""
import copy
import hashlib
import ipaddress
import json
import os
import re
import socket
import sqlite3
import threading
import time
import xml.etree.ElementTree as ET
from collections import OrderedDict
from contextlib import closing
from pathlib import Path
from urllib.error import HTTPError, URLError
from urllib.parse import urljoin, urlsplit
from urllib.request import HTTPRedirectHandler, Request, build_opener
from provider import ProviderError

MAX_MANIFEST = 2 * 1024 * 1024
NS = 'urn:mpeg:dash:schema:mpd:2011'
ET.register_namespace('', NS)


def local(tag):
    return tag.split('}')[-1]


def checked_url(url):
    parsed = urlsplit(url)
    if parsed.scheme != 'https' or not parsed.hostname or parsed.username or parsed.password or parsed.port not in (None,443):
        raise ProviderError('Unsupported media manifest address')
    try:
        addresses=socket.getaddrinfo(parsed.hostname,443,type=socket.SOCK_STREAM)
    except OSError:
        raise ProviderError('Media DNS lookup failed') from None
    if not addresses or any(not ipaddress.ip_address(item[4][0]).is_global for item in addresses):
        raise ProviderError('Non-public media address rejected')
    return url


class Redirect(HTTPRedirectHandler):
    def redirect_request(self, req, fp, code, msg, headers, newurl):
        checked_url(newurl)
        # Never forward a signed cookie to an unrelated redirected hostname.
        if urlsplit(req.full_url).hostname != urlsplit(newurl).hostname:
            raise ProviderError('Cross-host media redirect requires a fresh source')
        return super().redirect_request(req,fp,code,msg,headers,newurl)


def fetch_manifest(url, headers):
    checked_url(url)
    safe={k:v for k,v in headers.items() if k.lower() in ('cookie','user-agent','referer','origin','accept') and not any(c in str(v) for c in '\r\n')}
    safe['Accept']='application/dash+xml, application/xml;q=0.9, */*;q=0.5'
    try:
        with build_opener(Redirect()).open(Request(url,headers=safe),timeout=15) as response:
            data=response.read(MAX_MANIFEST+1)
            if len(data)>MAX_MANIFEST:
                raise ProviderError('Media manifest exceeds size limit')
            return data,response.url
    except HTTPError as error:
        raise ProviderError('Media manifest HTTP '+str(error.code),error.code,60) from None
    except (URLError,OSError,ValueError):
        raise ProviderError('Media manifest could not be fetched') from None


def parse_mpd(data):
    if len(data)>MAX_MANIFEST or b'<!DOCTYPE' in data.upper() or b'<!ENTITY' in data.upper():
        raise ProviderError('Unsafe or oversized DASH manifest')
    try:
        root=ET.fromstring(data)
    except ET.ParseError:
        raise ProviderError('Invalid DASH manifest') from None
    if local(root.tag)!='MPD' or root.get('type','static')!='static':
        raise ProviderError('Only static movie/episode DASH is supported')
    return root


def video_set(adaptation):
    return adaptation.get('contentType')=='video' or adaptation.get('mimeType','').startswith('video/') or any(rep.get('height') or rep.get('mimeType','').startswith('video/') for rep in adaptation if local(rep.tag)=='Representation')


def height_of(rep, adaptation):
    value=rep.get('height',adaptation.get('height',''))
    return int(value) if value.isdigit() and 1<=int(value)<=8640 else None


def available_heights(root):
    choices=[]
    for period in [item for item in root if local(item.tag)=='Period']:
        heights=set()
        for adaptation in period:
            if local(adaptation.tag)=='AdaptationSet' and video_set(adaptation):
                heights.update(h for rep in adaptation if local(rep.tag)=='Representation' and (h:=height_of(rep,adaptation)))
        if heights:
            choices.append(heights)
    return sorted(set.intersection(*choices),reverse=True) if choices else []


def filtered_manifest(data, origin, height, adaptive=False):
    root=parse_mpd(data)
    available=available_heights(root)
    allowed={h for h in available if h<=height} if adaptive else {height}
    if height not in available or (adaptive and len(allowed)<2):
        raise LookupError('Requested resolution is unavailable')
    for period in [item for item in root if local(item.tag)=='Period']:
        kept_video=False
        had_video=False
        for adaptation in list(period):
            if local(adaptation.tag)!='AdaptationSet' or not video_set(adaptation):
                continue
            had_video=True
            for rep in list(adaptation):
                if local(rep.tag)=='Representation':
                    if rep.get('dependencyId'):
                        raise ProviderError('Layered video representation cannot be safely separated')
                    if height_of(rep,adaptation) not in allowed:
                        adaptation.remove(rep)
            remaining=[rep for rep in adaptation if local(rep.tag)=='Representation']
            if not remaining:
                period.remove(adaptation)
            else:
                kept_video=True
                for key in (() if adaptive else ('height','minHeight','maxHeight')):
                    if key in adaptation.attrib:
                        adaptation.set(key,str(height))
                if adaptive:
                    adaptation.attrib.pop('height',None)
                    adaptation.set('minHeight',str(min(allowed)))
                    adaptation.set('maxHeight',str(max(allowed)))
        if had_video and not kept_video:
            raise LookupError('Resolution unavailable for a video period')
    # The new document lives on our domain, so anchor all relative segment paths
    # to the original CDN rather than our add-on. Child BaseURLs remain intact.
    namespace=root.tag.partition('}')[0]+'}' if root.tag.startswith('{') else ''
    roots=[item for item in root if local(item.tag)=='BaseURL']
    if roots:
        for item in roots:
            item.text=urljoin(origin,(item.text or '').strip())
    else:
        item=ET.Element(namespace+'BaseURL'); item.text=urljoin(origin,'.')
        index=next((i for i,v in enumerate(root) if local(v.tag) not in ('ProgramInformation',)),len(root))
        root.insert(index,item)
    # A Location refresh must not send the player back to the adaptive MPD.
    for item in list(root):
        if local(item.tag) in ('Location','PatchLocation'):
            root.remove(item)
    return ET.tostring(root,encoding='utf-8',xml_declaration=True)


class Descriptors:
    """Opaque IDs persist across restarts; signed URL/cookies never enter our URLs."""
    def __init__(self,state=None):
        self.guard=threading.Lock()
        self.memory=OrderedDict()
        self.path=None
        if state:
            directory=Path(state); directory.mkdir(parents=True,exist_ok=True,mode=0o700)
            self.path=directory/'quality-descriptors.sqlite'
            with closing(sqlite3.connect(self.path)) as db, db:
                db.execute('CREATE TABLE IF NOT EXISTS descriptors (id TEXT PRIMARY KEY, expires REAL NOT NULL, payload TEXT NOT NULL)')
            self.path.chmod(0o600)

    def put(self,payload):
        text=json.dumps(payload,sort_keys=True,separators=(',',':'))
        if len(text)>32768:
            raise ProviderError('Source descriptor is too large')
        ident=hashlib.sha256(text.encode()).hexdigest()[:32]
        expiry=time.time()+3600
        with self.guard:
            if self.path:
                with closing(sqlite3.connect(self.path,timeout=3)) as db, db:
                    db.execute('DELETE FROM descriptors WHERE expires<?',(time.time(),))
                    db.execute('INSERT OR REPLACE INTO descriptors VALUES(?,?,?)',(ident,expiry,text))
                    db.execute('DELETE FROM descriptors WHERE id NOT IN (SELECT id FROM descriptors ORDER BY expires DESC LIMIT 512)')
            else:
                self.memory[ident]=(expiry,text); self.memory.move_to_end(ident)
                while len(self.memory)>512:
                    self.memory.popitem(last=False)
        return ident

    def get(self,ident):
        with self.guard:
            if self.path:
                with closing(sqlite3.connect(self.path,timeout=3)) as db, db:
                    row=db.execute('SELECT expires,payload FROM descriptors WHERE id=?',(ident,)).fetchone()
            else:
                row=self.memory.get(ident)
        if not row or row[0]<time.time():
            raise LookupError('Source choice expired; refresh the source list')
        return json.loads(row[1])


class Quality:
    def __init__(self,fetch=fetch_manifest,state=None):
        self.fetch=fetch
        self.store=Descriptors(state)
        self.guard=threading.Lock()
        self.cache=OrderedDict()
        self.failures=OrderedDict()
        self.locks=[threading.Lock() for _ in range(16)]

    def load(self,payload):
        key=hashlib.sha256(json.dumps(payload,sort_keys=True).encode()).hexdigest()
        with self.locks[int(key[:4],16)%len(self.locks)]:
            with self.guard:
                failure=self.failures.get(key)
                if failure and failure[0]>time.monotonic():
                    raise ProviderError('Media manifest cooling down',failure[1],max(1,int(failure[0]-time.monotonic())))
                row=self.cache.get(key)
                if row and row[0]>time.monotonic():
                    self.cache.move_to_end(key)
                    return row[1],row[2]
            try:
                data,origin=self.fetch(payload['url'],payload['headers'])
                parse_mpd(data)
            except ProviderError as error:
                with self.guard:
                    self.failures[key]=(time.monotonic()+(max(error.retry_after,60) if error.status==429 else 15),error.status)
                    self.failures.move_to_end(key)
                    while len(self.failures)>32:
                        self.failures.popitem(last=False)
                raise
            with self.guard:
                self.cache[key]=(time.monotonic()+30,data,origin)
                self.cache.move_to_end(key)
                while len(self.cache)>16:
                    self.cache.popitem(last=False)
            return data,origin

    def expand(self,sources,base):
        result=[]; seen=set()
        for source in sources[:4]:
            if urlsplit(source['url']).path.lower().endswith('.mpd') and base:
                payload={'url':source['url'],'headers':source.get('behaviorHints',{}).get('proxyHeaders',{}).get('request',{})}
                try:
                    data,origin=self.load(payload)
                    heights=available_heights(parse_mpd(data))
                    if not heights:
                        raise ProviderError('Manifest has no selectable video resolutions')
                    allowed=[h for h in heights if h<=1080]
                    if not allowed:
                        raise LookupError('No video track at or below 1080p')
                    cap=max(allowed)
                    adaptive=len(allowed)>1
                    document=filtered_manifest(data,origin,cap,adaptive=adaptive)
                    signature=hashlib.sha256(document).hexdigest()
                    if signature not in seen:
                        seen.add(signature)
                        ident=self.store.put(payload)
                        selected=copy.deepcopy(source)
                        mode='auto'+str(cap) if adaptive else str(cap)
                        selected['url']=base+'/dash/'+ident+'/'+mode+'/manifest.mpd'
                        selected['name']='MovieBox Auto up to '+str(cap)+'p'
                        note='adjusts quality during playback' if adaptive else 'only available video quality'
                        selected['title']=source.get('title','').split('\n',1)[0]+'\nAuto up to '+str(cap)+'p | '+note
                        selected.setdefault('behaviorHints',{})['filename']=source.get('behaviorHints',{}).get('filename','MovieBox')+'.'+mode+'.mpd'
                        result.append(selected)
                    continue
                except (ProviderError,LookupError):
                    # Do not invent quality options if the real MPD cannot be checked.
                    source=copy.deepcopy(source)
                    source['name']='MovieBox Auto'
                    source['title']=source.get('title','').split('\n')[0]+'\nAuto (adaptive; quality not verified)'
            result.append(source)
        return result[:4]

    def manifest(self,route):
        match=re.fullmatch(r'/dash/([0-9a-f]{32})/(\d{3,4}|auto\d{3,4})/manifest\.mpd',route)
        if not match:
            raise LookupError('Invalid quality manifest route')
        ident,height=match.groups()
        data,origin=self.load(self.store.get(ident))
        return filtered_manifest(data,origin,int(height[4:]) if height.startswith('auto') else int(height),adaptive=height.startswith('auto'))